switch marshal from JSON to byte-based

Author: jackgopack4

exporter/exporterhelper/internal/queuebatch/persistent_queue.go

@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"strconv"
 	"sync"
 
@@ -243,11 +244,6 @@ 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,omitempty"`
-}
-
 type spanContextConfigWrapper struct {
 	TraceID    string
 	SpanID     string
@@ -289,31 +285,42 @@ func spanContextFromWrapper(wrapper spanContextConfigWrapper) (*trace.SpanContex
 	return &sc, nil
 }
 
-// unmarshalRequestWithSpanContext unmarshals a marshaledRequestWithSpanContext from bytes, returning the request
-// and a context with the restored SpanContext (if present).
+// unmarshalRequestWithSpanContext unmarshals a binary envelope, returning the request and a context with the restored SpanContext (if present).
 func unmarshalRequestWithSpanContext[T any](encoding Encoding[T], value []byte) (T, context.Context, error) {
 	var req T
 	restoredContext := context.Background()
-	var envelope marshaledRequestWithSpanContext
-	if err := json.Unmarshal(value, &envelope); err != nil {
-		return req, restoredContext, err
+	if len(value) < 8 {
+		return req, restoredContext, errors.New("envelope too short")
+	}
+	reqLen := binary.LittleEndian.Uint32(value[:4])
+	if len(value) < int(4+reqLen+4) {
+		return req, restoredContext, errors.New("envelope too short for request")
 	}
-	request, err := encoding.Unmarshal(envelope.RequestBytes)
+	reqBytes := value[4 : 4+reqLen]
+	scLen := binary.LittleEndian.Uint32(value[4+reqLen : 8+reqLen])
+	if len(value) < int(8+reqLen+scLen) {
+		return req, restoredContext, errors.New("envelope too short for span context")
+	}
+	scBytes := value[8+reqLen : 8+reqLen+scLen]
+	// Unmarshal request
+	r, err := encoding.Unmarshal(reqBytes)
 	if err != nil {
 		return req, restoredContext, err
 	}
-	if len(envelope.SpanContextJSON) > 0 {
+	req = r
+	// Unmarshal span context if present
+	if scLen > 0 {
 		var wrapper spanContextConfigWrapper
-		if err := json.Unmarshal(envelope.SpanContextJSON, &wrapper); err == nil {
+		if err := json.Unmarshal(scBytes, &wrapper); err == nil {
 			if sc, err := spanContextFromWrapper(wrapper); err == nil && sc != nil {
 				restoredContext = trace.ContextWithSpanContext(restoredContext, *sc)
 			}
 		}
 	}
-	return request, restoredContext, nil
+	return req, restoredContext, nil
 }
 
-// marshalRequestWithSpanContext marshals the request and the SpanContext from ctx into a marshaledRequestWithSpanContext envelope as bytes.
+// marshalRequestWithSpanContext marshals the request and the SpanContext from ctx into a binary envelope as bytes.
 func marshalRequestWithSpanContext[T any](ctx context.Context, encoding Encoding[T], req T) ([]byte, error) {
 	reqBuf, err := encoding.Marshal(req)
 	if err != nil {
@@ -326,14 +333,22 @@ func marshalRequestWithSpanContext[T any](ctx context.Context, encoding Encoding
 		if err != nil {
 			return nil, err
 		}
-	} else {
-		scJSON = nil // Will be omitted due to omitempty
-	}
-	envelope := marshaledRequestWithSpanContext{
-		RequestBytes:    reqBuf,
-		SpanContextJSON: scJSON,
 	}
-	return json.Marshal(envelope)
+	if len(reqBuf) > int(math.MaxUint32) {
+		return nil, fmt.Errorf("request too large to encode: %d bytes", len(reqBuf))
+	}
+	if len(scJSON) > int(math.MaxUint32) {
+		return nil, fmt.Errorf("span context too large to encode: %d bytes", len(scJSON))
+	}
+	// Compose binary envelope: [4 bytes reqLen][req][4 bytes scLen][scJSON]
+	buf := make([]byte, 0, 8+len(reqBuf)+len(scJSON))
+	//nolint:gosec // G115: integer overflow conversion int -> uint32
+	buf = binary.LittleEndian.AppendUint32(buf, uint32(len(reqBuf)))
+	buf = append(buf, reqBuf...)
+	//nolint:gosec // G115: integer overflow conversion int -> uint32
+	buf = binary.LittleEndian.AppendUint32(buf, uint32(len(scJSON)))
+	buf = append(buf, scJSON...)
+	return buf, nil
 }
 
 // putInternal is the internal version that requires caller to hold the mutex lock.

exporter/exporterhelper/internal/queuebatch/persistent_queue_test.go

@@ -6,7 +6,6 @@ package queuebatch
 import (
 	"context"
 	"encoding/binary"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"math"
@@ -1441,28 +1440,34 @@ func TestMarshalUnmarshalRequestWithSpanContext(t *testing.T) {
 	})
 
 	t.Run("corrupted span_context field", func(t *testing.T) {
-		// Manually create a bad envelope
-		envelope := marshaledRequestWithSpanContext{
-			RequestBytes:    []byte{0x2a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
-			SpanContextJSON: []byte(`{"TraceID":123}`), // invalid TraceID
-		}
-		data, err := json.Marshal(envelope)
-		require.NoError(t, err)
-		_, gotCtx, err := unmarshalRequestWithSpanContext(uint64Encoding{}, data)
+		// Manually create a bad envelope using the new binary format
+		requestBytes := []byte{0x2a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
+		badSpanContext := []byte(`{"TraceID":123}`) // invalid TraceID
+		buf := make([]byte, 0, 8+len(requestBytes)+len(badSpanContext))
+		//nolint:gosec // G115: integer overflow conversion int -> uint32
+		buf = binary.LittleEndian.AppendUint32(buf, uint32(len(requestBytes)))
+		buf = append(buf, requestBytes...)
+		//nolint:gosec // G115: integer overflow conversion int -> uint32
+		buf = binary.LittleEndian.AppendUint32(buf, uint32(len(badSpanContext)))
+		buf = append(buf, badSpanContext...)
+		_, gotCtx, err := unmarshalRequestWithSpanContext(uint64Encoding{}, buf)
 		require.NoError(t, err)
 		// Should not panic, should return background context
 		restoredSC := trace.SpanContextFromContext(gotCtx)
 		assert.False(t, restoredSC.IsValid())
 	})
 
-	t.Run("valid JSON, invalid RequestBytes returns error", func(t *testing.T) {
-		envelope := marshaledRequestWithSpanContext{
-			RequestBytes:    []byte{0x01, 0x02}, // too short for uint64Encoding.Unmarshal
-			SpanContextJSON: nil,
-		}
-		data, err := json.Marshal(envelope)
-		require.NoError(t, err)
-		_, _, err = unmarshalRequestWithSpanContext(uint64Encoding{}, data)
+	t.Run("valid binary, invalid RequestBytes returns error", func(t *testing.T) {
+		requestBytes := []byte{0x01, 0x02} // too short for uint64Encoding.Unmarshal
+		badSpanContext := []byte{}
+		buf := make([]byte, 0, 8+len(requestBytes)+len(badSpanContext))
+		//nolint:gosec // G115: integer overflow conversion int -> uint32
+		buf = binary.LittleEndian.AppendUint32(buf, uint32(len(requestBytes)))
+		buf = append(buf, requestBytes...)
+		//nolint:gosec // G115: integer overflow conversion int -> uint32
+		buf = binary.LittleEndian.AppendUint32(buf, uint32(len(badSpanContext)))
+		buf = append(buf, badSpanContext...)
+		_, _, err := unmarshalRequestWithSpanContext(uint64Encoding{}, buf)
 		require.Error(t, err)
 	})
 }