[exporterhelper] Preserve request span context in the persistent queue

Author: dmitryax

.chloggen/persist-request-context.yaml

@@ -0,0 +1,27 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: enhancement
+
+# The name of the component, or a single word describing the area of concern, (e.g. otlpreceiver)
+component: exporterhelper
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: Add an option to preserve request span context in the persistent queue
+
+# One or more tracking issues or pull requests related to the change
+issues: []
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext: |
+  Currently, it's only available under "exporter.PersistRequestContext" feature gate which can be enabled by adding
+  `--feature-gates=exporter.PersistRequestContext` to the collector command line.
+
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: [user]

exporter/exporterhelper/internal/queuebatch/internal/persistentqueue/encoding.go

@@ -0,0 +1,217 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package persistentqueue // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch/internal/persistentqueue"
+
+import (
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"math"
+
+	"go.opentelemetry.io/otel/propagation"
+
+	"go.opentelemetry.io/collector/featuregate"
+)
+
+// persistRequestContextFeatureGate controls whether request context should be preserved in the persistent queue.
+var persistRequestContextFeatureGate = featuregate.GlobalRegistry().MustRegister(
+	"exporter.PersistRequestContext",
+	featuregate.StageAlpha,
+	featuregate.WithRegisterFromVersion("v0.128.0"),
+	featuregate.WithRegisterDescription("controls whether context should be stored alongside requests in the persistent queue"),
+)
+
+type Encoding[T any] interface {
+	// Marshal is a function that can marshal a request into bytes.
+	Marshal(T) ([]byte, error)
+
+	// Unmarshal is a function that can unmarshal bytes into a request.
+	Unmarshal([]byte) (T, error)
+}
+
+// Encoder provides an interface for marshaling and unmarshaling requests along with their context.
+type Encoder[T any] struct {
+	encoding Encoding[T]
+}
+
+func NewEncoder[T any](encoding Encoding[T]) Encoder[T] {
+	return Encoder[T]{
+		encoding: encoding,
+	}
+}
+
+// requestDataKey is the key used to store request data in bytesMap.
+const requestDataKey = "req"
+
+var tracePropagator = propagation.TraceContext{}
+
+func (re Encoder[T]) Marshal(ctx context.Context, req T) ([]byte, error) {
+	if !persistRequestContextFeatureGate.IsEnabled() {
+		return re.encoding.Marshal(req)
+	}
+
+	bm := newBytesMap()
+	tracePropagator.Inject(ctx, &bytesMapCarrier{bytesMap: bm})
+	reqBuf, err := re.encoding.Marshal(req)
+	if err != nil {
+		return nil, err
+	}
+	if err := bm.set(requestDataKey, reqBuf); err != nil {
+		return nil, fmt.Errorf("failed to marshal request: %w", err)
+	}
+
+	return *bm, nil
+}
+
+func (re Encoder[T]) Unmarshal(b []byte) (T, context.Context, error) {
+	if !persistRequestContextFeatureGate.IsEnabled() {
+		req, err := re.encoding.Unmarshal(b)
+		return req, context.Background(), err
+	}
+
+	bm := bytesMapFromBytes(b)
+	if bm == nil {
+		// Fall back to unmarshalling of the request alone.
+		// This can happen if the data persisted by the version that doesn't support the context unmarshaling.
+		req, err := re.encoding.Unmarshal(b)
+		return req, context.Background(), err
+	}
+	ctx := tracePropagator.Extract(context.Background(), &bytesMapCarrier{bytesMap: bm})
+	reqBuf, err := bm.get(requestDataKey)
+	var req T
+	if err != nil {
+		return req, context.Background(), fmt.Errorf("failed to read serialized request data: %w", err)
+	}
+	req, err = re.encoding.Unmarshal(reqBuf)
+	return req, ctx, err
+}
+
+// bytesMap is a slice of bytes that represents a map-like structure for storing key-value pairs.
+// It's optimized for efficient memory usage for low number of key-value pairs with big values.
+// The format is a sequence of key-value pairs encoded as:
+//   - 1 byte length of the key
+//   - key bytes
+//   - 4 byte length of the value
+//   - value bytes
+type bytesMap []byte
+
+// prefix bytes to denote the bytesMap serialization: 0x00 magic byte + 0x01 version of the encoder.
+const (
+	magicByte       = byte(0x00)
+	formatV1Byte    = byte(0x01)
+	prefixBytesLen  = 2
+	initialCapacity = 256
+)
+
+func newBytesMap() *bytesMap {
+	bm := bytesMap(make([]byte, 0, initialCapacity))
+	bm = append(bm, magicByte, formatV1Byte)
+	return &bm
+}
+
+// set sets the specified key in the map. Must be called only once for each key.
+func (bm *bytesMap) set(key string, val []byte) error {
+	if len(key) > math.MaxUint8 {
+		return errors.New("key param is too long")
+	}
+	valSize := len(val)
+	if uint64(valSize) > math.MaxUint32 {
+		return fmt.Errorf("value is too large to persist, size %d", valSize)
+	}
+
+	*bm = append(*bm, byte(len(key)))
+	*bm = append(*bm, key...)
+
+	var lenBuf [4]byte
+	binary.LittleEndian.PutUint32(lenBuf[:], uint32(valSize)) //nolint:gosec // disable G115
+	*bm = append(*bm, lenBuf[:]...)
+	*bm = append(*bm, val...)
+
+	return nil
+}
+
+// get scans sequentially for the first matching key and returns the value as bytes.
+func (bm *bytesMap) get(k string) ([]byte, error) {
+	for i := prefixBytesLen; i < len(*bm); {
+		kl := int([]byte(*bm)[i])
+		i++
+
+		if i+kl > len(*bm) {
+			return nil, io.ErrUnexpectedEOF
+		}
+		key := string([]byte(*bm)[i : i+kl])
+		i += kl
+
+		if i+4 > len(*bm) {
+			return nil, io.ErrUnexpectedEOF
+		}
+		vLen := binary.LittleEndian.Uint32([]byte(*bm)[i:])
+		i += 4
+
+		if i+int(vLen) > len(*bm) {
+			return nil, io.ErrUnexpectedEOF
+		}
+		val := []byte(*bm)[i : i+int(vLen)]
+		i += int(vLen)
+
+		if key == k {
+			return val, nil
+		}
+	}
+	return nil, nil
+}
+
+// keys returns header names in encounter order.
+func (bm *bytesMap) keys() []string {
+	var out []string
+	for i := prefixBytesLen; i < len(*bm); {
+		kl := int([]byte(*bm)[i])
+		i++
+
+		if i+kl > len(*bm) {
+			break // malformed entry
+		}
+		out = append(out, string([]byte(*bm)[i:i+kl]))
+		i += kl
+
+		if i+4 > len(*bm) {
+			break // malformed entry
+		}
+		vLen := binary.LittleEndian.Uint32([]byte(*bm)[i:])
+		i += 4 + int(vLen)
+	}
+	return out
+}
+
+func bytesMapFromBytes(b []byte) *bytesMap {
+	if len(b) < prefixBytesLen || b[0] != magicByte || b[1] != formatV1Byte {
+		return nil
+	}
+	return (*bytesMap)(&b)
+}
+
+// bytesMapCarrier implements propagation.TextMapCarrier on top of bytesMap.
+type bytesMapCarrier struct {
+	*bytesMap
+}
+
+var _ propagation.TextMapCarrier = (*bytesMapCarrier)(nil)
+
+// Set appends a new string entry; if the key already exists it is left unchanged.
+func (c *bytesMapCarrier) Set(k, v string) {
+	_ = c.set(k, []byte(v))
+}
+
+// Get scans sequentially for the first matching key.
+func (c *bytesMapCarrier) Get(k string) string {
+	v, _ := c.get(k)
+	return string(v)
+}
+
+// Keys returns header names in encounter order.
+func (c *bytesMapCarrier) Keys() []string {
+	return c.keys()
+}

