Application Real-Time Monitoring Service:Custom instrumentation for Golang WebSocket applications

Core instrumentation logic

The following snippet shows only the instrumentation logic:

1. Get a tracer from the global TracerProvider.

2. Start a client span.

3. Inject the trace context into a MapCarrier.

4. Serialize the carrier as JSON and append it to the message body.

5. Set custom attributes on the span.

// Get the tracer
tracer := otel.GetTracerProvider().Tracer("")

// Start a client span
opts := append([]tracex.SpanStartOption{}, tracex.WithSpanKind(tracex.SpanKindClient))
ctx, span := tracer.Start(context.Background(), "Client/User defined span", opts...)
defer span.End()

// Inject trace context into a MapCarrier
var headerMap propagation.MapCarrier
headerMap = make(map[string]string)
otel.GetTextMapPropagator().Inject(ctx, headerMap)

// Serialize the carrier and append it to the message
traceJSON, _ := json.Marshal(headerMap)
message := t.String() + "|" + string(traceJSON)
err := c.WriteMessage(websocket.TextMessage, []byte(message))
if err != nil {
    log.Println("write:", err)
    return
}

// Set custom attributes on the span
span.SetAttributes(attribute.String("client", "client-with-ot"))
span.SetAttributes(attribute.Bool("user.defined", true))

Complete client code

The following example shows a complete WebSocket client that sends a message with embedded trace context every second.

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/trace"
	tracex "go.opentelemetry.io/otel/trace"
	"log"
	"net/url"
	"os"
	"os/signal"
	"time"
)

func init() {
	// Initialize TracerProvider with AlwaysSample to capture all traces
	tp := trace.NewTracerProvider(trace.WithSampler(trace.AlwaysSample()))
	otel.SetTracerProvider(tp)

	// Use W3C Trace Context propagation format
	prop := propagation.TraceContext{}
	otel.SetTextMapPropagator(prop)
}

var addr = flag.String("addr", "localhost:8080", "http service address")

func main() {
	flag.Parse()
	log.SetFlags(0)
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	u := url.URL{Scheme: "ws", Host: *addr, Path: "/echo"}
	log.Printf("connecting to %s", u.String())

	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial:", err)
	}
	defer c.Close()

	done := make(chan struct{})

	go func() {
		defer close(done)
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				return
			}
			log.Printf("recv: %s", message)
		}
	}()

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-done:
			return
		case t := <-ticker.C:
			tracer := otel.GetTracerProvider().Tracer("")
			opts := append([]tracex.SpanStartOption{}, tracex.WithSpanKind(tracex.SpanKindClient))
			// Embed trace context in the message body because WebSocket
			// does not support per-message HTTP headers for propagation.
			ctx, span := tracer.Start(context.Background(), "Client/User defined span", opts...)
			defer span.End()
			var headerMap propagation.MapCarrier
			headerMap = make(map[string]string)
			otel.GetTextMapPropagator().Inject(ctx, headerMap)
			xx, _ := json.Marshal(headerMap)
			y := t.String() + "|" + string(xx)
			err := c.WriteMessage(websocket.TextMessage, []byte(y))
			if err != nil {
				log.Println("write:", err)
				return
			}
			span.SetAttributes(attribute.String("client", "client-with-ot"))
			span.SetAttributes(attribute.Bool("user.defined", true))
		case <-interrupt:
			log.Println("interrupt")

			// Cleanly close the connection by sending a close message and then
			// waiting (with timeout) for the server to close the connection.
			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Println("write close:", err)
				return
			}
			select {
			case <-done:
			case <-time.After(time.Second):
			}
			return
		}
	}
}

Core instrumentation logic

The following snippet shows only the extraction and span creation logic:

1. Split the incoming message to isolate the JSON-encoded trace context.

2. Deserialize the trace context into a MapCarrier.

3. Extract the parent span context from the carrier.

4. Start a server span as a child of the extracted context.

5. End the span after processing the message.

// Split the message to get the trace context portion
xx := strings.Split(string(message), "|")

// Deserialize trace context from JSON into a MapCarrier
var headerMap propagation.MapCarrier
headerMap = make(map[string]string)
err = json.Unmarshal([]byte(xx[1]), &headerMap)
if err != nil {
    fmt.Println(err.Error())
}

