feat: Allow partial traces to be returned

cmd/tempo-cli/cmd-query-blocks.go

@@ -51,7 +51,7 @@ func (cmd *queryBlocksCmd) Run(ctx *globalOptions) error {
 	}
 
 	var (
-		combiner   = trace.NewCombiner(0)
+		combiner   = trace.NewCombiner(0, true)
 		marshaller = new(jsonpb.Marshaler)
 		jsonBytes  = bytes.Buffer{}
 	)

modules/frontend/combiner/trace_by_id.go

@@ -35,7 +35,7 @@ type traceByIDCombiner struct {
 // - encode the returned trace as either json or proto depending on the request
 func NewTraceByID(maxBytes int, contentType string) Combiner {
 	return &traceByIDCombiner{
-		c:           trace.NewCombiner(maxBytes),
+		c:           trace.NewCombiner(maxBytes, false),
 		code:        http.StatusNotFound,
 		contentType: contentType,
 	}

modules/frontend/combiner/trace_by_id_v2.go

@@ -6,7 +6,7 @@ import (
 )
 
 func NewTraceByIDV2(maxBytes int, marshalingFormat string) Combiner {
-	combiner := trace.NewCombiner(maxBytes)
+	combiner := trace.NewCombiner(maxBytes, true)
 	gc := &genericCombiner[*tempopb.TraceByIDResponse]{
 		combine: func(partial *tempopb.TraceByIDResponse, _ *tempopb.TraceByIDResponse, _ PipelineResponse) error {
 			_, err := combiner.Consume(partial.Trace)

modules/ingester/ingester.go

@@ -286,7 +286,7 @@ func (i *Ingester) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDRequ
 		return &tempopb.TraceByIDResponse{}, nil
 	}
 
-	trace, err := inst.FindTraceByID(ctx, req.TraceID)
+	trace, err := inst.FindTraceByID(ctx, req.TraceID, req.AllowPartialTrace)
 	if err != nil {
 		return nil, err
 	}

modules/ingester/instance.go

@@ -176,7 +176,7 @@ func (i *instance) addTraceError(errorsByTrace []tempopb.PushErrorReason, pushEr
 		errorsByTrace = append(errorsByTrace, tempopb.PushErrorReason_UNKNOWN_ERROR)
 		return errorsByTrace
 
-	} else if pushError == nil && len(errorsByTrace) > 0 {
+	} else if len(errorsByTrace) > 0 {
 		errorsByTrace = append(errorsByTrace, tempopb.PushErrorReason_NO_ERROR)
 	}
 
@@ -405,7 +405,7 @@ func (i *instance) ClearFlushedBlocks(completeBlockTimeout time.Duration) error
 	return err
 }
 
-func (i *instance) FindTraceByID(ctx context.Context, id []byte) (*tempopb.Trace, error) {
+func (i *instance) FindTraceByID(ctx context.Context, id []byte, allowPartialTrace bool) (*tempopb.Trace, error) {
 	span, ctx := opentracing.StartSpanFromContext(ctx, "instance.FindTraceByID")
 	defer span.Finish()
 
@@ -426,7 +426,7 @@ func (i *instance) FindTraceByID(ctx context.Context, id []byte) (*tempopb.Trace
 	maxBytes := i.limiter.limits.MaxBytesPerTrace(i.instanceID)
 	searchOpts := common.DefaultSearchOptionsWithMaxBytes(maxBytes)
 
-	combiner := trace.NewCombiner(maxBytes)
+	combiner := trace.NewCombiner(maxBytes, allowPartialTrace)
 	_, err = combiner.Consume(completeTrace)
 	if err != nil {
 		return nil, err

modules/ingester/instance_search_test.go

@@ -664,7 +664,7 @@ func TestInstanceSearchDoesNotRace(t *testing.T) {
 	})
 
 	go concurrent(func() {
-		_, err := i.FindTraceByID(context.Background(), []byte{0x01})
+		_, err := i.FindTraceByID(context.Background(), []byte{0x01}, false)
 		assert.NoError(t, err, "error finding trace by id")
 	})
 

modules/ingester/instance_test.go

@@ -150,7 +150,7 @@ func pushTracesToInstance(t *testing.T, i *instance, numTraces int) ([]*tempopb.
 
 func queryAll(t *testing.T, i *instance, ids [][]byte, traces []*tempopb.Trace) {
 	for j, id := range ids {
-		trace, err := i.FindTraceByID(context.Background(), id)
+		trace, err := i.FindTraceByID(context.Background(), id, false)
 		require.NoError(t, err)
 		require.Equal(t, traces[j], trace)
 	}
@@ -200,7 +200,7 @@ func TestInstanceDoesNotRace(t *testing.T) {
 	})
 
 	go concurrent(func() {
-		_, err := i.FindTraceByID(context.Background(), []byte{0x01})
+		_, err := i.FindTraceByID(context.Background(), []byte{0x01}, false)
 		require.NoError(t, err, "error finding trace by id")
 	})
 
@@ -625,25 +625,25 @@ func TestInstancePartialSuccess(t *testing.T) {
 	assert.Equal(t, true, traceTooLargeCount > 0)
 
 	// check that the two good ones actually made it
-	result, err := i.FindTraceByID(ctx, ids[0])
+	result, err := i.FindTraceByID(ctx, ids[0], false)
 	require.NoError(t, err, "error finding trace by id")
 	assert.Equal(t, 1, len(result.ResourceSpans))
 
-	result, err = i.FindTraceByID(ctx, ids[3])
+	result, err = i.FindTraceByID(ctx, ids[3], false)
 	require.NoError(t, err, "error finding trace by id")
 	assert.Equal(t, 1, len(result.ResourceSpans))
 
 	// check that the three traces that had errors did not actually make it
 	var expected *tempopb.Trace
-	result, err = i.FindTraceByID(ctx, ids[1])
+	result, err = i.FindTraceByID(ctx, ids[1], false)
 	require.NoError(t, err, "error finding trace by id")
 	assert.Equal(t, expected, result)
 
-	result, err = i.FindTraceByID(ctx, ids[2])
+	result, err = i.FindTraceByID(ctx, ids[2], false)
 	require.NoError(t, err, "error finding trace by id")
 	assert.Equal(t, expected, result)
 
-	result, err = i.FindTraceByID(ctx, ids[4])
+	result, err = i.FindTraceByID(ctx, ids[4], false)
 	require.NoError(t, err, "error finding trace by id")
 	assert.Equal(t, expected, result)
 }
@@ -745,7 +745,7 @@ func BenchmarkInstanceFindTraceByIDFromCompleteBlock(b *testing.B) {
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		trace, err := instance.FindTraceByID(context.Background(), traceID)
+		trace, err := instance.FindTraceByID(context.Background(), traceID, false)
 		require.NotNil(b, trace)
 		require.NoError(b, err)
 	}
@@ -899,7 +899,7 @@ func BenchmarkInstanceContention(t *testing.B) {
 	})
 
 	go concurrent(func() {
-		_, err := i.FindTraceByID(ctx, []byte{0x01})
+		_, err := i.FindTraceByID(ctx, []byte{0x01}, false)
 		require.NoError(t, err, "error finding trace by id")
 		finds++
 	})

modules/querier/http.go

@@ -63,7 +63,7 @@ func (q *Querier) TraceByIDHandler(w http.ResponseWriter, r *http.Request) {
 		BlockStart: blockStart,
 		BlockEnd:   blockEnd,
 		QueryMode:  queryMode,
-	}, timeStart, timeEnd)
+	}, timeStart, timeEnd, false)
 	if err != nil {
 		handleError(w, err)
 		return
@@ -112,7 +112,7 @@ func (q *Querier) TraceByIDHandlerV2(w http.ResponseWriter, r *http.Request) {
 		BlockStart: blockStart,
 		BlockEnd:   blockEnd,
 		QueryMode:  queryMode,
-	}, timeStart, timeEnd)
+	}, timeStart, timeEnd, true)
 	if err != nil {
 		handleError(w, err)
 		return

modules/querier/querier.go

@@ -221,7 +221,7 @@ func (q *Querier) stopping(_ error) error {
 }
 
 // FindTraceByID implements tempopb.Querier.
-func (q *Querier) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDRequest, timeStart int64, timeEnd int64) (*tempopb.TraceByIDResponse, error) {
+func (q *Querier) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDRequest, timeStart int64, timeEnd int64, allowPartialTrace bool) (*tempopb.TraceByIDResponse, error) {
 	if !validation.ValidTraceID(req.TraceID) {
 		return nil, errors.New("invalid trace id")
 	}
@@ -237,7 +237,7 @@ func (q *Querier) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDReque
 	span.SetTag("queryMode", req.QueryMode)
 
 	maxBytes := q.limits.MaxBytesPerTrace(userID)
-	combiner := trace.NewCombiner(maxBytes)
+	combiner := trace.NewCombiner(maxBytes, allowPartialTrace)
 
 	var spanCount, spanCountTotal, traceCountTotal int
 	if req.QueryMode == QueryModeIngesters || req.QueryMode == QueryModeAll {
@@ -252,6 +252,7 @@ func (q *Querier) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDReque
 		// get responses from all ingesters in parallel
 		span.LogFields(ot_log.String("msg", "searching ingesters"))
 		responses, err := q.forIngesterRings(ctx, userID, getRSFn, func(funcCtx context.Context, client tempopb.QuerierClient) (interface{}, error) {
+			req.AllowPartialTrace = allowPartialTrace
 			return client.FindTraceByID(funcCtx, req)
 		})
 		if err != nil {
@@ -285,6 +286,7 @@ func (q *Querier) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDReque
 
 		opts := common.DefaultSearchOptionsWithMaxBytes(maxBytes)
 		opts.BlockReplicationFactor = backend.DefaultReplicationFactor
+		opts.AllowPartialTraces = allowPartialTrace
 		partialTraces, blockErrs, err := q.store.Find(ctx, userID, req.TraceID, req.BlockStart, req.BlockEnd, timeStart, timeEnd, opts)
 		if err != nil {
 			retErr := fmt.Errorf("error querying store in Querier.FindTraceByID: %w", err)

pkg/model/trace/combine.go

@@ -41,15 +41,18 @@ var ErrTraceTooLarge = fmt.Errorf("trace exceeds max size")
 // * Only sort the final result once and if needed.
 // * Don't scan/hash the spans for the last input (final=true).
 type Combiner struct {
-	result       *tempopb.Trace
-	spans        map[token]struct{}
-	combined     bool
-	maxSizeBytes int
+	result              *tempopb.Trace
+	spans               map[token]struct{}
+	combined            bool
+	maxSizeBytes        int
+	allowPartialTrace   bool
+	maxTraceSizeReached bool
 }
 
-func NewCombiner(maxSizeBytes int) *Combiner {
+func NewCombiner(maxSizeBytes int, allowPartialTrace bool) *Combiner {
 	return &Combiner{
-		maxSizeBytes: maxSizeBytes,
+		maxSizeBytes:      maxSizeBytes,
+		allowPartialTrace: allowPartialTrace,
 	}
 }
 
@@ -63,7 +66,7 @@ func (c *Combiner) Consume(tr *tempopb.Trace) (int, error) {
 func (c *Combiner) ConsumeWithFinal(tr *tempopb.Trace, final bool) (int, error) {
 	var spanCount int
 	if tr == nil {
-		return spanCount, c.sizeError()
+		return spanCount, nil
 	}
 
 	h := newHash()
@@ -90,9 +93,17 @@ func (c *Combiner) ConsumeWithFinal(tr *tempopb.Trace, final bool) (int, error)
 				}
 			}
 		}
-		return spanCount, c.sizeError()
+		maxSizeErr := c.sizeError()
+		if maxSizeErr != nil && c.allowPartialTrace {
+			return spanCount, nil
+		}
+		return spanCount, maxSizeErr
 	}
 
+	// Do not combine more spans for now
+	if c.maxTraceSizeReached && c.allowPartialTrace {
+		return spanCount, nil
+	}
 	// loop through every span and copy spans in B that don't exist to A
 	for _, b := range tr.ResourceSpans {
 		notFoundILS := b.ScopeSpans[:0]
@@ -129,15 +140,23 @@ func (c *Combiner) ConsumeWithFinal(tr *tempopb.Trace, final bool) (int, error)
 	}
 
 	c.combined = true
-	return spanCount, c.sizeError()
+	maxSizeErr := c.sizeError()
+	if maxSizeErr != nil && c.allowPartialTrace {
+		return spanCount, nil
+	}
+
+	return spanCount, maxSizeErr
 }
 
 func (c *Combiner) sizeError() error {
+	// Should we allow a maxSizeBytes <= 0?
 	if c.result == nil || c.maxSizeBytes <= 0 {
 		return nil
 	}
 
 	if c.result.Size() > c.maxSizeBytes {
+		// To avoid recalculing the size
+		c.maxTraceSizeReached = true
 		return fmt.Errorf("%w (max bytes: %d)", ErrTraceTooLarge, c.maxSizeBytes)
 	}
 

pkg/model/trace/combine_test.go

@@ -17,7 +17,7 @@ import (
 func TestCombineProtoTotals(t *testing.T) {
 	methods := []func(a, b *tempopb.Trace) (*tempopb.Trace, int){
 		func(a, b *tempopb.Trace) (*tempopb.Trace, int) {
-			c := NewCombiner(0)
+			c := NewCombiner(0, false)
 			_, err := c.Consume(a)
 			require.NoError(t, err)
 			_, err = c.Consume(b)
@@ -65,7 +65,7 @@ func TestCombineProtoTotals(t *testing.T) {
 func TestCombinerChecksMaxBytes(t *testing.T) {
 	// Ensure that the combiner checks max bytes when consuming a trace.
 	for _, maxBytes := range []int{0, 100, 1000, 10000} {
-		c := NewCombiner(maxBytes)
+		c := NewCombiner(maxBytes, false)
 		curSize := 0
 
 		// attempt up to 20 traces to exceed max bytes
@@ -83,6 +83,27 @@ func TestCombinerChecksMaxBytes(t *testing.T) {
 	}
 }
 
+func TestCombinerReturnsAPartialTrace(t *testing.T) {
+	// Ensure that the combiner checks max bytes when consuming a trace.
+	for _, maxBytes := range []int{0, 100, 1000, 10000} {
+		c := NewCombiner(maxBytes, true)
+		curSize := 0
+
+		// attempt up to 20 traces to exceed max bytes
+		for i := 0; i < 20; i++ {
+			tr := test.MakeTraceWithSpanCount(1, 1, []byte{0x01})
+			curSize += tr.Size()
+
+			_, err := c.Consume(tr)
+			if curSize > maxBytes && maxBytes != 0 {
+				require.NoError(t, err)
+				continue
+			}
+			require.NoError(t, err)
+		}
+	}
+}
+
 func TestTokenForIDCollision(t *testing.T) {
 	// Estimate the hash collision rate of tokenForID.
 
@@ -147,7 +168,7 @@ func BenchmarkCombine(b *testing.B) {
 		{
 			"Combiner",
 			func(traces []*tempopb.Trace) int {
-				c := NewCombiner(0)
+				c := NewCombiner(0, false)
 				for i := range traces {
 					_, err := c.ConsumeWithFinal(traces[i], i == len(traces)-1)
 					require.NoError(b, err)

pkg/model/v1/object_decoder.go

@@ -40,7 +40,7 @@ func (d *ObjectDecoder) PrepareForRead(obj []byte) (*tempopb.Trace, error) {
 }
 
 func (d *ObjectDecoder) Combine(objs ...[]byte) ([]byte, error) {
-	c := trace.NewCombiner(0)
+	c := trace.NewCombiner(0, false)
 	for i, obj := range objs {
 		t, err := staticDecoder.PrepareForRead(obj)
 		if err != nil {

pkg/model/v1/segment_decoder.go

@@ -25,7 +25,7 @@ func (d *SegmentDecoder) PrepareForWrite(trace *tempopb.Trace, _, _ uint32) ([]b
 
 func (d *SegmentDecoder) PrepareForRead(segments [][]byte) (*tempopb.Trace, error) {
 	// each slice is a marshalled tempopb.Trace, unmarshal and combine
-	combiner := trace.NewCombiner(0)
+	combiner := trace.NewCombiner(0, false)
 	for i, s := range segments {
 		t := &tempopb.Trace{}
 		err := proto.Unmarshal(s, t)

pkg/model/v2/object_decoder.go

@@ -57,7 +57,7 @@ func (d *ObjectDecoder) Combine(objs ...[]byte) ([]byte, error) {
 	var minStart, maxEnd uint32
 	minStart = math.MaxUint32
 
-	c := trace.NewCombiner(0)
+	c := trace.NewCombiner(0, false)
 	for i, obj := range objs {
 		t, err := d.PrepareForRead(obj)
 		if err != nil {

pkg/model/v2/segment_decoder.go

@@ -28,7 +28,7 @@ func (d *SegmentDecoder) PrepareForWrite(trace *tempopb.Trace, start uint32, end
 }
 
 func (d *SegmentDecoder) PrepareForRead(segments [][]byte) (*tempopb.Trace, error) {
-	combiner := trace.NewCombiner(0)
+	combiner := trace.NewCombiner(0, false)
 	for i, obj := range segments {
 		obj, _, _, err := stripStartEnd(obj)
 		if err != nil {