unitel/span.go

package unitel

import (
	"context"
	"fmt"
	"time"

	"git.devminer.xyz/devminer/unitel/unitelutils"
	"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
}

type (
	SeverityLevel  string
	BreadcrumbType string
)

const (
	SeverityDebug   SeverityLevel = "debug"
	SeverityInfo    SeverityLevel = "info"
	SeverityWarning SeverityLevel = "warning"
	SeverityError   SeverityLevel = "error"
	SeverityFatal   SeverityLevel = "fatal"

	BreadcrumbTypeDefault BreadcrumbType = "default"
	// BreadcrumbCatagoryConsole is rendered as a BreadcrumbTypeDebug
	BreadcrumbCatagoryConsole string = "console"

	BreadcrumbTypeDebug BreadcrumbType = "debug"

	BreadcrumbTypeError     BreadcrumbType = "error"
	BreadcrumbCatagoryError string         = "error"

	// BreadcrumbTypeNavigration requires the `from` and `to` attributes
	BreadcrumbTypeNavigration     BreadcrumbType = "navigation"
	BreadcrumbCatagoryNavigration string         = "navigation"

	// BreadcrumbTypeHTTP requires the `url`, `method` and `status_code` attributes
	BreadcrumbTypeHTTP     BreadcrumbType = "http"
	BreadcrumbCatagoryHTTP string         = "http"

	BreadcrumbTypeInfo BreadcrumbType = "info"

	BreadcrumbTypeQuery BreadcrumbType = "query"

	BreadcrumbTypeTransaction           BreadcrumbType = "transaction"
	BreadcrumbCategorySentryTransaction string         = "sentry.transaction"

	// BreadcrumbTypeUI is also applicable for all `ui.*` categories
	BreadcrumbTypeUI BreadcrumbType = "ui"

	BreadcrumbTypeUser      BreadcrumbType = "user"
	BreadcrumbCategoryClick string         = "click"
)

func (s *Span) CaptureBreadcrumb(level SeverityLevel, type_ BreadcrumbType, category string, message string, data map[string]any) *Span {
	now := time.Now()

	if hub := sentry.GetHubFromContext(s.Context()); hub != nil {
		hub.AddBreadcrumb(&sentry.Breadcrumb{
			Level:     sentry.Level(level),
			Type:      string(type_),
			Category:  category,
			Message:   message,
			Data:      data,
			Timestamp: now,
		}, nil)
	}

	attrs := []attribute.KeyValue{
		attribute.String("category", category),
		attribute.String("type", string(type_)),
		attribute.String("level", string(level)),
	}
	for k, v := range data {
		attrs = append(attrs, attribute.String(k, fmt.Sprintf("%v", v)))
	}

	s.otel.AddEvent(message, trace.WithAttributes(attrs...), trace.WithTimestamp(now))

	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(unitelutils.Base+"/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)
	}
}