Warp/server/internal/tracing/tracing.go

// Package tracing is Warp's OpenTelemetry tracing: a span per HTTP
// request, spans around the work that request does in the components
// (persisting an event, building a sync, calling another server), and the
// W3C trace context passed on to the servers and services Warp calls, so
// that a federation request or an application service transaction shows up
// in the same trace as the client request that caused it.
//
// Tracing is off unless tracing.enabled is set. Off, the global tracer
// provider is OpenTelemetry's no-op one, and every span below costs a
// function call that returns at once — which is what lets the components
// start spans unconditionally rather than checking a flag.
//
// Spans are exported over OTLP/HTTP, which every collector and every
// tracing backend worth running (Jaeger, Tempo, Honeycomb, an OpenTelemetry
// Collector in front of anything else) accepts.
package tracing

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.43.0"
	"go.opentelemetry.io/otel/trace"

	"github.com/dariuszjaniak/warp/server/internal/httputil"
)

// instrumentation is the name Warp's spans are reported under.
const instrumentation = "github.com/dariuszjaniak/warp/server"

// Config is what Setup needs.
type Config struct {
	Enabled bool

	// Endpoint is the OTLP/HTTP collector, e.g. "http://localhost:4318".
	// Empty leaves it to the OTEL_EXPORTER_OTLP_* environment variables,
	// and to the exporter's default of localhost:4318 without them.
	Endpoint string

	// SampleRatio is the fraction of traces started here that are kept,
	// from 0 to 1. A request that arrives with a sampled parent is always
	// kept, whatever the ratio: the caller has already decided.
	SampleRatio float64

	ServiceName    string
	ServiceVersion string

	// ServerName is the Matrix server name, recorded on every span so that
	// traces from several homeservers sharing a collector can be told
	// apart.
	ServerName string
}

// Setup installs the global tracer provider and propagator, and returns
// the function that flushes and stops them.
func Setup(ctx context.Context, cfg Config) (func(context.Context) error, error) {
	// The propagator is installed even with tracing off: a Warp that does
	// not record still passes on the trace context it was given, so that a
	// trace through it is not broken in two.
	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
		propagation.TraceContext{}, propagation.Baggage{}))
	if !cfg.Enabled {
		return func(context.Context) error { return nil }, nil
	}

	var opts []otlptracehttp.Option
	if cfg.Endpoint != "" {
		u, err := url.Parse(cfg.Endpoint)
		if err != nil || u.Host == "" {
			return nil, fmt.Errorf("tracing: endpoint %q is not a URL", cfg.Endpoint)
		}
		opts = append(opts, otlptracehttp.WithEndpointURL(cfg.Endpoint))
	}
	exporter, err := otlptracehttp.New(ctx, opts...)
	if err != nil {
		return nil, fmt.Errorf("tracing: build the OTLP exporter: %w", err)
	}

	name := cfg.ServiceName
	if name == "" {
		name = "warp"
	}
	res, err := resource.Merge(resource.Default(), resource.NewWithAttributes(semconv.SchemaURL,
		semconv.ServiceName(name),
		semconv.ServiceVersion(cfg.ServiceVersion),
		attribute.String("matrix.server_name", cfg.ServerName),
	))
	if err != nil && !errors.Is(err, resource.ErrSchemaURLConflict) {
		return nil, fmt.Errorf("tracing: describe the service: %w", err)
	}

	provider := sdktrace.NewTracerProvider(
		sdktrace.WithBatcher(exporter),
		sdktrace.WithResource(res),
		sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(cfg.SampleRatio))),
	)
	otel.SetTracerProvider(provider)
	return provider.Shutdown, nil
}

// Tracer is the tracer Warp's components start spans with.
func Tracer() trace.Tracer { return otel.Tracer(instrumentation) }

// Start begins a span in the current trace. It is a no-op span when
// tracing is off.
func Start(ctx context.Context, name string, attrs ...attribute.KeyValue) (context.Context, trace.Span) {
	return Tracer().Start(ctx, name, trace.WithAttributes(attrs...))
}

// End finishes a span, recording err on it when there is one.
func End(span trace.Span, err error) {
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
	}
	span.End()
}

// Instrument wraps one route's handler in a server span named after the
// route's pattern — "PUT /_matrix/client/v3/rooms/{roomID}/send/..." —
// rather than its path, which would make every room its own span name.
//
// It has the shape of httputil.Instrumenter, so it composes with the
// metrics middleware at route registration.
func Instrument(route string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
		ctx, span := Tracer().Start(ctx, route,
			trace.WithSpanKind(trace.SpanKindServer),
			trace.WithAttributes(
				semconv.HTTPRequestMethodKey.String(r.Method),
				semconv.HTTPRoute(route),
				semconv.URLPath(r.URL.Path),
				semconv.UserAgentOriginal(r.UserAgent()),
			))
		defer span.End()

		rec := &httputil.ResponseRecorder{ResponseWriter: w}
		next.ServeHTTP(rec, r.WithContext(ctx))

		status := rec.Status
		if status == 0 {
			status = http.StatusOK
		}
		span.SetAttributes(semconv.HTTPResponseStatusCode(status))
		if status >= http.StatusInternalServerError {
			span.SetStatus(codes.Error, strconv.Itoa(status))
		}
	})
}

// Transport wraps an outbound HTTP transport so that every request it sends
// is a client span, carrying the trace context to the other side. A nil
// base is http.DefaultTransport.
func Transport(base http.RoundTripper) http.RoundTripper {
	if base == nil {
		base = http.DefaultTransport
	}
	return &transport{base: base}
}

type transport struct{ base http.RoundTripper }

func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
	ctx, span := Tracer().Start(req.Context(), req.Method+" "+req.URL.Host,
		trace.WithSpanKind(trace.SpanKindClient),
		trace.WithAttributes(
			semconv.HTTPRequestMethodKey.String(req.Method),
			semconv.ServerAddress(req.URL.Hostname()),
			semconv.URLPath(req.URL.Path),
		))
	defer span.End()

	// RoundTrip must not modify the request it was given.
	out := req.Clone(ctx)
	otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(out.Header))
	resp, err := t.base.RoundTrip(out)
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		return nil, err
	}
	span.SetAttributes(semconv.HTTPResponseStatusCode(resp.StatusCode))
	if resp.StatusCode >= http.StatusInternalServerError {
		span.SetStatus(codes.Error, strconv.Itoa(resp.StatusCode))
	}
	return resp, nil
}