diff --git a/pkg/util/trace.go b/pkg/util/trace.go index b27eaca2507..afdc633a47a 100644 --- a/pkg/util/trace.go +++ b/pkg/util/trace.go @@ -2,6 +2,8 @@ package util import ( "bytes" + "hash" + "hash/fnv" "github.com/cortexproject/cortex/pkg/util" "github.com/go-kit/kit/log/level" @@ -72,14 +74,16 @@ func CombineTraceProtos(traceA, traceB *tempopb.Trace) (*tempopb.Trace, int, int spanCountB := 0 spanCountTotal := 0 + h := fnv.New32() + spansInA := make(map[uint32]struct{}) for _, batchA := range traceA.Batches { for _, ilsA := range batchA.InstrumentationLibrarySpans { for _, spanA := range ilsA.Spans { - spanCountA++ - spanCountTotal++ - spansInA[TokenForTraceID(spanA.SpanId)] = struct{}{} + spansInA[tokenForID(h, spanA.SpanId)] = struct{}{} } + spanCountA += len(ilsA.Spans) + spanCountTotal += len(ilsA.Spans) } } @@ -90,13 +94,13 @@ func CombineTraceProtos(traceA, traceB *tempopb.Trace) (*tempopb.Trace, int, int for _, ilsB := range batchB.InstrumentationLibrarySpans { notFoundSpans := ilsB.Spans[:0] for _, spanB := range ilsB.Spans { - spanCountB++ // if found in A, remove from the batch - _, ok := spansInA[TokenForTraceID(spanB.SpanId)] + _, ok := spansInA[tokenForID(h, spanB.SpanId)] if !ok { notFoundSpans = append(notFoundSpans, spanB) } } + spanCountB += len(ilsB.Spans) if len(notFoundSpans) > 0 { spanCountTotal += len(notFoundSpans) @@ -114,3 +118,9 @@ func CombineTraceProtos(traceA, traceB *tempopb.Trace) (*tempopb.Trace, int, int return traceA, spanCountA, spanCountB, spanCountTotal } + +func tokenForID(h hash.Hash32, b []byte) uint32 { + h.Reset() + _, _ = h.Write(b) + return h.Sum32() +} diff --git a/pkg/util/trace_test.go b/pkg/util/trace_test.go index 7519b2f2090..1d1d5125bef 100644 --- a/pkg/util/trace_test.go +++ b/pkg/util/trace_test.go @@ -2,6 +2,7 @@ package util import ( "bytes" + "fmt" "math/rand" "sort" "testing" @@ -186,3 +187,19 @@ func BenchmarkCombineTracesIdentical(b *testing.B) { CombineTraces(b1, b2) } } + +func BenchmarkCombineTraceProtos(b *testing.B) { + sizes := []int{1, 10, 1000, 10000, 100000} + + for _, size := range sizes { + b.Run(fmt.Sprint(size), func(b *testing.B) { + t1 := test.MakeTraceWithSpanCount(1, size, []byte{0x01, 0x02}) + t2 := test.MakeTraceWithSpanCount(1, size, []byte{0x01, 0x03}) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + CombineTraceProtos(t1, t2) + } + }) + } +}