anthropic-proxy/internal/telemetry/telemetry.go

package telemetry

import (
	"context"
	"io"
	"net/http"

	"github.com/fujin/anthropic-proxy/internal/config"
	"github.com/fujin/anthropic-proxy/internal/ratelimit"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc"
	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
	promexporter "go.opentelemetry.io/otel/exporters/prometheus"
	otellog "go.opentelemetry.io/otel/log/global"
	"go.opentelemetry.io/otel/sdk/log"
	sdkmetric "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.26.0"
)

func Setup(ctx context.Context, cfg config.TelemetryConfig, tracker *ratelimit.Tracker) (shutdown func(context.Context) error, logWriter io.Writer, metricsHandler http.Handler, err error) {
	res, err := resource.New(ctx,
		resource.WithAttributes(
			semconv.ServiceName(cfg.ServiceName),
		),
	)
	if err != nil {
		return nil, nil, nil, err
	}

	var readers []sdkmetric.Option
	readers = append(readers, sdkmetric.WithResource(res))

	var promHandler http.Handler
	if cfg.Embedded.Enabled {
		exporter, pErr := promexporter.New()
		if pErr != nil {
			return nil, nil, nil, pErr
		}
		readers = append(readers, sdkmetric.WithReader(exporter))
		promHandler = promhttp.Handler()
	}

	if !cfg.Export.Enabled() {
		mp := sdkmetric.NewMeterProvider(readers...)
		otel.SetMeterProvider(mp)
		InitMetrics(mp.Meter(cfg.ServiceName), tracker)
		return func(ctx context.Context) error { return mp.Shutdown(ctx) }, nil, promHandler, nil
	}

	traceOpts := []otlptracegrpc.Option{otlptracegrpc.WithEndpoint(cfg.Export.Endpoint)}
	metricOpts := []otlpmetricgrpc.Option{
		otlpmetricgrpc.WithEndpoint(cfg.Export.Endpoint),
		otlpmetricgrpc.WithTemporalitySelector(sdkmetric.CumulativeTemporalitySelector),
	}
	logOpts := []otlploggrpc.Option{otlploggrpc.WithEndpoint(cfg.Export.Endpoint)}
	if cfg.Export.Insecure {
		traceOpts = append(traceOpts, otlptracegrpc.WithInsecure())
		metricOpts = append(metricOpts, otlpmetricgrpc.WithInsecure())
		logOpts = append(logOpts, otlploggrpc.WithInsecure())
	}

	traceExp, err := otlptracegrpc.New(ctx, traceOpts...)
	if err != nil {
		return nil, nil, nil, err
	}
	tp := trace.NewTracerProvider(
		trace.WithBatcher(traceExp),
		trace.WithResource(res),
	)
	otel.SetTracerProvider(tp)

	metricExp, err := otlpmetricgrpc.New(ctx, metricOpts...)
	if err != nil {
		return nil, nil, nil, err
	}
	readers = append(readers, sdkmetric.WithReader(sdkmetric.NewPeriodicReader(metricExp)))
	mp := sdkmetric.NewMeterProvider(readers...)
	otel.SetMeterProvider(mp)
	InitMetrics(mp.Meter(cfg.ServiceName), tracker)

	logExp, err := otlploggrpc.New(ctx, logOpts...)
	if err != nil {
		return nil, nil, nil, err
	}
	lp := log.NewLoggerProvider(
		log.WithProcessor(log.NewBatchProcessor(logExp)),
		log.WithResource(res),
	)
	otellog.SetLoggerProvider(lp)

	bridge := &LogBridge{provider: lp}

	shutdownFn := func(ctx context.Context) error {
		var firstErr error
		if e := tp.Shutdown(ctx); e != nil && firstErr == nil {
			firstErr = e
		}
		if e := mp.Shutdown(ctx); e != nil && firstErr == nil {
			firstErr = e
		}
		if e := lp.Shutdown(ctx); e != nil && firstErr == nil {
			firstErr = e
		}
		return firstErr
	}

	return shutdownFn, bridge, promHandler, nil
}