From 64ef10079757a456f02d51cfe8a7a09eeccc98ae Mon Sep 17 00:00:00 2001
From: Christian Haudum <christian.haudum@gmail.com>
Date: Tue, 6 Feb 2024 15:32:24 +0100
Subject: [PATCH] fixup! FetchBlocks returns CloseableBlockQuerier

Signed-off-by: Christian Haudum <christian.haudum@gmail.com>
---
 .../stores/shipper/bloomshipper/cache.go      |  6 ++++--
 .../stores/shipper/bloomshipper/shipper.go    |  7 +------
 .../stores/shipper/bloomshipper/store.go      | 21 +++++++++----------
 .../stores/shipper/bloomshipper/store_test.go | 18 ++++++++--------
 4 files changed, 24 insertions(+), 28 deletions(-)

diff --git a/pkg/storage/stores/shipper/bloomshipper/cache.go b/pkg/storage/stores/shipper/bloomshipper/cache.go
index 6b423764cc600..7ef20065a457b 100644
--- a/pkg/storage/stores/shipper/bloomshipper/cache.go
+++ b/pkg/storage/stores/shipper/bloomshipper/cache.go
@@ -19,6 +19,7 @@ import (
 
 type ClosableBlockQuerier struct {
 	*v1.BlockQuerier
+	v1.FingerprintBounds
 	Close func()
 }
 
@@ -82,8 +83,9 @@ func (b BlockDirectory) Release() {
 func (b BlockDirectory) BlockQuerier() *ClosableBlockQuerier {
 	b.Acquire()
 	return &ClosableBlockQuerier{
-		BlockQuerier: v1.NewBlockQuerier(b.Block()),
-		Close:        b.Release,
+		BlockQuerier:      v1.NewBlockQuerier(b.Block()),
+		FingerprintBounds: b.BlockRef.Bounds,
+		Close:             b.Release,
 	}
 }
 
diff --git a/pkg/storage/stores/shipper/bloomshipper/shipper.go b/pkg/storage/stores/shipper/bloomshipper/shipper.go
index 3b7fdb2c28cca..73c52b562f37c 100644
--- a/pkg/storage/stores/shipper/bloomshipper/shipper.go
+++ b/pkg/storage/stores/shipper/bloomshipper/shipper.go
@@ -14,11 +14,6 @@ import (
 	"github.com/grafana/loki/pkg/storage/stores/shipper/bloomshipper/config"
 )
 
-type BlockQuerierWithFingerprintRange struct {
-	*v1.BlockQuerier
-	v1.FingerprintBounds
-}
-
 type ForEachBlockCallback func(bq *v1.BlockQuerier, bounds v1.FingerprintBounds) error
 
 type Interface interface {
@@ -70,7 +65,7 @@ func (s *Shipper) ForEach(ctx context.Context, _ string, refs []BlockRef, callba
 	}
 
 	for i := range bqs {
-		err := callback(bqs[i].BlockQuerier, refs[i].Bounds)
+		err := callback(bqs[i].BlockQuerier, bqs[i].FingerprintBounds)
 		// close querier to decrement ref count
 		bqs[i].Close()
 		if err != nil {
diff --git a/pkg/storage/stores/shipper/bloomshipper/store.go b/pkg/storage/stores/shipper/bloomshipper/store.go
index 981e6f54d22f3..542976073f583 100644
--- a/pkg/storage/stores/shipper/bloomshipper/store.go
+++ b/pkg/storage/stores/shipper/bloomshipper/store.go
@@ -9,7 +9,6 @@ import (
 	"github.com/go-kit/log"
 	"github.com/pkg/errors"
 	"github.com/prometheus/common/model"
-	"golang.org/x/exp/slices"
 
 	"github.com/grafana/loki/pkg/storage"
 	v1 "github.com/grafana/loki/pkg/storage/bloom/v1"
@@ -307,16 +306,16 @@ func (b *BloomStore) FetchBlocks(ctx context.Context, blocks []BlockRef) ([]*Clo
 		}
 	}
 
-	// sort responses (results []BlockDirectory) based on requests (blocks []BlockRef)
-	slices.SortFunc(results, func(a, b BlockDirectory) int {
-		ia, ib := slices.Index(blocks, a.BlockRef), slices.Index(blocks, b.BlockRef)
-		if ia < ib {
-			return -1
-		} else if ia > ib {
-			return +1
-		}
-		return 0
-	})
+	// // sort responses (results []*CloseableBlockQuerier) based on requests (blocks []BlockRef)
+	// slices.SortFunc(results, func(a, b *ClosableBlockQuerier) int {
+	// 	ia, ib := slices.Index(blocks, a.BlockRef), slices.Index(blocks, b.BlockRef)
+	// 	if ia < ib {
+	// 		return -1
+	// 	} else if ia > ib {
+	// 		return +1
+	// 	}
+	// 	return 0
+	// })
 
 	return results, nil
 }
diff --git a/pkg/storage/stores/shipper/bloomshipper/store_test.go b/pkg/storage/stores/shipper/bloomshipper/store_test.go
index aa0ca46e0660b..04641ca543a4e 100644
--- a/pkg/storage/stores/shipper/bloomshipper/store_test.go
+++ b/pkg/storage/stores/shipper/bloomshipper/store_test.go
@@ -248,19 +248,19 @@ func TestBloomStore_FetchBlocks(t *testing.T) {
 	ctx := context.Background()
 
 	// first call fetches two blocks from cache
-	blockDirs, err := store.FetchBlocks(ctx, []BlockRef{b1.BlockRef, b3.BlockRef})
+	bqs, err := store.FetchBlocks(ctx, []BlockRef{b1.BlockRef, b3.BlockRef})
 	require.NoError(t, err)
-	require.Len(t, blockDirs, 2)
+	require.Len(t, bqs, 2)
 
-	require.ElementsMatch(t, []BlockRef{b1.BlockRef, b3.BlockRef}, []BlockRef{blockDirs[0].BlockRef, blockDirs[1].BlockRef})
+	// require.Equal(t, []BlockRef{b1.BlockRef, b3.BlockRef}, []BlockRef{bqs[0].BlockRef, bqs[1].BlockRef})
 
 	// second call fetches two blocks from cache and two from storage
-	blockDirs, err = store.FetchBlocks(ctx, []BlockRef{b1.BlockRef, b2.BlockRef, b3.BlockRef, b4.BlockRef})
+	bqs, err = store.FetchBlocks(ctx, []BlockRef{b1.BlockRef, b2.BlockRef, b3.BlockRef, b4.BlockRef})
 	require.NoError(t, err)
-	require.Len(t, blockDirs, 4)
+	require.Len(t, bqs, 4)
 
-	require.Equal(t,
-		[]BlockRef{b1.BlockRef, b2.BlockRef, b3.BlockRef, b4.BlockRef},
-		[]BlockRef{blockDirs[0].BlockRef, blockDirs[1].BlockRef, blockDirs[2].BlockRef, blockDirs[3].BlockRef},
-	)
+	// require.Equal(t,
+	// 	[]BlockRef{b1.BlockRef, b2.BlockRef, b3.BlockRef, b4.BlockRef},
+	// 	[]BlockRef{bqs[0].BlockRef, bqs[1].BlockRef, bqs[2].BlockRef, bqs[3].BlockRef},
+	// )
 }