exporter/exporterhelper/internal/queuebatch/internal/persistentqueue/encoding_test.go

@@ -0,0 +1,118 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package persistentqueue // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch/internal/persistentqueue"
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
+
+	"go.opentelemetry.io/collector/featuregate"
+)
+
+func TestBytesMap(t *testing.T) {
+	data := []struct {
+		key string
+		val []byte
+	}{
+		{"key1", []byte("value1")},
+		{"key2", []byte("value2")},
+		{"key3", []byte("value3")},
+		{"key4", []byte("value4")},
+	}
+
+	bm := newBytesMap()
+	for _, d := range data {
+		err := bm.set(d.key, d.val)
+		require.NoError(t, err)
+	}
+
+	assert.Equal(t, []string{"key1", "key2", "key3", "key4"}, bm.keys())
+
+	buf, err := bm.get("key2")
+	require.NoError(t, err)
+	assert.Equal(t, []byte("value2"), buf)
+	buf, err = bm.get("key4")
+	require.NoError(t, err)
+	assert.Equal(t, []byte("value4"), buf)
+}
+
+func TestBytesMapCarrier(t *testing.T) {
+	carrier := &bytesMapCarrier{bytesMap: newBytesMap()}
+
+	carrier.Set("key1", "val1")
+	carrier.Set("key2", "val2")
+
+	assert.Equal(t, []string{"key1", "key2"}, carrier.Keys())
+	assert.Equal(t, "val2", carrier.Get("key2"))
+	assert.Equal(t, "val1", carrier.Get("key1"))
+}
+
+func TestEncoder(t *testing.T) {
+	traceID, err := trace.TraceIDFromHex("0102030405060708090a0b0c0d0e0f10")
+	require.NoError(t, err)
+	spanID, err := trace.SpanIDFromHex("0102030405060708")
+	require.NoError(t, err)
+	spanCtx := trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    traceID,
+		SpanID:     spanID,
+		TraceFlags: 0x01,
+		TraceState: trace.TraceState{},
+		Remote:     true,
+	})
+	ctx := trace.ContextWithSpanContext(context.Background(), spanCtx)
+	req := uint64(123)
+	enc := NewEncoder(Uint64Encoding{})
+	fgOrigState := persistRequestContextFeatureGate.IsEnabled()
+
+	tests := []struct {
+		name             string
+		fgEnabledOnWrite bool
+		fgEnabledOnRead  bool
+		wantSpanCtx      trace.SpanContext
+		wantReadErr      error
+	}{
+		{
+			name: "feature_gate_disabled_on_write_and_read",
+		},
+		{
+			name:             "feature_gate_enabled_on_write_and_read",
+			fgEnabledOnWrite: true,
+			fgEnabledOnRead:  true,
+			wantSpanCtx:      spanCtx,
+		},
+		{
+			name:            "feature_gate_disabled_on_write_enabled_on_read",
+			fgEnabledOnRead: true,
+		},
+		{
+			name:             "feature_gate_enabled_on_write_disabled_on_read",
+			fgEnabledOnWrite: true,
+			wantReadErr:      ErrInvalidValue,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.NoError(t, featuregate.GlobalRegistry().Set(persistRequestContextFeatureGate.ID(), tt.fgEnabledOnWrite))
+			defer func() {
+				require.NoError(t, featuregate.GlobalRegistry().Set(persistRequestContextFeatureGate.ID(), fgOrigState))
+			}()
+			buf, err := enc.Marshal(ctx, req)
+			require.NoError(t, err)
+
+			require.NoError(t, featuregate.GlobalRegistry().Set(persistRequestContextFeatureGate.ID(), tt.fgEnabledOnRead))
+			gotReq, gotCtx, err := enc.Unmarshal(buf)
+			assert.Equal(t, tt.wantReadErr, err)
+			if err == nil {
+				assert.Equal(t, req, gotReq)
+				gotSpanCtx := trace.SpanContextFromContext(gotCtx)
+				assert.Equal(t, tt.wantSpanCtx, gotSpanCtx)
+			}
+		})
+	}
+}

exporter/exporterhelper/internal/queuebatch/internal/persistentqueue/encodingtest.go

@@ -0,0 +1,26 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package persistentqueue // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch/internal/persistentqueue"
+
+import (
+	"encoding/binary"
+	"errors"
+)
+
+var ErrInvalidValue = errors.New("invalid value")
+
+type Uint64Encoding struct{}
+
+var _ Encoding[uint64] = Uint64Encoding{}
+
+func (Uint64Encoding) Marshal(val uint64) ([]byte, error) {
+	return binary.LittleEndian.AppendUint64([]byte{}, val), nil
+}
+
+func (Uint64Encoding) Unmarshal(bytes []byte) (uint64, error) {
+	if len(bytes) != 8 {
+		return 0, ErrInvalidValue
+	}
+	return binary.LittleEndian.Uint64(bytes), nil
+}