fix(ingest-limits): Query all assigned partitions (#17473)

Author: periklis

pkg/limits/frontend/ring.go

@@ -71,7 +71,7 @@ func (g *RingStreamUsageGatherer) forGivenReplicaSet(ctx context.Context, rs rin
 		return nil, err
 	}
 
-	instancesToQuery := make(map[string][]uint64)
+	ownedStreamHeashes := make(map[string][]uint64)
 	for _, hash := range r.StreamHashes {
 		partitionID := int32(hash % uint64(g.numPartitions))
 		addr, ok := partitionConsumers[partitionID]
@@ -80,34 +80,36 @@ func (g *RingStreamUsageGatherer) forGivenReplicaSet(ctx context.Context, rs rin
 			level.Warn(g.logger).Log("msg", "no instance found for partition", "partition", partitionID)
 			continue
 		}
-		instancesToQuery[addr] = append(instancesToQuery[addr], hash)
+		ownedStreamHeashes[addr] = append(ownedStreamHeashes[addr], hash)
 	}
 
 	errg, ctx := errgroup.WithContext(ctx)
-	responses := make([]GetStreamUsageResponse, len(instancesToQuery))
+	responses := make([]GetStreamUsageResponse, len(rs.Instances))
 
-	// Query each instance for stream usage
-	i := 0
-	for addr, hashes := range instancesToQuery {
-		j := i
-		i++
+	// Query all healthy instances in parallel,
+	// send requested stream hahes only to owning instances.
+	for i, instance := range rs.Instances {
 		errg.Go(func() error {
-			client, err := g.pool.GetClientFor(addr)
+			client, err := g.pool.GetClientFor(instance.Addr)
 			if err != nil {
 				return err
 			}
 
-			protoReq := &logproto.GetStreamUsageRequest{
-				Tenant:       r.Tenant,
-				StreamHashes: hashes,
+			protoReq := &logproto.GetStreamUsageRequest{Tenant: r.Tenant}
+
+			// Only send stream hashes to the instance that owns them.
+			// This eliminates the need in downstream callers to filter
+			// duplicate unknown streams.
+			if hashes, ok := ownedStreamHeashes[instance.Addr]; ok {
+				protoReq.StreamHashes = hashes
 			}
 
 			resp, err := client.(logproto.IngestLimitsClient).GetStreamUsage(ctx, protoReq)
 			if err != nil {
 				return err
 			}
 
-			responses[j] = GetStreamUsageResponse{Addr: addr, Response: resp}
+			responses[i] = GetStreamUsageResponse{Addr: instance.Addr, Response: resp}
 			return nil
 		})
 	}

pkg/limits/frontend/ring_test.go

@@ -81,8 +81,9 @@ func TestRingStreamUsageGatherer_GetStreamUsage(t *testing.T) {
 		}},
 	}, {
 		// When there is one stream, and two instances each owning separate
-		// partitions, only the instance owning the partition for the stream hash
-		// should be queried.
+		// partitions, all instances should be queried. However, only the
+		// instance owning the partition for the stream hash should be queried
+		// for the requested stream hashes.
 		name: "one stream two instances",
 		getStreamUsageRequest: GetStreamUsageRequest{
 			Tenant:       "test",
@@ -104,16 +105,29 @@ func TestRingStreamUsageGatherer_GetStreamUsage(t *testing.T) {
 				1: time.Now().UnixNano(),
 			},
 		}},
-		expectedStreamUsageRequests: []*logproto.GetStreamUsageRequest{nil, {
+		expectedStreamUsageRequests: []*logproto.GetStreamUsageRequest{{
+			Tenant: "test",
+		}, {
 			Tenant:       "test",
 			StreamHashes: []uint64{0x1},
 		}},
-		getStreamUsageResponses: []*logproto.GetStreamUsageResponse{nil, {
+		getStreamUsageResponses: []*logproto.GetStreamUsageResponse{{
+			Tenant:        "test",
+			ActiveStreams: 1,
+			Rate:          10,
+		}, {
 			Tenant:        "test",
 			ActiveStreams: 1,
 			Rate:          10,
 		}},
 		expectedResponses: []GetStreamUsageResponse{{
+			Addr: "instance-0",
+			Response: &logproto.GetStreamUsageResponse{
+				Tenant:        "test",
+				ActiveStreams: 1,
+				Rate:          10,
+			},
+		}, {
 			Addr: "instance-1",
 			Response: &logproto.GetStreamUsageResponse{
 				Tenant:        "test",
@@ -124,7 +138,7 @@ func TestRingStreamUsageGatherer_GetStreamUsage(t *testing.T) {
 	}, {
 		// When there is one stream, and two instances owning overlapping
 		// partitions, only the instance with the latest timestamp for the relevant
-		// partition should be queried.
+		// partition should be queried with the requested stream hash.
 		name: "one stream two instances, overlapping partition ownership",
 		getStreamUsageRequest: GetStreamUsageRequest{
 			Tenant:       "test",
@@ -146,16 +160,29 @@ func TestRingStreamUsageGatherer_GetStreamUsage(t *testing.T) {
 				1: time.Now().UnixNano(),
 			},
 		}},
-		expectedStreamUsageRequests: []*logproto.GetStreamUsageRequest{nil, {
+		expectedStreamUsageRequests: []*logproto.GetStreamUsageRequest{{
+			Tenant: "test",
+		}, {
 			Tenant:       "test",
 			StreamHashes: []uint64{0x1},
 		}},
-		getStreamUsageResponses: []*logproto.GetStreamUsageResponse{nil, {
+		getStreamUsageResponses: []*logproto.GetStreamUsageResponse{{
+			Tenant:        "test",
+			ActiveStreams: 1,
+			Rate:          10,
+		}, {
 			Tenant:        "test",
 			ActiveStreams: 1,
 			Rate:          10,
 		}},
 		expectedResponses: []GetStreamUsageResponse{{
+			Addr: "instance-0",
+			Response: &logproto.GetStreamUsageResponse{
+				Tenant:        "test",
+				ActiveStreams: 1,
+				Rate:          10,
+			},
+		}, {
 			Addr: "instance-1",
 			Response: &logproto.GetStreamUsageResponse{
 				Tenant:        "test",