|
| 1 | +/* |
| 2 | + * TencentBlueKing is pleased to support the open source community by making |
| 3 | + * 蓝鲸智云 - API 网关(BlueKing - APIGateway) available. |
| 4 | + * Copyright (C) 2025 Tencent. All rights reserved. |
| 5 | + * Licensed under the MIT License (the "License"); you may not use this file except |
| 6 | + * in compliance with the License. You may obtain a copy of the License at |
| 7 | + * |
| 8 | + * http://opensource.org/licenses/MIT |
| 9 | + * |
| 10 | + * Unless required by applicable law or agreed to in writing, software distributed under |
| 11 | + * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, |
| 12 | + * either express or implied. See the License for the specific language governing permissions and |
| 13 | + * limitations under the License. |
| 14 | + * |
| 15 | + * We undertake not to change the open source license (MIT license) applicable |
| 16 | + * to the current version of the project delivered to anyone in the future. |
| 17 | + */ |
| 18 | + |
| 19 | +// Package bkaidevtrace provides an independent OpenTelemetry trace pipeline |
| 20 | +// for BKAIDev Agent trace reporting. It is fully isolated from the project's |
| 21 | +// own global tracing to avoid interference. |
| 22 | +package bkaidevtrace |
| 23 | + |
| 24 | +import ( |
| 25 | + "context" |
| 26 | + "crypto/rand" |
| 27 | + "encoding/binary" |
| 28 | + "fmt" |
| 29 | + mrand "math/rand" |
| 30 | + "sync" |
| 31 | + "time" |
| 32 | + |
| 33 | + "go.opentelemetry.io/otel/attribute" |
| 34 | + "go.opentelemetry.io/otel/exporters/otlp/otlptrace" |
| 35 | + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" |
| 36 | + "go.opentelemetry.io/otel/propagation" |
| 37 | + "go.opentelemetry.io/otel/sdk/resource" |
| 38 | + sdktrace "go.opentelemetry.io/otel/sdk/trace" |
| 39 | + semconv "go.opentelemetry.io/otel/semconv/v1.10.0" |
| 40 | + tc "go.opentelemetry.io/otel/trace" |
| 41 | + |
| 42 | + "mcp_proxy/pkg/config" |
| 43 | +) |
| 44 | + |
| 45 | +var ( |
| 46 | + globalProvider *sdktrace.TracerProvider |
| 47 | + globalTracer tc.Tracer |
| 48 | + propagator propagation.TextMapPropagator |
| 49 | + once sync.Once |
| 50 | +) |
| 51 | + |
| 52 | +// Init initializes the independent BKAIDev trace pipeline. |
| 53 | +func Init(cfg config.BkAIDevTrace) error { |
| 54 | + var initErr error |
| 55 | + once.Do(func() { |
| 56 | + // WithInsecure uses HTTP instead of HTTPS for the OTLP exporter. |
| 57 | + // This is intentional: the BKAIDev trace collector runs in the same |
| 58 | + // internal network, so TLS is not required. The bk.data.token in the |
| 59 | + // resource attributes provides authentication. |
| 60 | + client := otlptracehttp.NewClient( |
| 61 | + otlptracehttp.WithEndpoint(cfg.Endpoint), |
| 62 | + otlptracehttp.WithInsecure(), |
| 63 | + ) |
| 64 | + exporter, err := otlptrace.New(context.Background(), client) |
| 65 | + if err != nil { |
| 66 | + initErr = fmt.Errorf("create bkaidevtrace exporter: %w", err) |
| 67 | + return |
| 68 | + } |
| 69 | + |
| 70 | + tp := sdktrace.NewTracerProvider( |
| 71 | + sdktrace.WithBatcher(exporter), |
| 72 | + sdktrace.WithResource(resource.NewWithAttributes( |
| 73 | + semconv.SchemaURL, |
| 74 | + semconv.ServiceNameKey.String(cfg.ServiceName), |
| 75 | + attribute.Key("bk.data.token").String(cfg.Token), |
| 76 | + )), |
| 77 | + sdktrace.WithSampler(sdktrace.AlwaysSample()), |
| 78 | + ) |
| 79 | + |
| 80 | + globalProvider = tp |
| 81 | + globalTracer = tp.Tracer(cfg.ServiceName) |
| 82 | + propagator = propagation.NewCompositeTextMapPropagator( |
| 83 | + propagation.TraceContext{}, |
| 84 | + propagation.Baggage{}, |
| 85 | + ) |
| 86 | + }) |
| 87 | + return initErr |
| 88 | +} |
| 89 | + |
| 90 | +// Enabled returns whether the BKAIDev trace pipeline is initialized. |
| 91 | +func Enabled() bool { |
| 92 | + return globalTracer != nil |
| 93 | +} |
| 94 | + |
| 95 | +// StartSpan starts a new span using the independent tracer. |
| 96 | +func StartSpan(ctx context.Context, name string, opts ...tc.SpanStartOption) (context.Context, tc.Span) { |
| 97 | + if globalTracer == nil { |
| 98 | + return ctx, nil |
| 99 | + } |
| 100 | + return globalTracer.Start(ctx, name, opts...) |
| 101 | +} |
| 102 | + |
| 103 | +// Extract extracts trace context from carrier using the independent propagator. |
| 104 | +func Extract(ctx context.Context, carrier propagation.TextMapCarrier) context.Context { |
| 105 | + if propagator == nil { |
| 106 | + return ctx |
| 107 | + } |
| 108 | + return propagator.Extract(ctx, carrier) |
| 109 | +} |
| 110 | + |
| 111 | +// Inject injects trace context into carrier using the independent propagator. |
| 112 | +func Inject(ctx context.Context, carrier propagation.TextMapCarrier) { |
| 113 | + if propagator == nil { |
| 114 | + return |
| 115 | + } |
| 116 | + propagator.Inject(ctx, carrier) |
| 117 | +} |
| 118 | + |
| 119 | +// GetTraceIDFromContext extracts the trace ID from the active span in context. |
| 120 | +func GetTraceIDFromContext(ctx context.Context) string { |
| 121 | + span := tc.SpanFromContext(ctx) |
| 122 | + if span == nil { |
| 123 | + return "" |
| 124 | + } |
| 125 | + sc := span.SpanContext() |
| 126 | + if !sc.TraceID().IsValid() { |
| 127 | + return "" |
| 128 | + } |
| 129 | + return sc.TraceID().String() |
| 130 | +} |
| 131 | + |
| 132 | +// GetSpanIDFromContext extracts the span ID from the active span in context. |
| 133 | +func GetSpanIDFromContext(ctx context.Context) string { |
| 134 | + span := tc.SpanFromContext(ctx) |
| 135 | + if span == nil { |
| 136 | + return "" |
| 137 | + } |
| 138 | + sc := span.SpanContext() |
| 139 | + if !sc.SpanID().IsValid() { |
| 140 | + return "" |
| 141 | + } |
| 142 | + return sc.SpanID().String() |
| 143 | +} |
| 144 | + |
| 145 | +// NewSpanContext creates a span context with a randomly-generated trace ID and span ID. |
| 146 | +// This is used to carry a valid trace context without creating an actual span. |
| 147 | +func NewSpanContext() tc.SpanContext { |
| 148 | + var traceID tc.TraceID |
| 149 | + var spanID tc.SpanID |
| 150 | + |
| 151 | + // Try crypto/rand first for both traceID and spanID. |
| 152 | + // If crypto/rand fails (extremely rare), fallback to a single math/rand instance |
| 153 | + // to avoid seed collision when two separate instances are created with the same |
| 154 | + // time-based seed under high concurrency. |
| 155 | + traceOK := true |
| 156 | + if _, err := rand.Read(traceID[:]); err != nil { |
| 157 | + traceOK = false |
| 158 | + } |
| 159 | + if _, err := rand.Read(spanID[:]); err != nil { |
| 160 | + if traceOK { |
| 161 | + // Only traceID succeeded via crypto/rand; generate spanID with math/rand |
| 162 | + r := mrand.New(mrand.NewSource(time.Now().UnixNano())) |
| 163 | + binary.BigEndian.PutUint64(spanID[:], uint64(r.Int63())) |
| 164 | + } |
| 165 | + } |
| 166 | + if !traceOK { |
| 167 | + // crypto/rand failed for traceID; use a single math/rand instance for both |
| 168 | + r := mrand.New(mrand.NewSource(time.Now().UnixNano())) |
| 169 | + binary.BigEndian.PutUint64(traceID[:8], uint64(r.Int63())) |
| 170 | + binary.BigEndian.PutUint64(traceID[8:], uint64(r.Int63())) |
| 171 | + binary.BigEndian.PutUint64(spanID[:], uint64(r.Int63())) |
| 172 | + } |
| 173 | + |
| 174 | + return tc.NewSpanContext(tc.SpanContextConfig{ |
| 175 | + TraceID: traceID, |
| 176 | + SpanID: spanID, |
| 177 | + TraceFlags: tc.FlagsSampled, |
| 178 | + }) |
| 179 | +} |
| 180 | + |
| 181 | +// Shutdown flushes and shuts down the tracer provider. |
| 182 | +func Shutdown(ctx context.Context) error { |
| 183 | + if globalProvider == nil { |
| 184 | + return nil |
| 185 | + } |
| 186 | + return globalProvider.Shutdown(ctx) |
| 187 | +} |
0 commit comments