unitel/telemetry.go

package unitel

import (
	"context"
	"os"
	"time"

	"github.com/getsentry/sentry-go"
	"github.com/rs/zerolog/log"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
	"go.opentelemetry.io/otel/exporters/stdout/stdoutmetric"
	"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
	"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"
	semconv "go.opentelemetry.io/otel/semconv/v1.26.0"
	"go.opentelemetry.io/otel/trace"
	"go.opentelemetry.io/otel/trace/noop"
)

type telemetryContextKey struct{}

type Opts struct {
	Environment  string
	ServiceName  string
	OTLPEndpoint string
	SentryDSN    string
}

func ParseOpts(serviceName string) Opts {
	return Opts{
		Environment:  os.Getenv("ENVIRONMENT"),
		ServiceName:  serviceName,
		OTLPEndpoint: os.Getenv("OTLP_ENDPOINT"),
		SentryDSN:    os.Getenv("SENTRY_DSN"),
	}
}

type Telemetry struct {
	Opts           Opts
	propagator     propagation.TextMapPropagator
	tracerProvider trace.TracerProvider
	meterProvider  metric.MeterProvider
	Tracer         trace.Tracer
	Meter          metric.Meter
}

func (o *Telemetry) Shutdown(ctx context.Context) {
	if tp, ok := o.tracerProvider.(*sdktrace.TracerProvider); ok {
		if err := tp.Shutdown(ctx); err != nil {
			log.Error().Err(err).Msg("failed to shutdown TracerProvider")
		}
	}

	if mp, ok := o.meterProvider.(*sdkmetric.MeterProvider); ok {
		if err := mp.Shutdown(ctx); err != nil {
			log.Error().Err(err).Msg("failed to stop MeterProvider")
		}
	}

	if o.Opts.SentryDSN != "" {
		if ok := sentry.Flush(5 * time.Second); !ok {
			log.Error().Msg("failed to flush Sentry")
		}
	}
}

func SetOnContext(o *Telemetry, ctx context.Context) context.Context {
	return context.WithValue(ctx, telemetryContextKey{}, o)
}
func FromContext(ctx context.Context) *Telemetry {
	return ctx.Value(telemetryContextKey{}).(*Telemetry)
}

func newSpanExporter(ctx context.Context, opts Opts) (sdktrace.SpanExporter, error) {
	return otlptracegrpc.New(
		ctx,
		otlptracegrpc.WithInsecure(),
		otlptracegrpc.WithEndpoint(opts.OTLPEndpoint),
		otlptracegrpc.WithReconnectionPeriod(time.Second*5),
	)
}

func newMetricExporter(ctx context.Context, opts Opts) (*otlpmetricgrpc.Exporter, error) {
	return otlpmetricgrpc.New(
		ctx,
		otlpmetricgrpc.WithInsecure(),
		otlpmetricgrpc.WithEndpoint(opts.OTLPEndpoint),
		otlpmetricgrpc.WithReconnectionPeriod(time.Second*5),
	)
}

func newTextMapPropagator() propagation.TextMapPropagator {
	return propagation.NewCompositeTextMapPropagator(
		propagation.TraceContext{},
		propagation.Baggage{},
	)
}

func getSampleRate(opts Opts) float64 {
	switch opts.Environment {
	case "development":
		return 1.0
	case "staging":
		return 0.5
	case "production":
		return 0.2
	default:
		return 1.0
	}
}

func initializeSentry(opts Opts) error {
	if opts.SentryDSN == "" {
		return nil
	}

	return sentry.Init(sentry.ClientOptions{
		Dsn:                opts.SentryDSN,
		Release:            libVersion,
		AttachStacktrace:   true,
		SampleRate:         getSampleRate(opts),
		EnableTracing:      true,
		TracesSampleRate:   1.0,
		ProfilesSampleRate: 1.0,
		Debug:              false,
		Environment:        opts.Environment,
		BeforeSend: func(event *sentry.Event, hint *sentry.EventHint) *sentry.Event {
			// Only send stuff to Sentry when it includes exceptions
			if len(event.Exception) == 0 {
				return nil
			}

			return event
		},
		// BeforeSendTransaction: func(event *sentry.Event, hint *sentry.EventHint) *sentry.Event {
		// 	// Only send stuff to Sentry when it includes exceptions
		// 	if len(event.Exception) == 0 {
		// 		return nil
		// 	}

		// 	return event
		// },
	})
}

func Initialize(opts Opts) (*Telemetry, error) {
	propagator := newTextMapPropagator()
	otel.SetTextMapPropagator(propagator)

	r, err := resource.Merge(
		resource.Default(),
		resource.NewWithAttributes(
			semconv.SchemaURL,
			semconv.ServiceName(opts.ServiceName),
			semconv.ServiceVersion(libVersion),
			semconv.DeploymentEnvironment(opts.Environment),
		),
	)
	if err != nil {
		return nil, err
	}

	var tracerProvider trace.TracerProvider
	if opts.OTLPEndpoint == "" {
		tracerProvider = noop.NewTracerProvider()
	} else {
		var exp sdktrace.SpanExporter
		var err error
		if opts.OTLPEndpoint == "console" {
			exp, err = stdouttrace.New(stdouttrace.WithPrettyPrint())
		} else {
			exp, err = newSpanExporter(context.Background(), opts)
		}

		if err != nil {
			return nil, err
		}

		tracerProvider = sdktrace.NewTracerProvider(
			sdktrace.WithBatcher(exp),
			sdktrace.WithResource(r),
			sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(getSampleRate(opts)))),
		)
	}

	var metricProvider metric.MeterProvider
	if opts.OTLPEndpoint == "" {
		// noop meter provider
		metricProvider = sdkmetric.NewMeterProvider(sdkmetric.WithResource(r))
	} else {
		var exp sdkmetric.Exporter
		if opts.OTLPEndpoint == "console" {
			exp, err = stdoutmetric.New(stdoutmetric.WithPrettyPrint())
		} else {
			exp, err = newMetricExporter(context.Background(), opts)
		}

		if err != nil {
			return nil, err
		}

		metricProvider = sdkmetric.NewMeterProvider(
			sdkmetric.WithResource(r),
			sdkmetric.WithReader(sdkmetric.NewPeriodicReader(exp, sdkmetric.WithInterval(15*time.Second))),
		)
	}

	otel.SetTracerProvider(tracerProvider)
	otel.SetMeterProvider(metricProvider)

	if err := initializeSentry(opts); err != nil {
		return nil, err
	}

	return &Telemetry{
		Opts:           opts,
		propagator:     propagator,
		tracerProvider: tracerProvider,
		meterProvider:  metricProvider,
		Tracer:         tracerProvider.Tracer(opts.ServiceName),
		Meter:          metricProvider.Meter(opts.ServiceName),
	}, nil
}

func (t *Telemetry) HandleError(ctx context.Context, err error) {
	if span := SpanFromContext(ctx); span != nil {
		span.CaptureError(err)
	} else {
		if hub := sentry.CurrentHub().Clone(); hub != nil {
			hub.CaptureException(err)
		}

		otel.Handle(err)
	}
}