fix: avoid recalculating the segmentation key hash twice (#19961)

Author: grobinson-grafana

pkg/distributor/dataobj_tee.go

@@ -98,7 +98,8 @@ func NewDataObjTee(
 // A SegmentedStream is a KeyedStream with a segmentation key.
 type SegmentedStream struct {
 	KeyedStream
-	SegmentationKey SegmentationKey
+	SegmentationKey     SegmentationKey
+	SegmentationKeyHash uint64
 }
 
 // Duplicate implements the [Tee] interface.
@@ -112,8 +113,9 @@ func (t *DataObjTee) Duplicate(ctx context.Context, tenant string, streams []Key
 			return
 		}
 		segmentationKeyStreams = append(segmentationKeyStreams, SegmentedStream{
-			KeyedStream:     stream,
-			SegmentationKey: segmentationKey,
+			KeyedStream:         stream,
+			SegmentationKey:     segmentationKey,
+			SegmentationKeyHash: segmentationKey.Sum64(),
 		})
 	}
 	rates, err := t.limitsClient.UpdateRates(ctx, tenant, segmentationKeyStreams)
@@ -127,7 +129,7 @@ func (t *DataObjTee) Duplicate(ctx context.Context, tenant string, streams []Key
 		fastRates[rate.StreamHash] = rate.Rate
 	}
 	for _, s := range segmentationKeyStreams {
-		go t.duplicate(ctx, tenant, s, fastRates[s.SegmentationKey.Sum64()])
+		go t.duplicate(ctx, tenant, s, fastRates[s.SegmentationKeyHash])
 	}
 }
 

pkg/distributor/ingest_limits.go

@@ -218,7 +218,7 @@ func newUpdateRatesRequest(tenant string, streams []SegmentedStream) (*proto.Upd
 	for _, stream := range streams {
 		entriesSize, structuredMetadataSize := calculateStreamSizes(stream.Stream)
 		streamMetadata = append(streamMetadata, &proto.StreamMetadata{
-			StreamHash:      stream.SegmentationKey.Sum64(),
+			StreamHash:      stream.SegmentationKeyHash,
 			TotalSize:       entriesSize + structuredMetadataSize,
 			IngestionPolicy: stream.Policy,
 		})

pkg/distributor/ingest_limits_test.go

@@ -320,22 +320,23 @@ func TestIngestLimits_UpdateRates(t *testing.T) {
 		name:   "updates rates",
 		tenant: "test",
 		streams: []SegmentedStream{{
-			SegmentationKey: "test",
+			SegmentationKey:     "test",
+			SegmentationKeyHash: 13113208752873574959,
 		}},
 		expectedRequest: &proto.UpdateRatesRequest{
 			Tenant: "test",
 			Streams: []*proto.StreamMetadata{{
-				StreamHash: 0xb5fb79e24c92922f,
+				StreamHash: 13113208752873574959,
 			}},
 		},
 		response: &proto.UpdateRatesResponse{
 			Results: []*proto.UpdateRatesResult{{
-				StreamHash: 0xb5fb79e24c92922f,
+				StreamHash: 13113208752873574959,
 				Rate:       1024,
 			}},
 		},
 		expectedResult: []*proto.UpdateRatesResult{{
-			StreamHash: 0xb5fb79e24c92922f,
+			StreamHash: 13113208752873574959,
 			Rate:       1024,
 		}},
 	}}

pkg/distributor/segment_test.go

@@ -54,9 +54,9 @@ func TestGetSegmentationKey(t *testing.T) {
 
 func TestSegmentationKey_Sum64(t *testing.T) {
 	k1 := SegmentationKey("")
-	require.Equal(t, uint64(0x552129d0d55dcd1b), k1.Sum64())
+	require.Equal(t, uint64(6134230144364956955), k1.Sum64())
 	k2 := SegmentationKey("abc")
-	require.Equal(t, uint64(0x4a3a160be83aefc5), k2.Sum64())
+	require.Equal(t, uint64(5348611747852513221), k2.Sum64())
 	// The same key always produces the same 64 bit sum.
 	k3 := SegmentationKey("abc")
 	require.Equal(t, k2.Sum64(), k3.Sum64())