From 261161b9118df4bbc083910347c238e3b1366313 Mon Sep 17 00:00:00 2001
From: DevMiner <devminer@devminer.xyz>
Date: Sun, 11 Aug 2024 02:54:27 +0200
Subject: [PATCH] feat(tracing): map injection support

---
 fiber_middleware.go |  2 +-
 tracing.go          | 50 +++++++++++++++++++++++++++++++++++++++------
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/fiber_middleware.go b/fiber_middleware.go
index 5ca869e..b5ec754 100644
--- a/fiber_middleware.go
+++ b/fiber_middleware.go
@@ -146,7 +146,7 @@ func (t *Telemetry) FiberMiddleware(config ...FiberMiddlewareConfig) fiber.Handl
 			description,
 			WithOtelOptions(trace.WithSpanKind(trace.SpanKindServer)),
 			WithOtelTracer(tracer),
-			t.ContinueFromRequest(&stdRequest),
+			t.ContinueFromHeaders(stdRequest.Header),
 		)
 		defer func() {
 			if err := recover(); err != nil {
diff --git a/tracing.go b/tracing.go
index f9c16c1..5a3f129 100644
--- a/tracing.go
+++ b/tracing.go
@@ -14,6 +14,11 @@ type contextKey struct {
 	Key string
 }
 
+const (
+	SentryTraceHeader   = "sentry-trace"
+	SentryBaggageHeader = "sentry-baggage"
+)
+
 var (
 	tracerContextKey = contextKey{"tracer"}
 )
@@ -63,16 +68,49 @@ func (t *Telemetry) InjectIntoHeaders(ctx context.Context, headers http.Header)
 		// 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())
+		headers.Set(SentryTraceHeader, sentrySpan.ToSentryTrace())
+		headers.Set(SentryBaggageHeader, 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))
+func (t *Telemetry) InjectIntoMap(ctx context.Context, m map[string]string) {
+	t.propagator.Inject(ctx, propagation.MapCarrier(m))
 
-		return ctx, []trace.SpanStartOption{}, []sentry.SpanOption{sentry.ContinueFromRequest(r), sentry.WithTransactionSource(sentry.SourceURL)}
+	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 {
+		m[SentryTraceHeader] = sentrySpan.ToSentryTrace()
+		m[SentryBaggageHeader] = sentrySpan.ToBaggage()
+	}
+}
+
+func (t *Telemetry) ContinueFromHeaders(h http.Header) ConfigureSpanStartFunc {
+	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
+		ctx = t.propagator.Extract(ctx, propagation.HeaderCarrier(h))
+
+		return ctx,
+			[]trace.SpanStartOption{},
+			[]sentry.SpanOption{
+				sentry.ContinueFromHeaders(h.Get(sentry.SentryTraceHeader), h.Get(sentry.SentryBaggageHeader)),
+				sentry.WithTransactionSource(sentry.SourceURL),
+			}
+	}
+}
+
+func (t *Telemetry) ContinueFromMap(m map[string]string) ConfigureSpanStartFunc {
+	return func(ctx context.Context) (context.Context, []trace.SpanStartOption, []sentry.SpanOption) {
+		ctx = t.propagator.Extract(ctx, propagation.MapCarrier(m))
+
+		sentryTrace := m[sentry.SentryTraceHeader]
+		sentryBaggage := m[sentry.SentryBaggageHeader]
+
+		return ctx,
+			[]trace.SpanStartOption{},
+			[]sentry.SpanOption{
+				sentry.ContinueFromHeaders(sentryTrace, sentryBaggage),
+				sentry.WithTransactionSource(sentry.SourceURL),
+			}
 	}
 }