package trace import ( "context" "errors" "fmt" "os" "runtime" "sync" "time" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "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/propagation" "go.opentelemetry.io/otel/sdk/resource" sdktrace "go.opentelemetry.io/otel/sdk/trace" semconv "go.opentelemetry.io/otel/semconv/v1.26.0" "go.opentelemetry.io/otel/trace" coltracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1" ) type systemContextKeyType struct{} var systemContextKey systemContextKeyType type Options struct { DebugLevel int } type systemContext struct { Options tpm *tracerProviderManager exporterServer *ExporterServer } func systemContextFromContext(ctx context.Context) *systemContext { return ctx.Value(systemContextKey).(*systemContext) } func init() { otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})) otel.SetTracerProvider(panicTracerProvider{}) } var _ trace.Tracer = panicTracer{} type tracerProviderManager struct { mu sync.Mutex tracerProviders []*sdktrace.TracerProvider } func (tpm *tracerProviderManager) ShutdownAll(ctx context.Context) error { tpm.mu.Lock() defer tpm.mu.Unlock() var errs []error for _, tp := range tpm.tracerProviders { errs = append(errs, tp.ForceFlush(ctx)) } for _, tp := range tpm.tracerProviders { errs = append(errs, tp.Shutdown(ctx)) } clear(tpm.tracerProviders) return errors.Join(errs...) } func (tpm *tracerProviderManager) Add(tp *sdktrace.TracerProvider) { tpm.mu.Lock() defer tpm.mu.Unlock() tpm.tracerProviders = append(tpm.tracerProviders, tp) } func (op Options) NewContext(ctx context.Context) context.Context { var remoteClient otlptrace.Client if os.Getenv("OTEL_EXPORTER_OTLP_PROTOCOL") == "http/protobuf" { remoteClient = otlptracehttp.NewClient() } else { remoteClient = otlptracegrpc.NewClient() } sys := &systemContext{ Options: op, tpm: &tracerProviderManager{}, } ctx = context.WithValue(ctx, systemContextKey, sys) sys.exporterServer = NewServer(ctx, remoteClient) sys.exporterServer.Start(ctx) return ctx } func NewContext(ctx context.Context) context.Context { return Options{}.NewContext(ctx) } func NewTracerProvider(ctx context.Context, serviceName string) trace.TracerProvider { _, file, line, _ := runtime.Caller(1) sys := systemContextFromContext(ctx) exp, err := otlptrace.New(ctx, sys.exporterServer.NewClient()) if err != nil { panic(err) } r, err := resource.Merge( resource.Default(), resource.NewWithAttributes( semconv.SchemaURL, semconv.ServiceName(serviceName), attribute.String("provider.created_at", fmt.Sprintf("%s:%d", file, line)), ), ) if err != nil { panic(err) } options := []sdktrace.TracerProviderOption{ sdktrace.WithBatcher(exp), sdktrace.WithResource(r), } for _, proc := range sys.exporterServer.SpanProcessors() { options = append(options, sdktrace.WithSpanProcessor(proc)) } if sys.DebugLevel >= 1 { options = append(options, sdktrace.WithSpanProcessor(&stackTraceProcessor{}), ) } tp := sdktrace.NewTracerProvider(options...) sys.tpm.Add(tp) return tp } func ShutdownContext(ctx context.Context) error { var errs []error sys := systemContextFromContext(ctx) if err := sys.tpm.ShutdownAll(context.Background()); err != nil { errs = append(errs, fmt.Errorf("(*tracerProviderManager).ShutdownAll: %w", err)) } if err := sys.exporterServer.Shutdown(context.Background()); err != nil { errs = append(errs, fmt.Errorf("(*Server).Shutdown: %w", err)) } return errors.Join(errs...) } func ExporterServerFromContext(ctx context.Context) coltracepb.TraceServiceServer { return systemContextFromContext(ctx).exporterServer } func WaitForSpans(ctx context.Context, maxDuration time.Duration) error { return systemContextFromContext(ctx).exporterServer.spanExportQueue.WaitForSpans(maxDuration) }