Pass and Receive Traces

clues.go

@@ -72,9 +72,37 @@ func AddMap[K comparable, V any](
 }
 
 // ---------------------------------------------------------------------------
-// traces
+// spans and traces
 // ---------------------------------------------------------------------------
 
+// PassTrace adds the current trace details to the provided
+// headers.  If otel is not initialized, no-ops.
+//
+// The mapCarrier is mutated by this request.  The passed
+// reference is returned mostly as a quality-of-life step
+// so that callers don't need to declare the map outside of
+// this call.
+func PassTrace[C traceMapCarrierBase](
+	ctx context.Context,
+	mapCarrier C,
+) C {
+	nodeFromCtx(ctx).
+		passTrace(ctx, asTraceMapCarrier(mapCarrier))
+
+	return mapCarrier
+}
+
+// ReceiveTrace extracts the current trace details from the
+// headers and adds them to the context.  If otel is not
+// initialized, no-ops.
+func ReceiveTrace[C traceMapCarrierBase](
+	ctx context.Context,
+	mapCarrier C,
+) context.Context {
+	return nodeFromCtx(ctx).
+		receiveTrace(ctx, asTraceMapCarrier(mapCarrier))
+}
+
 // AddSpan stacks a clues node onto this context and uses the provided
 // name for the trace id, instead of a randomly generated hash. AddSpan
 // can be called without additional values if you only want to add a trace
@@ -105,10 +133,8 @@ func AddSpan(
 // CloseSpan closes the current span in the clues node.  Should only be called
 // following a `clues.AddSpan()` call.
 func CloseSpan(ctx context.Context) context.Context {
-	nc := nodeFromCtx(ctx)
-	node := nc.closeSpan(ctx)
-
-	return setNodeInCtx(ctx, node)
+	nc := nodeFromCtx(ctx).closeSpan(ctx)
+	return setNodeInCtx(ctx, nc)
 }
 
 // ---------------------------------------------------------------------------

datanode.go

@@ -3,14 +3,17 @@ package clues
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"path"
 	"runtime"
 	"strings"
 
 	"github.com/alcionai/clues/internal/stringify"
 	"github.com/google/uuid"
+	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	otellog "go.opentelemetry.io/otel/log"
+	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/exp/maps"
 	"golang.org/x/exp/slices"
@@ -403,6 +406,66 @@ func setNodeInCtx(ctx context.Context, dn *dataNode) context.Context {
 // span handling
 // ------------------------------------------------------------
 
+// traceMapCarrierBase defines the structures that support
+// otel traceMapCarrier behavior.  A traceMapCarrier is used
+// to pass and receive traces using message delivery headers
+// and other metadata.
+type traceMapCarrierBase interface {
+	map[string]string | http.Header
+}
+
+// asTraceMapCarrier converts a traceMapCarrier interface to
+// its propagation package implementation for that structure.
+// ie: map becomes a MapCarrier, headers become HeaderCarriers.
+func asTraceMapCarrier[C traceMapCarrierBase](
+	carrier C,
+) propagation.TextMapCarrier {
+	if carrier == nil {
+		return propagation.MapCarrier{}
+	}
+
+	if mss, ok := any(carrier).(map[string]string); ok {
+		return propagation.MapCarrier(mss)
+	}
+
+	if hh, ok := any(carrier).(http.Header); ok {
+		return propagation.HeaderCarrier(hh)
+	}
+
+	return propagation.MapCarrier{}
+}
+
+// passTrace adds the current trace details to the provided
+// carrier.  If otel is not initialized, no-ops.
+//
+// The carrier data is mutated by this call.
+func (dn *dataNode) passTrace(
+	ctx context.Context,
+	carrier propagation.TextMapCarrier,
+) {
+	if dn == nil {
+		return
+	}
+
+	otel.GetTextMapPropagator().Inject(ctx, carrier)
+}
+
+// receiveTrace extracts the current trace details from the
+// carrier and adds them to the context.  If otel is not
+// initialized, no-ops.
+//
+// The carrier data is mutated by this call.
+func (dn *dataNode) receiveTrace(
+	ctx context.Context,
+	carrier propagation.TextMapCarrier,
+) context.Context {
+	if dn == nil {
+		return ctx
+	}
+
+	return otel.GetTextMapPropagator().Extract(ctx, carrier)
+}
+
 // addSpan adds a new otel span.  If the otel client is nil, no-ops.
 // Attrs can be added to the span with addSpanAttrs.  This span will
 // continue to be used for that purpose until replaced with another