Add deep copy to avoid race (#2532)

* add deep copy to avoid race

Signed-off-by: Joe Elliott <[email protected]>

* conditionally copy spansets

Signed-off-by: Joe Elliott <[email protected]>

---------

Signed-off-by: Joe Elliott <[email protected]>

modules/frontend/search_progress.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/grafana/tempo/pkg/tempopb"
+	v1 "github.com/grafana/tempo/pkg/tempopb/common/v1"
 	"github.com/grafana/tempo/pkg/traceql"
 )
 
@@ -135,6 +136,31 @@ func (r *searchProgress) result() *shardedSearchResults {
 		finishedRequests: r.finishedRequests,
 	}
 
+	// copy metadata b/c the resultsCombiner holds a pointer to the data and continues
+	// to modify it. this may race with anything getting results
+	md := r.resultsCombiner.Metadata()
+	mdCopy := make([]*tempopb.TraceSearchMetadata, 0, len(md))
+	for _, m := range md {
+		mCopy := &tempopb.TraceSearchMetadata{
+			TraceID:           m.TraceID,
+			RootServiceName:   m.RootServiceName,
+			RootTraceName:     m.RootTraceName,
+			StartTimeUnixNano: m.StartTimeUnixNano,
+			DurationMs:        m.DurationMs,
+			SpanSet:           copySpanset(m.SpanSet),
+		}
+
+		// now copy spansets
+		if len(m.SpanSets) > 0 {
+			mCopy.SpanSets = make([]*tempopb.SpanSet, 0, len(m.SpanSets))
+			for _, ss := range m.SpanSets {
+				mCopy.SpanSets = append(mCopy.SpanSets, copySpanset(ss))
+			}
+		}
+
+		mdCopy = append(mdCopy, mCopy)
+	}
+
 	searchRes := &tempopb.SearchResponse{
 		// clone search metrics to avoid race conditions on the pointer
 		Metrics: &tempopb.SearchMetrics{
@@ -145,10 +171,24 @@ func (r *searchProgress) result() *shardedSearchResults {
 			TotalJobs:       r.resultsMetrics.TotalJobs,
 			TotalBlockBytes: r.resultsMetrics.TotalBlockBytes,
 		},
-		Traces: r.resultsCombiner.Metadata(),
+		Traces: mdCopy,
 	}
 
 	res.response = searchRes
 
 	return res
 }
+
+func copySpanset(ss *tempopb.SpanSet) *tempopb.SpanSet {
+	if ss == nil {
+		return nil
+	}
+
+	// the metadata results combiner considers the spans and attributes immutable. it does not attempt to change them
+	// so just copying the slices should be safe
+	return &tempopb.SpanSet{
+		Spans:      append([]*tempopb.Span(nil), ss.Spans...),
+		Matched:    ss.Matched,
+		Attributes: append([]*v1.KeyValue(nil), ss.Attributes...),
+	}
+}

modules/frontend/search_progress_test.go

@@ -3,10 +3,12 @@ package frontend
 import (
 	"context"
 	"errors"
+	"math"
 	"testing"
 	"time"
 
 	"github.com/grafana/tempo/pkg/tempopb"
+	"github.com/grafana/tempo/pkg/traceql"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -121,3 +123,61 @@ func TestSearchProgressCombineResults(t *testing.T) {
 	assert.Equal(t, expected, sr.result())
 
 }
+
+func TestInstanceDoesNotRace(t *testing.T) {
+
+	end := make(chan struct{})
+	concurrent := func(f func()) {
+		for {
+			select {
+			case <-end:
+				return
+			default:
+				f()
+			}
+		}
+	}
+
+	traceID := "1234"
+
+	progress := newSearchProgress(context.Background(), 10, 0, 0, 0)
+	i := 0
+	go concurrent(func() {
+		i++
+		resp := &tempopb.SearchResponse{
+			Traces: []*tempopb.TraceSearchMetadata{
+				{
+					TraceID:           traceID,
+					StartTimeUnixNano: math.MaxUint64 - uint64(i),
+					DurationMs:        uint32(i),
+					SpanSets: []*tempopb.SpanSet{{
+						Matched: uint32(i),
+					}},
+				},
+			},
+			Metrics: &tempopb.SearchMetrics{
+				InspectedTraces: 1,
+				InspectedBytes:  1,
+				TotalBlocks:     1,
+				TotalJobs:       1,
+				CompletedJobs:   1,
+			},
+		}
+		progress.addResponse(resp)
+	})
+
+	combiner := traceql.NewMetadataCombiner()
+	go concurrent(func() {
+		res := progress.result()
+		if len(res.response.Traces) > 0 {
+			// by using a combiner we are testing and changing the entire response
+			combiner.AddMetadata(res.response.Traces[0])
+		}
+	})
+
+	time.Sleep(100 * time.Millisecond)
+	close(end)
+	// Wait for go funcs to quit before
+	// exiting and cleaning up
+	time.Sleep(2 * time.Second)
+}