kurious/cmd/kuriweb/trace.go

package main

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"time"

	"git.loyso.art/frx/kurious/internal/common/config"

	"github.com/gorilla/mux"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/exporters/stdout/stdoutmetric"
	"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/resource"
	"go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.24.0"
)

var webtracer = otel.Tracer("kuriweb.http")

type shutdownFunc func(context.Context) error

func setupOtelSDK(ctx context.Context, cfg config.Trace) (shutdown shutdownFunc, err error) {
	var shutdownFuncs []shutdownFunc

	shutdown = func(ctx context.Context) error {
		var err error
		for _, f := range shutdownFuncs {
			err = errors.Join(err, f(ctx))
		}
		shutdownFuncs = nil
		return err
	}

	handleError := func(inErr error) error {
		err = errors.Join(inErr, shutdown(ctx))
		return err
	}

	prop := newPropagator()
	otel.SetTextMapPropagator(prop)

	tracerProvider, err := newCommonTraceProvider(ctx, TraceProviderParams{
		Endpoint:   cfg.Endpoint,
		Type:       cfg.Type,
		AuthHeader: cfg.APIHeader,
		APIKey:     cfg.APIKey,
	})
	if err != nil {
		return nil, handleError(err)
	}
	shutdownFuncs = append(shutdownFuncs, tracerProvider.Shutdown)
	otel.SetTracerProvider(tracerProvider)

	if cfg.ShowMetrics {
		meterProvider, err := newMeterProvider()
		if err != nil {
			return nil, handleError(err)
		}
		shutdownFuncs = append(shutdownFuncs, meterProvider.Shutdown)
		otel.SetMeterProvider(meterProvider)
	}

	return shutdown, nil
}

func newPropagator() propagation.TextMapPropagator {
	return propagation.NewCompositeTextMapPropagator(
		propagation.TraceContext{},
		propagation.Baggage{},
	)
}

type TraceProviderParams struct {
	Type       config.TraceClientType
	Endpoint   string
	APIKey     string
	AuthHeader string
}

func newCommonTraceProvider(ctx context.Context, params TraceProviderParams) (tp *trace.TracerProvider, err error) {
	r, err := resource.New(
		ctx,
		resource.WithDetectors(
			resource.StringDetector(semconv.SchemaURL, semconv.ServiceNameKey, func() (string, error) {
				return "bigstats:kuriweb", nil
			}),
		),
	)
	if err != nil {
		return nil, fmt.Errorf("making new resource: %w", err)
	}

	r, err = resource.Merge(resource.Default(), r)
	if err != nil {
		return nil, fmt.Errorf("merging resources: %w", err)
	}

	opts := make([]trace.TracerProviderOption, 0, 4)
	opts = append(
		opts,
		trace.WithSampler(trace.AlwaysSample()),
		trace.WithResource(r),
	)

	if params.Type != config.TraceClientTypeUnset {
		var spanExporter trace.SpanExporter
		var headers map[string]string

		if params.AuthHeader != "" {
			headers = make(map[string]string, 1)
			headers[params.AuthHeader] = params.APIKey
		}

		switch params.Type {
		case config.TraceClientTypeGRPC:
			spanExporter, err = otlptracegrpc.New(
				ctx,
				otlptracegrpc.WithEndpointURL(params.Endpoint),
				otlptracegrpc.WithHeaders(headers),
			)
		case config.TraceClientTypeHTTP:
			httpClient := otlptracehttp.NewClient(
				otlptracehttp.WithEndpointURL(params.Endpoint),
				otlptracehttp.WithHeaders(headers),
			)
			spanExporter, err = otlptrace.New(ctx, httpClient)
		case config.TraceClientTypeStdout:
			spanExporter, err = stdouttrace.New(stdouttrace.WithPrettyPrint())
		default:
			return nil, fmt.Errorf("unsupported provider type")
		}
		if err != nil {
			return nil, fmt.Errorf("making trace exporter: %w", err)
		}

		opts = append(opts, trace.WithBatcher(spanExporter))
	}

	tp = trace.NewTracerProvider(opts...)

	return tp, nil
}

func newMeterProvider() (*metric.MeterProvider, error) {
	metricExporter, err := stdoutmetric.New()
	if err != nil {
		return nil, err
	}

	meterProvider := metric.NewMeterProvider(
		metric.WithReader(metric.NewPeriodicReader(metricExporter,
			// Default is 1m. Set to 3s for demonstrative purposes.
			metric.WithInterval(60*time.Second))),
	)
	return meterProvider, nil
}

func muxHandleFunc(router *mux.Router, name, path string, hf http.HandlerFunc) *mux.Route {
	// h := otelhttp.WithRouteTag(path, hf)
	return router.Handle(path, hf).Name(name)
}