l_ai_knowledge/internal/tracing/init.go

106 lines
2.6 KiB
Go

package tracing
import (
"context"
"log"
"os"
"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/stdout/stdouttrace"
"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.24.0"
"go.opentelemetry.io/otel/trace"
)
const (
AppName = "WeKnoraApp"
)
type Tracer struct {
Cleanup func(context.Context) error
}
var tracer trace.Tracer
// InitTracer initializes OpenTelemetry tracer
func InitTracer() (*Tracer, error) {
// Create resource description
labels := []attribute.KeyValue{
semconv.TelemetrySDKLanguageGo,
semconv.ServiceNameKey.String(AppName),
}
res := resource.NewWithAttributes(semconv.SchemaURL, labels...)
var err error
// First try to create OTLP exporter (can connect to Jaeger, Zipkin, etc.)
var traceExporter sdktrace.SpanExporter
if endpoint := os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT"); endpoint != "" {
// Use gRPC exporter
client := otlptracegrpc.NewClient(
otlptracegrpc.WithEndpoint(endpoint),
otlptracegrpc.WithInsecure(),
)
traceExporter, err = otlptrace.New(context.Background(), client)
if err != nil {
return nil, err
}
} else {
// If no OTLP endpoint is set, default to standard output
traceExporter, err = stdouttrace.New()
if err != nil {
return nil, err
}
}
// Create batch SpanProcessor
bsp := sdktrace.NewBatchSpanProcessor(traceExporter)
sampler := sdktrace.AlwaysSample()
// Create and register TracerProvider
tp := sdktrace.NewTracerProvider(
sdktrace.WithSampler(sampler),
sdktrace.WithResource(res),
sdktrace.WithSpanProcessor(bsp),
)
otel.SetTracerProvider(tp)
// Set global propagator
otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
propagation.TraceContext{},
propagation.Baggage{},
))
// Create Tracer for project use
tracer = tp.Tracer(AppName)
// Return cleanup function
return &Tracer{
Cleanup: func(ctx context.Context) error {
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
if err := tp.Shutdown(ctx); err != nil {
log.Printf("Error shutting down tracer provider: %v", err)
return err
}
return nil
},
}, nil
}
// GetTracer gets global Tracer
func GetTracer() trace.Tracer {
return tracer
}
// Create context with span
func ContextWithSpan(ctx context.Context, name string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
return GetTracer().Start(ctx, name, opts...)
}