unitel/span.go

package unitel

import (
	"context"
	"fmt"
	"time"

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

var (
	spanContextKey = contextKey{"span"}
)

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

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

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

	return nil
}

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) CaptureMessage(message string) *Span {
	if hub := sentry.GetHubFromContext(s.Context()); hub != nil {
		hub.CaptureMessage(message)
	}

	s.otel.AddEvent(message)

	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()
}

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)
	}
}