go-omni/telemetry/telemetry.go

// Package telemetry adds opt-in OpenTelemetry without process-global providers.
package telemetry

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

	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/logx"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/metric"
	"go.opentelemetry.io/otel/propagation"
	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"
)

type Config struct {
	Service     string                `json:"service"`
	SampleRatio float64               `json:"sampleRatio"`
	Exporter    sdktrace.SpanExporter `json:"-"`
	Reader      sdkmetric.Reader      `json:"-"`
}
type Runtime struct {
	traces     *sdktrace.TracerProvider
	meters     *sdkmetric.MeterProvider
	tracer     trace.Tracer
	propagator propagation.TextMapPropagator
	duration   metric.Float64Histogram
	requests   metric.Int64Counter
}

func New(cfg Config) (*Runtime, error) {
	if cfg.Service == "" || math.IsNaN(cfg.SampleRatio) || cfg.SampleRatio < 0 || cfg.SampleRatio > 1 {
		return nil, errors.New("telemetry: service and sampleRatio in 0..1 required")
	}
	res := resource.NewSchemaless(attribute.String("service.name", cfg.Service))
	traceOptions := []sdktrace.TracerProviderOption{sdktrace.WithResource(res), sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(cfg.SampleRatio)))}
	if cfg.Exporter != nil {
		traceOptions = append(traceOptions, sdktrace.WithBatcher(cfg.Exporter, sdktrace.WithMaxQueueSize(512), sdktrace.WithMaxExportBatchSize(128), sdktrace.WithExportTimeout(3*time.Second)))
	}
	metricOptions := []sdkmetric.Option{sdkmetric.WithResource(res)}
	if cfg.Reader != nil {
		metricOptions = append(metricOptions, sdkmetric.WithReader(cfg.Reader))
	}
	r := &Runtime{traces: sdktrace.NewTracerProvider(traceOptions...), meters: sdkmetric.NewMeterProvider(metricOptions...), propagator: propagation.TraceContext{}}
	r.tracer = r.traces.Tracer("git.yuntuo.hk/k2safe/go-omni")
	meter := r.meters.Meter("git.yuntuo.hk/k2safe/go-omni")
	r.duration, _ = meter.Float64Histogram("http.server.request.duration", metric.WithUnit("s"), metric.WithExplicitBucketBoundaries(.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10))
	r.requests, _ = meter.Int64Counter("http.server.requests")
	return r, nil
}
func (r *Runtime) Start(ctx context.Context, name string, options ...trace.SpanStartOption) (context.Context, trace.Span) {
	ctx, span := r.tracer.Start(ctx, name, options...)
	ctx = logx.WithTraceID(ctx, span.SpanContext().TraceID().String())
	return ctx, span
}
func (r *Runtime) Do(ctx context.Context, name string, work func(context.Context) error) (err error) {
	if ctx == nil || work == nil || name == "" {
		return errors.New("telemetry: context, operation name and work required")
	}
	ctx, span := r.Start(ctx, name)
	defer span.End()
	defer func() {
		if p := recover(); p != nil {
			span.SetStatus(codes.Error, "panic")
			panic(p)
		}
	}()
	err = work(ctx)
	if err != nil {
		span.SetStatus(codes.Error, "operation failed")
	}
	return err
}
func (r *Runtime) Middleware() httpx.Middleware {
	return func(next httpx.Handler) httpx.Handler {
		return func(c *httpx.Context) (err error) {
			parent := c.Context()
			ctx := r.propagator.Extract(parent, propagation.HeaderCarrier(c.Request().Header))
			route := httpx.RoutePattern(parent)
			if route == "" {
				route = "unmatched"
			}
			method := c.Request().Method
			switch method {
			case "GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS", "CONNECT", "TRACE":
			default:
				method = "OTHER"
			}
			ctx, span := r.Start(ctx, method+" "+route, trace.WithSpanKind(trace.SpanKindServer))
			c.WithContext(ctx)
			started := time.Now()
			defer func() {
				status := c.Status()
				if err != nil && !c.Responded() {
					status = c.ResolveError(err).Status
				}
				p := recover()
				if p != nil {
					status = 500
				}
				attrs := []attribute.KeyValue{attribute.String("http.request.method", method), attribute.String("http.route", route), attribute.Int("http.response.status_code", status)}
				failure := c.Failure()
				if failure == nil && err != nil {
					failure = c.ResolveError(err)
				}
				if failure != nil {
					attrs = append(attrs, attribute.Bool("omni.request.failed", true), attribute.String("omni.error.key", failure.Code))
				}
				span.SetAttributes(attrs...)
				if status >= 500 || (failure != nil && status < 400) {
					span.SetStatus(codes.Error, "request failed")
				}
				r.requests.Add(ctx, 1, metric.WithAttributes(attrs...))
				r.duration.Record(ctx, time.Since(started).Seconds(), metric.WithAttributes(attrs...))
				span.End()
				// Preserve correlation for outer access-log middleware without
				// exposing an ended span as its active span.
				c.WithContext(logx.WithTraceID(parent, span.SpanContext().TraceID().String()))
				if p != nil {
					panic(p)
				}
			}()
			return next(c)
		}
	}
}
func (r *Runtime) Shutdown(ctx context.Context) error {
	return errors.Join(r.traces.Shutdown(ctx), r.meters.Shutdown(ctx))
}
func (r *Runtime) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	return r.Shutdown(ctx)
}
func (r *Runtime) Flush(ctx context.Context) error {
	return errors.Join(r.traces.ForceFlush(ctx), r.meters.ForceFlush(ctx))
}

// Transport borrows Base. Headers are injected on a clone, and the span lives
// until EOF/Close. URL paths, query strings and credentials are never recorded.
type Transport struct {
	Runtime *Runtime
	Base    http.RoundTripper
}

func (t Transport) RoundTrip(request *http.Request) (*http.Response, error) {
	if t.Runtime == nil {
		return nil, errors.New("telemetry: runtime required")
	}
	ctx, span := t.Runtime.Start(request.Context(), "http.client", trace.WithSpanKind(trace.SpanKindClient))
	defer func() {
		if p := recover(); p != nil {
			span.SetStatus(codes.Error, "panic")
			span.End()
			panic(p)
		}
	}()
	req := request.Clone(ctx)
	if req.Header == nil {
		req.Header = make(http.Header)
	}
	t.Runtime.propagator.Inject(ctx, propagation.HeaderCarrier(req.Header))
	base := t.Base
	if base == nil {
		base = http.DefaultTransport
	}
	response, err := base.RoundTrip(req)
	if err != nil {
		span.SetStatus(codes.Error, "transport failed")
		span.End()
		if response != nil && response.Body != nil {
			response.Body.Close()
		}
		return nil, err
	}
	if response == nil || response.Body == nil {
		span.SetStatus(codes.Error, "invalid response")
		span.End()
		return nil, errors.New("telemetry: invalid transport response")
	}
	span.SetAttributes(attribute.Int("http.response.status_code", response.StatusCode))
	if response.StatusCode >= 500 {
		span.SetStatus(codes.Error, "upstream failed")
	}
	response.Body = &spanBody{ReadCloser: response.Body, span: span}
	return response, nil
}