[exporterhelper] Context unmarshaling along with the request

Author: dmitryax

exporter/exporterhelper/internal/base_exporter_test.go

@@ -162,8 +162,15 @@ func newFakeQueueBatch() QueueBatchSettings[request.Request] {
 
 type fakeEncoding struct{}
 
-func (f fakeEncoding) Marshal(request.Request) ([]byte, error) {
-	return []byte("mockRequest"), nil
+func (f fakeEncoding) MarshalTo(_ request.Request, b []byte) (int, error) {
+	for i, ch := range "mockRequest" {
+		b[i] = byte(ch)
+	}
+	return len("mockRequest"), nil
+}
+
+func (f fakeEncoding) MarshalSize(request.Request) int {
+	return len("mockRequest")
 }
 
 func (f fakeEncoding) Unmarshal([]byte) (request.Request, error) {

exporter/exporterhelper/internal/persistentqueue/encoding.go

@@ -0,0 +1,239 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package persistentqueue // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/persistentqueue"
+
+import (
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"math"
+	"slices"
+
+	"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 {
+	// MarshalTo marshals a request into a preallocated byte slice.
+	// The size of the byte slice must be at least MarshalSize(T) bytes.
+	MarshalTo(T, []byte) (int, error)
+
+	// MarshalSize returns the size of the marshaled request.
+	MarshalSize(T) int
+
+	// Unmarshal unmarshals 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,
+	}
+}
+
+const (
+	// requestDataKey is the key used to store request data in bytesMap.
+	requestDataKey = "request_data"
+
+	// traceContextBufCap is a capacity of the bytesMap buffer that is enough for trace context key-value pairs.
+	traceContextBufCap = 128
+)
+
+var tracePropagator = propagation.TraceContext{}
+
+func (re Encoder[T]) Marshal(ctx context.Context, req T) ([]byte, error) {
+	reqSize := re.encoding.MarshalSize(req)
+
+	if !persistRequestContextFeatureGate.IsEnabled() {
+		b := make([]byte, reqSize)
+		_, err := re.encoding.MarshalTo(req, b)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal request: %w", err)
+		}
+		return b, nil
+	}
+
+	bm := newBytesMap(traceContextBufCap + reqSize)
+	tracePropagator.Inject(ctx, &bytesMapCarrier{bytesMap: *bm})
+	reqBuf, err := bm.setEmptyBytes(requestDataKey, reqSize)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request: %w", err)
+	}
+	_, err = re.encoding.MarshalTo(req, reqBuf)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request: %w", err)
+	}
+
+	return reqBuf, 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 := context.Background()
+	tracePropagator.Extract(ctx, &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.
+const (
+	magicByte      = byte(0x00)
+	formatV1Byte   = byte(0x01)
+	prefixBytesLen = 2
+)
+
+func newBytesMap(initSize int) *bytesMap {
+	bm := bytesMap(make([]byte, 0, prefixBytesLen+initSize))
+	bm = append(bm, magicByte, formatV1Byte)
+	return &bm
+}
+
+// setEmptyBytes sets the specified key in the map, reserves the given number of bytes for the value,
+// and returns a byte slice for the value. Must be called only once for each key.
+func (bm *bytesMap) setEmptyBytes(key string, size int) ([]byte, error) {
+	if len(key) > math.MaxUint8 {
+		return nil, errors.New("key param is too long")
+	}
+	if size < 0 || uint64(size) > math.MaxUint32 {
+		return nil, fmt.Errorf("invalid value size %d", size)
+	}
+
+	*bm = append(*bm, byte(len(key)))
+	*bm = append(*bm, key...)
+
+	var lenBuf [4]byte
+	binary.LittleEndian.PutUint32(lenBuf[:], uint32(size)) //nolint:gosec // disable G115
+	*bm = append(*bm, lenBuf[:]...)
+
+	start := len(*bm)
+	*bm = slices.Grow(*bm, size)
+	*bm = []byte(*bm)[:start+size]
+
+	return []byte(*bm)[start:], 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) {
+	buf, _ := c.setEmptyBytes(k, len(v))
+	copy(buf, 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/persistentqueue/encoding_test.go

@@ -0,0 +1,50 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package persistentqueue // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/persistentqueue"
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+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(0)
+	for _, d := range data {
+		buf, err := bm.setEmptyBytes(d.key, len(d.val))
+		require.NoError(t, err)
+		copy(buf, d.val)
+	}
+
+	assert.Equal(t, []string{"key1", "key2", "key3", "key4"}, bm.keys())
+
+	buf, err := bm.get("key2")
+	assert.NoError(t, err)
+	assert.Equal(t, []byte("value2"), buf)
+	buf, err = bm.get("key4")
+	assert.NoError(t, err)
+	assert.Equal(t, []byte("value4"), buf)
+}
+
+func TestBytesMapCarrier(t *testing.T) {
+	bm := newBytesMap(0)
+	carrier := &bytesMapCarrier{bytesMap: *bm}
+
+	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"))
+}

exporter/exporterhelper/internal/queuebatch/internal/persistentqueue/meta.pb.go renamed to exporter/exporterhelper/internal/persistentqueue/meta.pb.go

@@ -2,7 +2,7 @@ syntax = "proto3";
 
 package opentelemetry.collector.exporter.exporterhelper.internal.queuebatch.internal.persistentqueue;
 
-option go_package = "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch/internal/persistentqueue";
+option go_package = "go.opentelemetry.io/collector/exporter/exporterhelper/internal/persistentqueue";
 
 // Sizer type configuration
 enum SizerType {

exporter/exporterhelper/internal/queuebatch/internal/persistentqueue/meta.proto renamed to exporter/exporterhelper/internal/persistentqueue/meta.proto

@@ -1,7 +1,7 @@
 // Copyright The OpenTelemetry Authors
 // SPDX-License-Identifier: Apache-2.0
 
-package persistentqueue // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch/internal/persistentqueue"
+package persistentqueue // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/persistentqueue"
 
 import (
 	"fmt"

exporter/exporterhelper/internal/queuebatch/internal/persistentqueue/sizer_conversion.go renamed to exporter/exporterhelper/internal/persistentqueue/sizer_conversion.go

@@ -13,14 +13,15 @@ import (
 
 	"go.uber.org/zap"
 
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/persistentqueue"
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch"
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/request"
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sender"
 )
 
 // QueueBatchSettings is a subset of the queuebatch.Settings that are needed when used within an Exporter.
 type QueueBatchSettings[T any] struct {
-	Encoding    queuebatch.Encoding[T]
+	Encoding    persistentqueue.Encoding[T]
 	Sizers      map[request.SizerType]request.Sizer[T]
 	Partitioner queuebatch.Partitioner[T]
 }