From b2238de7d70d98f923d1441ebb41ef95455793fb Mon Sep 17 00:00:00 2001
From: Christian Haudum <christian.haudum@gmail.com>
Date: Mon, 22 Jan 2024 09:57:30 +0100
Subject: [PATCH] Fix panic in chunk removal

Signed-off-by: Christian Haudum <christian.haudum@gmail.com>
---
 pkg/bloomgateway/bloomgateway.go      |  9 +++--
 pkg/bloomgateway/bloomgateway_test.go | 58 +++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/pkg/bloomgateway/bloomgateway.go b/pkg/bloomgateway/bloomgateway.go
index 5c2cc9dad003..bb58f79288ae 100644
--- a/pkg/bloomgateway/bloomgateway.go
+++ b/pkg/bloomgateway/bloomgateway.go
@@ -355,7 +355,7 @@ outer:
 			continue
 		}
 		// we must not remove items from req.Refs as long as the worker may iterater over them
-		g.removeNotMatchingChunks(req, o)
+		removeNotMatchingChunks(req, o, g.logger)
 	}
 
 	g.metrics.addUnfilteredCount(numChunksUnfiltered)
@@ -365,7 +365,7 @@ outer:
 	return &logproto.FilterChunkRefResponse{ChunkRefs: req.Refs}, nil
 }
 
-func (g *Gateway) removeNotMatchingChunks(req *logproto.FilterChunkRefRequest, res v1.Output) {
+func removeNotMatchingChunks(req *logproto.FilterChunkRefRequest, res v1.Output, logger log.Logger) {
 	// binary search index of fingerprint
 	idx := sort.Search(len(req.Refs), func(i int) bool {
 		return req.Refs[i].Fingerprint >= uint64(res.Fp)
@@ -373,7 +373,7 @@ func (g *Gateway) removeNotMatchingChunks(req *logproto.FilterChunkRefRequest, r
 
 	// fingerprint not found
 	if idx >= len(req.Refs) {
-		level.Error(g.logger).Log("msg", "index out of range", "idx", idx, "len", len(req.Refs), "fp", uint64(res.Fp))
+		level.Error(logger).Log("msg", "index out of range", "idx", idx, "len", len(req.Refs), "fp", uint64(res.Fp))
 		return
 	}
 
@@ -387,10 +387,11 @@ func (g *Gateway) removeNotMatchingChunks(req *logproto.FilterChunkRefRequest, r
 
 	for i := range res.Removals {
 		toRemove := res.Removals[i]
-		for j := range req.Refs[idx].Refs {
+		for j := 0; j < len(req.Refs[idx].Refs); j++ {
 			if toRemove.Checksum == req.Refs[idx].Refs[j].Checksum {
 				req.Refs[idx].Refs[j] = nil // avoid leaking pointer
 				req.Refs[idx].Refs = append(req.Refs[idx].Refs[:j], req.Refs[idx].Refs[j+1:]...)
+				j-- // since we removed the current item at index, we have to redo the same index
 			}
 		}
 	}
diff --git a/pkg/bloomgateway/bloomgateway_test.go b/pkg/bloomgateway/bloomgateway_test.go
index 1e85e7d2089c..e6d2b910c12f 100644
--- a/pkg/bloomgateway/bloomgateway_test.go
+++ b/pkg/bloomgateway/bloomgateway_test.go
@@ -340,6 +340,64 @@ func TestBloomGateway_FilterChunkRefs(t *testing.T) {
 	})
 }
 
+func TestBloomGateway_RemoveNotMatchingChunks(t *testing.T) {
+	t.Run("removing chunks partially", func(t *testing.T) {
+		req := &logproto.FilterChunkRefRequest{
+			Refs: []*logproto.GroupedChunkRefs{
+				{Fingerprint: 0x00, Tenant: "fake", Refs: []*logproto.ShortRef{
+					{Checksum: 0x1},
+					{Checksum: 0x2},
+					{Checksum: 0x3},
+					{Checksum: 0x4},
+					{Checksum: 0x5},
+				}},
+			},
+		}
+		res := v1.Output{
+			Fp: 0x00, Removals: v1.ChunkRefs{
+				{Checksum: 0x2},
+				{Checksum: 0x4},
+			},
+		}
+		expected := &logproto.FilterChunkRefRequest{
+			Refs: []*logproto.GroupedChunkRefs{
+				{Fingerprint: 0x00, Tenant: "fake", Refs: []*logproto.ShortRef{
+					{Checksum: 0x1},
+					{Checksum: 0x3},
+					{Checksum: 0x5},
+				}},
+			},
+		}
+		removeNotMatchingChunks(req, res, log.NewNopLogger())
+		require.Equal(t, expected, req)
+	})
+
+	t.Run("removing all chunks removed fingerprint ref", func(t *testing.T) {
+		req := &logproto.FilterChunkRefRequest{
+			Refs: []*logproto.GroupedChunkRefs{
+				{Fingerprint: 0x00, Tenant: "fake", Refs: []*logproto.ShortRef{
+					{Checksum: 0x1},
+					{Checksum: 0x2},
+					{Checksum: 0x3},
+				}},
+			},
+		}
+		res := v1.Output{
+			Fp: 0x00, Removals: v1.ChunkRefs{
+				{Checksum: 0x1},
+				{Checksum: 0x2},
+				{Checksum: 0x2},
+			},
+		}
+		expected := &logproto.FilterChunkRefRequest{
+			Refs: []*logproto.GroupedChunkRefs{},
+		}
+		removeNotMatchingChunks(req, res, log.NewNopLogger())
+		require.Equal(t, expected, req)
+	})
+
+}
+
 func createBlockQueriers(t *testing.T, numBlocks int, from, through model.Time, minFp, maxFp model.Fingerprint) ([]bloomshipper.BlockQuerierWithFingerprintRange, [][]v1.SeriesWithBloom) {
 	t.Helper()
 	step := (maxFp - minFp) / model.Fingerprint(numBlocks)