feat: Allow partial traces to be returned

CHANGELOG.md

@@ -3,6 +3,8 @@
 * [FEATURE] Discarded span logging `log_discarded_spans` [#3957](https://github.com/grafana/tempo/issues/3957) (@dastrobu)
 * [ENHANCEMENT] TraceQL: Attribute iterators collect matched array values [#3867](https://github.com/grafana/tempo/pull/3867) (@electron0zero, @stoewer)
 * [ENHANCEMENT] Add bytes and spans received to usage stats [#3983](https://github.com/grafana/tempo/pull/3983) (@joe-elliott)
+* [ENHANCEMENT] Allow returning partial traces that exceed the MaxBytes limit for V2 [#3941](https://github.com/grafana/tempo/pull/3941) (@javiermolinar)
+* [ENHANCEMENT] Added new middleware to validate request query values [#3993](https://github.com/grafana/tempo/pull/3993) (@javiermolinar)
 * [ENHANCEMENT] Prevent massive allocations in the frontend if there is not sufficient pressure from the query pipeline. [#3996](https://github.com/grafana/tempo/pull/3996) (@joe-elliott)
   **BREAKING CHANGE** Removed `querier_forget_delay` setting from the frontend. This configuration option did nothing.
 * [ENHANCEMENT] Update metrics-generator config in Tempo distributed docker compose example to serve TraceQL metrics [#4003](https://github.com/grafana/tempo/pull/4003) (@javiermolinar)

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

@@ -1,14 +1,20 @@
 package combiner
 
 import (
+	"fmt"
+
 	"github.com/grafana/tempo/pkg/model/trace"
 	"github.com/grafana/tempo/pkg/tempopb"
 )
 
 func NewTraceByIDV2(maxBytes int, marshalingFormat string) Combiner {
-	combiner := trace.NewCombiner(maxBytes)
+	combiner := trace.NewCombiner(maxBytes, true)
+	var partialTrace bool
 	gc := &genericCombiner[*tempopb.TraceByIDResponse]{
 		combine: func(partial *tempopb.TraceByIDResponse, _ *tempopb.TraceByIDResponse, _ PipelineResponse) error {
+			if partial.Status == tempopb.TraceByIDResponse_PARTIAL {
+				partialTrace = true
+			}
 			_, err := combiner.Consume(partial.Trace)
 			return err
 		},
@@ -21,8 +27,13 @@ func NewTraceByIDV2(maxBytes int, marshalingFormat string) Combiner {
 			// dedupe duplicate span ids
 			deduper := newDeduper()
 			traceResult = deduper.dedupe(traceResult)
-
 			resp.Trace = traceResult
+
+			if partialTrace || combiner.IsPartialTrace() {
+				resp.Status = tempopb.TraceByIDResponse_PARTIAL
+				resp.Message = fmt.Sprintf("Trace exceeds maximum size of %d bytes, a partial trace is returned", maxBytes)
+			}
+
 			return resp, nil
 		},
 		new:     func() *tempopb.TraceByIDResponse { return &tempopb.TraceByIDResponse{} },

modules/frontend/combiner/trace_by_id_v2_test.go

@@ -27,6 +27,61 @@ func (m MockResponse) RequestData() any {
 	return nil
 }
 
+func TestNewTraceByIdV2ReturnsAPartialTrace(t *testing.T) {
+	traceResponse := &tempopb.TraceByIDResponse{
+		Trace:   test.MakeTrace(2, []byte{0x01, 0x02}),
+		Metrics: &tempopb.TraceByIDMetrics{},
+	}
+	resBytes, err := proto.Marshal(traceResponse)
+	require.NoError(t, err)
+	response := http.Response{
+		StatusCode: 200,
+		Header: map[string][]string{
+			"Content-Type": {"application/protobuf"},
+		},
+		Body: io.NopCloser(bytes.NewReader(resBytes)),
+	}
+	combiner := NewTraceByIDV2(10, api.HeaderAcceptJSON)
+	err = combiner.AddResponse(MockResponse{&response})
+	require.NoError(t, err)
+
+	res, err := combiner.HTTPFinal()
+	require.NoError(t, err)
+
+	actualResp := &tempopb.TraceByIDResponse{}
+	err = new(jsonpb.Unmarshaler).Unmarshal(res.Body, actualResp)
+	require.NoError(t, err)
+	assert.Equal(t, actualResp.Status, tempopb.TraceByIDResponse_PARTIAL)
+}
+
+func TestNewTraceByIdV2ReturnsAPartialTraceOnPartialTraceReturnedByQuerier(t *testing.T) {
+	traceResponse := &tempopb.TraceByIDResponse{
+		Trace:   test.MakeTrace(2, []byte{0x01, 0x02}),
+		Status:  tempopb.TraceByIDResponse_PARTIAL,
+		Metrics: &tempopb.TraceByIDMetrics{},
+	}
+	resBytes, err := proto.Marshal(traceResponse)
+	require.NoError(t, err)
+	response := http.Response{
+		StatusCode: 200,
+		Header: map[string][]string{
+			"Content-Type": {"application/protobuf"},
+		},
+		Body: io.NopCloser(bytes.NewReader(resBytes)),
+	}
+	combiner := NewTraceByIDV2(10, api.HeaderAcceptJSON)
+	err = combiner.AddResponse(MockResponse{&response})
+	require.NoError(t, err)
+
+	res, err := combiner.HTTPFinal()
+	require.NoError(t, err)
+
+	actualResp := &tempopb.TraceByIDResponse{}
+	err = new(jsonpb.Unmarshaler).Unmarshal(res.Body, actualResp)
+	require.NoError(t, err)
+	assert.Equal(t, actualResp.Status, tempopb.TraceByIDResponse_PARTIAL)
+}
+
 func TestNewTraceByIDV2(t *testing.T) {
 	traceResponse := &tempopb.TraceByIDResponse{
 		Trace:   test.MakeTrace(2, []byte{0x01, 0x02}),

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

@@ -690,7 +690,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

@@ -108,10 +108,11 @@ func (q *Querier) TraceByIDHandlerV2(w http.ResponseWriter, r *http.Request) {
 		ot_log.String("apiVersion", "v2"))
 
 	resp, err := q.FindTraceByID(ctx, &tempopb.TraceByIDRequest{
-		TraceID:    byteID,
-		BlockStart: blockStart,
-		BlockEnd:   blockEnd,
-		QueryMode:  queryMode,
+		TraceID:           byteID,
+		BlockStart:        blockStart,
+		BlockEnd:          blockEnd,
+		QueryMode:         queryMode,
+		AllowPartialTrace: true,
 	}, timeStart, timeEnd)
 	if err != nil {
 		handleError(w, err)

modules/querier/querier.go

@@ -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, req.AllowPartialTrace)
 
 	var spanCount, spanCountTotal, traceCountTotal int
 	if req.QueryMode == QueryModeIngesters || req.QueryMode == QueryModeAll {
@@ -309,11 +309,17 @@ func (q *Querier) FindTraceByID(ctx context.Context, req *tempopb.TraceByIDReque
 	}
 
 	completeTrace, _ := combiner.Result()
-
-	return &tempopb.TraceByIDResponse{
+	resp := &tempopb.TraceByIDResponse{
 		Trace:   completeTrace,
 		Metrics: &tempopb.TraceByIDMetrics{},
-	}, nil
+	}
+
+	if combiner.IsPartialTrace() {
+		resp.Status = tempopb.TraceByIDResponse_PARTIAL
+		resp.Message = fmt.Sprintf("Trace exceeds maximum size of %d bytes, a partial trace is returned", maxBytes)
+	}
+
+	return resp, nil
 }
 
 type (

pkg/model/trace/combine.go

@@ -41,15 +41,20 @@ 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 {
+// It creates a new Trace combiner. If maxSizeBytes is 0, the final trace size is not checked
+// when allowPartialTrace is set to true a partial trace that exceed the max size may be returned
+func NewCombiner(maxSizeBytes int, allowPartialTrace bool) *Combiner {
 	return &Combiner{
-		maxSizeBytes: maxSizeBytes,
+		maxSizeBytes:      maxSizeBytes,
+		allowPartialTrace: allowPartialTrace,
 	}
 }
 
@@ -62,8 +67,8 @@ func (c *Combiner) Consume(tr *tempopb.Trace) (int, error) {
 // it is known that this is the last expected input trace.
 func (c *Combiner) ConsumeWithFinal(tr *tempopb.Trace, final bool) (int, error) {
 	var spanCount int
-	if tr == nil {
-		return spanCount, c.sizeError()
+	if tr == nil || c.IsPartialTrace() {
+		return spanCount, nil
 	}
 
 	h := newHash()
@@ -90,7 +95,12 @@ func (c *Combiner) ConsumeWithFinal(tr *tempopb.Trace, final bool) (int, error)
 				}
 			}
 		}
-		return spanCount, c.sizeError()
+
+		maxSizeErr := c.sizeError()
+		if c.IsPartialTrace() {
+			return spanCount, nil
+		}
+		return spanCount, maxSizeErr
 	}
 
 	// loop through every span and copy spans in B that don't exist to A
@@ -129,7 +139,12 @@ func (c *Combiner) ConsumeWithFinal(tr *tempopb.Trace, final bool) (int, error)
 	}
 
 	c.combined = true
-	return spanCount, c.sizeError()
+	maxSizeErr := c.sizeError()
+	if c.IsPartialTrace() {
+		return spanCount, nil
+	}
+
+	return spanCount, maxSizeErr
 }
 
 func (c *Combiner) sizeError() error {
@@ -138,6 +153,7 @@ func (c *Combiner) sizeError() error {
 	}
 
 	if c.result.Size() > c.maxSizeBytes {
+		c.maxTraceSizeReached = true
 		return fmt.Errorf("%w (max bytes: %d)", ErrTraceTooLarge, c.maxSizeBytes)
 	}
 
@@ -156,3 +172,8 @@ func (c *Combiner) Result() (*tempopb.Trace, int) {
 
 	return c.result, spanCount
 }
+
+// Returns true if the combined trace is a partial one if partal trace is enabled
+func (c *Combiner) IsPartialTrace() bool {
+	return c.maxTraceSizeReached && c.allowPartialTrace
+}