OTEL-2540 Add SpanContext to persistent queue

Author: jackgopack4

exporter/exporterhelper/internal/queuebatch/persistent_queue.go

@@ -6,11 +6,14 @@ package queuebatch // import "go.opentelemetry.io/collector/exporter/exporterhel
 import (
 	"context"
 	"encoding/binary"
+	"encoding/hex"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"strconv"
 	"sync"
 
+	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 
 	"go.opentelemetry.io/collector/component"
@@ -29,6 +32,8 @@ const (
 	writeIndexKey               = "wi"
 	currentlyDispatchedItemsKey = "di"
 	queueSizeKey                = "si"
+
+	errInvalidTraceFlagsLength = "trace flags must only be 1 byte"
 )
 
 var (
@@ -238,6 +243,56 @@ func (pq *persistentQueue[T]) Offer(ctx context.Context, req T) error {
 	return pq.putInternal(ctx, req)
 }
 
+type marshaledRequestWithSpanContext struct {
+	RequestBytes    []byte          `json:"request"`
+	SpanContextJSON json.RawMessage `json:"span_context"`
+}
+
+type spanContextConfigWrapper struct {
+	TraceID    string
+	SpanID     string
+	TraceFlags string
+	TraceState string
+	Remote     bool
+}
+
+func SpanContextFromWrapper(wrapper spanContextConfigWrapper) (*trace.SpanContext, error) {
+	traceID, err := trace.TraceIDFromHex(wrapper.TraceID)
+	if err != nil {
+		return nil, err
+	}
+	spanID, err := trace.SpanIDFromHex(wrapper.SpanID)
+	if err != nil {
+		return nil, err
+	}
+	decoded, err := hex.DecodeString(wrapper.TraceFlags)
+	if err != nil {
+		return nil, err
+	}
+	if len(decoded) != 1 {
+		return nil, errors.New(errInvalidTraceFlagsLength)
+	}
+	traceFlags := trace.TraceFlags(decoded[0])
+	traceState, err := trace.ParseTraceState(wrapper.TraceState)
+	if err != nil {
+		return nil, err
+	}
+
+	sc := trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    traceID,
+		SpanID:     spanID,
+		TraceFlags: traceFlags,
+		TraceState: traceState,
+		Remote:     wrapper.Remote,
+	})
+
+	if !sc.IsValid() {
+		return nil, nil
+	}
+
+	return &sc, nil
+}
+
 // putInternal is the internal version that requires caller to hold the mutex lock.
 func (pq *persistentQueue[T]) putInternal(ctx context.Context, req T) error {
 	reqSize := pq.set.sizer.Sizeof(req)
@@ -254,11 +309,24 @@ func (pq *persistentQueue[T]) putInternal(ctx context.Context, req T) error {
 	if err != nil {
 		return err
 	}
-
+	// Retrieve SpanContext object from provided context, and store alongside the request
+	sc := trace.SpanContextFromContext(ctx)
+	scJSON, err := json.Marshal(sc)
+	if err != nil {
+		return err
+	}
+	envelope := marshaledRequestWithSpanContext{
+		RequestBytes:    reqBuf,
+		SpanContextJSON: scJSON,
+	}
+	envelopeBytes, err := json.Marshal(envelope)
+	if err != nil {
+		return err
+	}
 	// Carry out a transaction where we both add the item and update the write index
 	ops := []*storage.Operation{
 		storage.SetOperation(writeIndexKey, itemIndexToBytes(pq.writeIndex+1)),
-		storage.SetOperation(getItemKey(pq.writeIndex), reqBuf),
+		storage.SetOperation(getItemKey(pq.writeIndex), envelopeBytes),
 	}
 	if err = pq.client.Batch(ctx, ops...); err != nil {
 		return err
@@ -291,7 +359,7 @@ func (pq *persistentQueue[T]) Read(ctx context.Context) (context.Context, T, Don
 
 		// Read until either a successful retrieved element or no more elements in the storage.
 		for pq.readIndex != pq.writeIndex {
-			index, req, consumed := pq.getNextItem(ctx)
+			index, req, consumed, restoredContext := pq.getNextItem(ctx)
 			// Ensure the used size and the channel size are in sync.
 			if pq.readIndex == pq.writeIndex {
 				pq.queueSize = 0
@@ -300,7 +368,7 @@ func (pq *persistentQueue[T]) Read(ctx context.Context) (context.Context, T, Don
 			if consumed {
 				id := indexDonePool.Get().(*indexDone)
 				id.reset(index, pq.set.sizer.Sizeof(req), pq)
-				return context.Background(), req, id, true
+				return restoredContext, req, id, true
 			}
 		}
 
@@ -313,7 +381,7 @@ func (pq *persistentQueue[T]) Read(ctx context.Context) (context.Context, T, Don
 // getNextItem pulls the next available item from the persistent storage along with its index. Once processing is
 // finished, the index should be called with onDone to clean up the storage. If no new item is available,
 // returns false.
-func (pq *persistentQueue[T]) getNextItem(ctx context.Context) (uint64, T, bool) {
+func (pq *persistentQueue[T]) getNextItem(ctx context.Context) (uint64, T, bool, context.Context) {
 	index := pq.readIndex
 	// Increase here, so even if errors happen below, it always iterates
 	pq.readIndex++
@@ -325,8 +393,24 @@ func (pq *persistentQueue[T]) getNextItem(ctx context.Context) (uint64, T, bool)
 		getOp)
 
 	var request T
+	restoredContext := context.Background()
 	if err == nil {
-		request, err = pq.set.encoding.Unmarshal(getOp.Value)
+		var envelope marshaledRequestWithSpanContext
+		if err = json.Unmarshal(getOp.Value, &envelope); err == nil {
+			// Unmarshal the request using the specified encoding
+			if request, err = pq.set.encoding.Unmarshal(envelope.RequestBytes); err == nil {
+				// Unmarshal the SpanContext from JSON
+				var wrapper spanContextConfigWrapper
+				if len(envelope.SpanContextJSON) > 0 {
+					if err = json.Unmarshal(envelope.SpanContextJSON, &wrapper); err == nil {
+						var sc *trace.SpanContext
+						if sc, err = SpanContextFromWrapper(wrapper); err == nil && sc != nil {
+							restoredContext = trace.ContextWithSpanContext(restoredContext, *sc)
+						}
+					}
+				}
+			}
+		}
 	}
 
 	if err != nil {
@@ -336,14 +420,14 @@ func (pq *persistentQueue[T]) getNextItem(ctx context.Context) (uint64, T, bool)
 			pq.logger.Error("Error deleting item from queue", zap.Error(err))
 		}
 
-		return 0, request, false
+		return 0, request, false, restoredContext
 	}
 
 	// Increase the reference count, so the client is not closed while the request is being processed.
 	// The client cannot be closed because we hold the lock since last we checked `stopped`.
 	pq.refClient++
 
-	return index, request, true
+	return index, request, true, restoredContext
 }
 
 // onDone should be called to remove the item of the given index from the queue once processing is finished.

exporter/exporterhelper/internal/queuebatch/persistent_queue_test.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
 
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/component/componenttest"
@@ -1205,3 +1206,181 @@ func requireCurrentlyDispatchedItemsEqual(t *testing.T, pq *persistentQueue[uint
 	defer pq.mu.Unlock()
 	assert.ElementsMatch(t, compare, pq.currentlyDispatchedItems)
 }
+
+func TestSpanContextFromWrapper(t *testing.T) {
+	testCases := []struct {
+		name          string
+		wrapper       spanContextConfigWrapper
+		expectErr     bool
+		errContains   string
+		expectNil     bool
+		expectValid   bool
+		expectTraceID string
+		expectSpanID  string
+		expectFlags   string
+		expectState   string
+		expectRemote  bool
+	}{
+		{
+			name: "invalid trace id",
+			wrapper: spanContextConfigWrapper{
+				TraceID:    "invalidtraceid",
+				SpanID:     "0102030405060708",
+				TraceFlags: "01",
+				TraceState: "",
+				Remote:     false,
+			},
+			expectErr: true,
+			expectNil: true,
+		},
+		{
+			name: "invalid span id",
+			wrapper: spanContextConfigWrapper{
+				TraceID:    "0102030405060708090a0b0c0d0e0f10",
+				SpanID:     "invalidspanid",
+				TraceFlags: "01",
+				TraceState: "",
+				Remote:     false,
+			},
+			expectErr: true,
+			expectNil: true,
+		},
+		{
+			name: "invalid trace flags hex",
+			wrapper: spanContextConfigWrapper{
+				TraceID:    "0102030405060708090a0b0c0d0e0f10",
+				SpanID:     "0102030405060708",
+				TraceFlags: "zz",
+				TraceState: "",
+				Remote:     false,
+			},
+			expectErr: true,
+			expectNil: true,
+		},
+		{
+			name: "invalid trace flags length",
+			wrapper: spanContextConfigWrapper{
+				TraceID:    "0102030405060708090a0b0c0d0e0f10",
+				SpanID:     "0102030405060708",
+				TraceFlags: "0102",
+				TraceState: "",
+				Remote:     false,
+			},
+			expectErr:   true,
+			expectNil:   true,
+			errContains: errInvalidTraceFlagsLength,
+		},
+		{
+			name: "invalid trace state",
+			wrapper: spanContextConfigWrapper{
+				TraceID:    "0102030405060708090a0b0c0d0e0f10",
+				SpanID:     "0102030405060708",
+				TraceFlags: "01",
+				TraceState: "invalid=tracestate,=bad",
+				Remote:     false,
+			},
+			expectErr: true,
+			expectNil: true,
+		},
+		{
+			name: "valid span context",
+			wrapper: spanContextConfigWrapper{
+				TraceID:    "0102030405060708090a0b0c0d0e0f10",
+				SpanID:     "0102030405060708",
+				TraceFlags: "01",
+				TraceState: "vendor=value",
+				Remote:     true,
+			},
+			expectErr:     false,
+			expectNil:     false,
+			expectValid:   true,
+			expectTraceID: "0102030405060708090a0b0c0d0e0f10",
+			expectSpanID:  "0102030405060708",
+			expectFlags:   "01",
+			expectState:   "vendor=value",
+			expectRemote:  true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			scc, err := SpanContextFromWrapper(tc.wrapper)
+			if tc.expectErr {
+				require.Error(t, err)
+				if tc.errContains != "" {
+					assert.Contains(t, err.Error(), tc.errContains)
+				}
+			} else {
+				require.NoError(t, err)
+			}
+			if tc.expectNil {
+				assert.Nil(t, scc)
+			} else {
+				assert.NotNil(t, scc)
+				if tc.expectValid {
+					assert.True(t, scc.IsValid())
+					assert.Equal(t, tc.expectTraceID, scc.TraceID().String())
+					assert.Equal(t, tc.expectSpanID, scc.SpanID().String())
+					assert.Equal(t, tc.expectFlags, scc.TraceFlags().String())
+					assert.Equal(t, tc.expectState, scc.TraceState().String())
+					assert.Equal(t, tc.expectRemote, scc.IsRemote())
+				}
+			}
+		})
+	}
+}
+
+func TestPersistentQueue_SpanContextRoundTrip(t *testing.T) {
+	// Setup a minimal persistent queue using uint64Encoding and uint64
+	pq := newPersistentQueue[uint64](persistentQueueSettings[uint64]{
+		sizer:     request.RequestsSizer[uint64]{},
+		capacity:  10,
+		signal:    pipeline.SignalTraces,
+		storageID: component.ID{},
+		encoding:  uint64Encoding{},
+		id:        component.NewID(exportertest.NopType),
+		telemetry: componenttest.NewNopTelemetrySettings(),
+	}).(*persistentQueue[uint64])
+
+	ext := storagetest.NewMockStorageExtension(nil)
+	client, err := ext.GetClient(context.Background(), component.KindExporter, pq.set.id, pq.set.signal.String())
+	require.NoError(t, err)
+	pq.initClient(context.Background(), client)
+
+	// Create a valid SpanContext
+	traceID, _ := trace.TraceIDFromHex("0102030405060708090a0b0c0d0e0f10")
+	spanID, _ := trace.SpanIDFromHex("0102030405060708")
+	sc := trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    traceID,
+		SpanID:     spanID,
+		TraceFlags: 0x01,
+		TraceState: trace.TraceState{},
+		Remote:     true,
+	})
+	ctxWithSC := trace.ContextWithSpanContext(context.Background(), sc)
+
+	// Offer a request with this context
+	req := uint64(42)
+	require.NoError(t, pq.Offer(ctxWithSC, req))
+
+	// Read the request and restored context
+	restoredCtx, gotReq, _, ok := pq.Read(context.Background())
+	require.True(t, ok)
+	assert.Equal(t, req, gotReq)
+	restoredSC := trace.SpanContextFromContext(restoredCtx)
+	assert.True(t, restoredSC.IsValid())
+	assert.Equal(t, sc.TraceID(), restoredSC.TraceID())
+	assert.Equal(t, sc.SpanID(), restoredSC.SpanID())
+	assert.Equal(t, sc.TraceFlags(), restoredSC.TraceFlags())
+	assert.Equal(t, sc.TraceState().String(), restoredSC.TraceState().String())
+	assert.Equal(t, sc.IsRemote(), restoredSC.IsRemote())
+
+	// Also test with a context with no SpanContext
+	req2 := uint64(99)
+	require.NoError(t, pq.Offer(context.Background(), req2))
+	restoredCtx2, gotReq2, _, ok2 := pq.Read(context.Background())
+	require.True(t, ok2)
+	assert.Equal(t, req2, gotReq2)
+	restoredSC2 := trace.SpanContextFromContext(restoredCtx2)
+	assert.False(t, restoredSC2.IsValid())
+}