unitel/tracing.go

package unitel

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"github.com/getsentry/sentry-go"
	"github.com/rs/zerolog/log"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/propagation"
	semconv "go.opentelemetry.io/otel/semconv/v1.26.0"
	"go.opentelemetry.io/otel/trace"
)

type contextKey struct {
	Key string
}

var (
	tracerContextKey = contextKey{"tracer"}
	spanContextKey   = contextKey{"span"}
)

type Span struct {
	otel   trace.Span
	sentry *sentry.Span
}

func (t *Telemetry) StartSpan(ctx context.Context, operation, name string, cfgs ...ConfigureSpanStartFunc) *Span {
	otelStartOpts := make([]trace.SpanStartOption, 0)
	sentryStartOpts := []sentry.SpanOption{sentry.WithTransactionName(name), sentry.WithDescription(name)}

	for _, opt := range cfgs {
		var otel []trace.SpanStartOption
		var sentry []sentry.SpanOption

		ctx, otel, sentry = opt(ctx)

		otelStartOpts = append(otelStartOpts, otel...)
		sentryStartOpts = append(sentryStartOpts, sentry...)
	}

	var otelTracer trace.Tracer
	if tracer, ok := ctx.Value(tracerContextKey).(trace.Tracer); ok && tracer != nil {
		otelTracer = tracer
	} else {
		otelTracer = t.Tracer
	}

	ctx, otelSpan := otelTracer.Start(ctx, name, otelStartOpts...)
	sentrySpan := sentry.StartSpan(ctx, operation, sentryStartOpts...)

	return &Span{
		otel:   otelSpan,
		sentry: sentrySpan,
	}
}

func (s *Span) Context() context.Context {
	return context.WithValue(s.sentry.Context(), spanContextKey, s)
}

func (s *Span) AddAttributes(attributes ...attribute.KeyValue) *Span {
	for _, attr := range attributes {
		s.sentry.SetData(string(attr.Key), attr.Value.AsInterface())
	}

	s.otel.SetAttributes(attributes...)

	return s
}

func (s *Span) AddAttribute(key string, value any) *Span {
	var attr attribute.KeyValue
	switch v := value.(type) {
	case string:
		attr = attribute.String(key, v)
	case fmt.Stringer:
		attr = attribute.String(key, v.String())
	default:
		attr = attribute.String(key, fmt.Sprintf("%v", v))
	}

	return s.AddAttributes(attr)
}

type SpanStatus struct {
	Otel   codes.Code
	Sentry sentry.SpanStatus
}

func (s *Span) SetStatus(status SpanStatus, description string) *Span {
	s.otel.SetStatus(status.Otel, description)
	s.sentry.Status = status.Sentry

	return s
}

type SimpleStatus uint8

const (
	Unset SimpleStatus = 0
	Error SimpleStatus = 1
	Ok    SimpleStatus = 2
)

func (s *Span) SetSimpleStatus(status SimpleStatus, description string) *Span {
	var otelStatus codes.Code
	var sentryStatus sentry.SpanStatus

	switch status {
	case Error:
		otelStatus = codes.Error
		sentryStatus = sentry.SpanStatusInternalError
	case Ok:
		otelStatus = codes.Ok
		sentryStatus = sentry.SpanStatusOK
	default:
		otelStatus = codes.Unset
		sentryStatus = sentry.SpanStatusUndefined
	}

	return s.SetStatus(SpanStatus{Otel: otelStatus, Sentry: sentryStatus}, description)
}

func (s *Span) CaptureError(err error) *Span {
	if hub := sentry.GetHubFromContext(s.Context()); hub != nil {
		hub.CaptureException(err)
	}

	s.otel.RecordError(err, trace.WithStackTrace(true))

	return s
}

func (s *Span) SetName(name string) *Span {
	s.otel.SetName(name)
	s.sentry.Name = name

	return s
}

func (s *Span) SetUser(id uint64, username, ip, permissions string) *Span {
	uid := fmt.Sprintf("%d", id)

	if hub := sentry.GetHubFromContext(s.Context()); hub != nil {
		hub.Scope().SetUser(sentry.User{
			ID:        uid,
			Username:  username,
			IPAddress: ip,
		})
	}

	s.otel.SetAttributes(
		semconv.EnduserID(uid),
		semconv.EnduserScope(permissions),
		attribute.String(libBase+"/username", username),
		semconv.ClientAddress(ip),
	)

	return s
}

func (s *Span) End() {
	s.otel.End()
	s.sentry.Finish()
}

type ConfigureSpanStartFunc = func(context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption)

func WithOtelOptions(opts ...trace.SpanStartOption) ConfigureSpanStartFunc {
	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
		return ctx, opts, []sentry.SpanOption{}
	}
}

func (t *Telemetry) InjectIntoHeaders(ctx context.Context, headers http.Header) {
	t.propagator.Inject(ctx, propagation.HeaderCarrier(headers))

	if sentrySpan := sentry.SpanFromContext(ctx); sentrySpan == nil {
		// this should never happen...
		log.Trace().Msgf("failed to inject Sentry span ID, Sentry span could not be extracted from the context.Context")
	} else {
		headers.Set("sentry-trace", sentrySpan.ToSentryTrace())
		headers.Set("sentry-baggage", sentrySpan.ToBaggage())
	}
}

func (t *Telemetry) ContinueFromRequest(r *http.Request) ConfigureSpanStartFunc {
	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
		ctx = t.propagator.Extract(ctx, propagation.HeaderCarrier(r.Header))

		return ctx, []trace.SpanStartOption{}, []sentry.SpanOption{sentry.ContinueFromRequest(r), sentry.WithTransactionSource(sentry.SourceURL)}
	}
}

func WithOtelTracer(tracer trace.Tracer) ConfigureSpanStartFunc {
	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
		return context.WithValue(ctx, tracerContextKey, tracer), []trace.SpanStartOption{}, []sentry.SpanOption{}
	}
}

func SpanFromContext(ctx context.Context) *Span {
	if span, ok := ctx.Value(spanContextKey).(*Span); ok {
		return span
	}

	return nil
}

func (s *Span) Recover(ctx context.Context, err error, flushTimeout *time.Duration) {
	s.CaptureError(err)

	hub := sentry.GetHubFromContext(s.Context())
	if hub == nil {
		return
	}

	// context.WithValue(context.Background(), sentry.RequestContextKey, c),
	eventID := hub.RecoverWithContext(ctx, err)

	if flushTimeout != nil && eventID != nil {
		hub.Flush(*flushTimeout)
	}
}