tailsampler: Combine batches of spans into a single batch

processor/samplingprocessor/tailsamplingprocessor/processor.go

@@ -167,9 +167,16 @@ func (tsp *tailSamplingSpanProcessor) samplingPolicyOnTick() {
 		trace.Unlock()
 
 		if decision == sampling.Sampled {
+
+			// Combine all individual batches into a single batch so
+			// consumers may operate on the entire trace
+			allSpans := pdata.NewTraces()
 			for j := 0; j < len(traceBatches); j++ {
-				_ = tsp.nextConsumer.ConsumeTraces(policy.ctx, internaldata.OCToTraceData(traceBatches[j]))
+				batch := internaldata.OCToTraceData(traceBatches[j])
+				batch.ResourceSpans().MoveAndAppendTo(allSpans.ResourceSpans())
 			}
+
+			_ = tsp.nextConsumer.ConsumeTraces(policy.ctx, allSpans)
 		}
 	}
 

processor/samplingprocessor/tailsamplingprocessor/processor_test.go

@@ -15,8 +15,10 @@
 package tailsamplingprocessor
 
 import (
+	"bytes"
 	"context"
 	"errors"
+	"sort"
 	"sync"
 	"testing"
 	"time"
@@ -331,8 +333,103 @@ func TestSamplingPolicyDecisionNotSampled(t *testing.T) {
 	require.Equal(t, 2, mpe.LateArrivingSpansCount, "policy was not notified of the late span")
 }
 
+func TestMultipleBatchesAreCombinedIntoOne(t *testing.T) {
+	const maxSize = 100
+	const decisionWaitSeconds = 1
+	// For this test explicitly control the timer calls and batcher, and set a mock
+	// sampling policy evaluator.
+	msp := new(exportertest.SinkTraceExporter)
+	mpe := &mockPolicyEvaluator{}
+	mtt := &manualTTicker{}
+	tsp := &tailSamplingSpanProcessor{
+		ctx:             context.Background(),
+		nextConsumer:    msp,
+		maxNumTraces:    maxSize,
+		logger:          zap.NewNop(),
+		decisionBatcher: newSyncIDBatcher(decisionWaitSeconds),
+		policies:        []*Policy{{Name: "mock-policy", Evaluator: mpe, ctx: context.TODO()}},
+		deleteChan:      make(chan traceKey, maxSize),
+		policyTicker:    mtt,
+	}
+
+	mpe.NextDecision = sampling.Sampled
+
+	traceIds, batches := generateIdsAndBatches(3)
+	for _, batch := range batches {
+		require.NoError(t, tsp.ConsumeTraces(context.Background(), batch))
+	}
+
+	tsp.samplingPolicyOnTick()
+	tsp.samplingPolicyOnTick()
+
+	require.EqualValues(t, 3, len(msp.AllTraces()), "There should be three batches, one for each trace")
+
+	expectedSpanIds := make(map[int][]pdata.SpanID)
+	expectedSpanIds[0] = []pdata.SpanID{
+		pdata.NewSpanID(tracetranslator.UInt64ToByteSpanID(uint64(1))),
+	}
+	expectedSpanIds[1] = []pdata.SpanID{
+		pdata.NewSpanID(tracetranslator.UInt64ToByteSpanID(uint64(2))),
+		pdata.NewSpanID(tracetranslator.UInt64ToByteSpanID(uint64(3))),
+	}
+	expectedSpanIds[2] = []pdata.SpanID{
+		pdata.NewSpanID(tracetranslator.UInt64ToByteSpanID(uint64(4))),
+		pdata.NewSpanID(tracetranslator.UInt64ToByteSpanID(uint64(5))),
+		pdata.NewSpanID(tracetranslator.UInt64ToByteSpanID(uint64(6))),
+	}
+
+	receivedTraces := msp.AllTraces()
+	for i, traceID := range traceIds {
+		trace := findTrace(receivedTraces, traceID)
+		require.NotNil(t, trace, "Trace was not received. TraceId %s", traceID.HexString())
+		require.EqualValues(t, i+1, trace.SpanCount(), "The trace should have all of its spans in a single batch")
+
+		expected := expectedSpanIds[i]
+		got := collectSpanIds(trace)
+
+		// might have received out of order, sort for comparison
+		sort.Slice(got, func(i, j int) bool {
+			a, _ := tracetranslator.BytesToInt64SpanID(got[i])
+			b, _ := tracetranslator.BytesToInt64SpanID(got[j])
+			return a < b
+		})
+
+		require.EqualValues(t, expected, got)
+	}
+}
+
+func collectSpanIds(trace *pdata.Traces) []pdata.SpanID {
+	spanIDs := make([]pdata.SpanID, 0)
+
+	for i := 0; i < trace.ResourceSpans().Len(); i++ {
+		ilss := trace.ResourceSpans().At(i).InstrumentationLibrarySpans()
+
+		for j := 0; j < ilss.Len(); j++ {
+			ils := ilss.At(j)
+
+			for k := 0; k < ils.Spans().Len(); k++ {
+				span := ils.Spans().At(k)
+				spanIDs = append(spanIDs, span.SpanID())
+			}
+		}
+	}
+
+	return spanIDs
+}
+
+func findTrace(a []pdata.Traces, traceID pdata.TraceID) *pdata.Traces {
+	for _, batch := range a {
+		id := batch.ResourceSpans().At(0).InstrumentationLibrarySpans().At(0).Spans().At(0).TraceID()
+		if bytes.Equal(traceID.Bytes(), id.Bytes()) {
+			return &batch
+		}
+	}
+	return nil
+}
+
 func generateIdsAndBatches(numIds int) ([]pdata.TraceID, []pdata.Traces) {
 	traceIds := make([]pdata.TraceID, numIds)
+	spanID := 0
 	var tds []pdata.Traces
 	for i := 0; i < numIds; i++ {
 		traceIds[i] = tracetranslator.UInt64ToTraceID(1, uint64(i+1))
@@ -341,7 +438,9 @@ func generateIdsAndBatches(numIds int) ([]pdata.TraceID, []pdata.Traces) {
 			td := testdata.GenerateTraceDataOneSpan()
 			span := td.ResourceSpans().At(0).InstrumentationLibrarySpans().At(0).Spans().At(0)
 			span.SetTraceID(traceIds[i])
-			span.SetSpanID(tracetranslator.UInt64ToByteSpanID(uint64(i + 1)))
+
+			spanID++
+			span.SetSpanID(tracetranslator.UInt64ToByteSpanID(uint64(spanID)))
 			tds = append(tds, td)
 		}
 	}