// Extract parent span context
ctxRequest := context.Background()
xxCtx := otel.GetTextMapPropagator().Extract(ctxRequest, headerMap)

// Start a server span linked to the client span
tracer := otel.GetTracerProvider().Tracer("")
opts := append([]trace.SpanStartOption{}, trace.WithSpanKind(trace.SpanKindServer))
_, span := tracer.Start(xxCtx, "Server/User defined span", opts...)

// Process the message (echo back)
err = c.WriteMessage(mt, message)
if err != nil {
    log.Println("write:", err)
    break
}

// Print trace ID for verification
fmt.Println(span.SpanContext().TraceID())
span.End()

Complete server code

The following example shows a complete WebSocket server that extracts trace context from incoming messages and creates server spans.

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	tracex "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"
	"html/template"
	"log"
	"net/http"
	"strings"
)

var addr = flag.String("addr", "localhost:8080", "http service address")

var upgrader = websocket.Upgrader{} // use default options

func init() {
	// Initialize TracerProvider
	tp := tracex.NewTracerProvider()
	otel.SetTracerProvider(tp)

	// Use W3C Trace Context propagation format
	prop := propagation.TraceContext{}
	otel.SetTextMapPropagator(prop)
}
func echo(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		log.Printf("recv: %s", message)

		// Extract trace context from the message body
		var headerMap propagation.MapCarrier
		headerMap = make(map[string]string)
		ctxRequest := context.Background()
		xx := strings.Split(string(message), "|")
		err = json.Unmarshal([]byte(xx[1]), &headerMap)
		if err != nil {
			fmt.Println(err.Error())
		}

		// Restore parent context and start a server span
		xxCtx := otel.GetTextMapPropagator().Extract(ctxRequest, headerMap)
		tracer := otel.GetTracerProvider().Tracer("")
		opts := append([]trace.SpanStartOption{}, trace.WithSpanKind(trace.SpanKindServer))
		_, span := tracer.Start(xxCtx, "Server/User defined span", opts...)

		// Echo the message back
		err = c.WriteMessage(mt, message)
		if err != nil {
			log.Println("write:", err)
			break
		}
		fmt.Println(span.SpanContext().TraceID())
		span.End()
	}
}

func home(w http.ResponseWriter, r *http.Request) {
	homeTemplate.Execute(w, "ws://"+r.Host+"/echo")
}

func main() {
	flag.Parse()
	log.SetFlags(0)
	http.HandleFunc("/echo", echo)
	http.HandleFunc("/", home)
	log.Fatal(http.ListenAndServe(*addr, nil))
}

var homeTemplate = template.Must(template.New("").Parse(`
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script>
window.addEventListener("load", function(evt) {

    var output = document.getElementById("output");
    var input = document.getElementById("input");
    var ws;

    var print = function(message) {
        var d = document.createElement("div");
        d.textContent = message;
        output.appendChild(d);
        output.scroll(0, output.scrollHeight);
    };

    document.getElementById("open").onclick = function(evt) {
        if (ws) {
            return false;
        }
        ws = new WebSocket("{{.}}");
        ws.onopen = function(evt) {
            print("OPEN");
        }
        ws.onclose = function(evt) {
            print("CLOSE");
            ws = null;
        }
        ws.onmessage = function(evt) {
            print("RESPONSE: " + evt.data);
        }
        ws.onerror = function(evt) {
            print("ERROR: " + evt.data);
        }
        return false;
    };

    document.getElementById("send").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        print("SEND: " + input.value);
        ws.send(input.value);
        return false;
    };

    document.getElementById("close").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        ws.close();
        return false;
    };

});
</script>
</head>
<body>
<table>
<tr><td valign="top" width="50%">
<p>Click "Open" to create a connection to the server,
"Send" to send a message to the server and "Close" to close the connection.
You can change the message and send multiple times.
<p>
<form>
<button id="open">Open</button>
<button id="close">Close</button>
<p><input id="input" type="text" value="Hello world!">
<button id="send">Send</button>
</form>
</td><td valign="top" width="50%">
<div id="output" style="max-height: 70vh;overflow-y: scroll;"></div>
</td></tr></table>
</body>
</html>
`))