From b39115fa9b07e1cc72c816ad960a9bdc58547db2 Mon Sep 17 00:00:00 2001
From: Dustin Wilson <dustin.wilson@grafana.com>
Date: Fri, 24 Jan 2025 14:21:31 -0500
Subject: [PATCH] add test for shuffled user order

---
 pkg/compactor/blocks_cleaner.go        |  5 ----
 pkg/storage/tsdb/users_scanner.go      |  5 ++++
 pkg/storage/tsdb/users_scanner_test.go | 39 +++++++++++++++++++++++---
 3 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/pkg/compactor/blocks_cleaner.go b/pkg/compactor/blocks_cleaner.go
index 0e921864cf3..55bfeffcdf1 100644
--- a/pkg/compactor/blocks_cleaner.go
+++ b/pkg/compactor/blocks_cleaner.go
@@ -8,7 +8,6 @@ package compactor
 import (
 	"context"
 	"fmt"
-	"math/rand"
 	"strconv"
 	"strings"
 	"sync"
@@ -194,10 +193,6 @@ func (c *BlocksCleaner) runCleanup(ctx context.Context, async bool) {
 		return
 	}
 
-	rand.Shuffle(len(allUsers), func(i, j int) {
-		allUsers[i], allUsers[j] = allUsers[j], allUsers[i]
-	})
-
 	doCleanup := func() {
 		err := c.cleanUsers(ctx, allUsers, isDeleted, logger)
 		c.instrumentFinishedCleanupRun(err, logger)
diff --git a/pkg/storage/tsdb/users_scanner.go b/pkg/storage/tsdb/users_scanner.go
index 0ade0f7a106..1b7ca534c8f 100644
--- a/pkg/storage/tsdb/users_scanner.go
+++ b/pkg/storage/tsdb/users_scanner.go
@@ -7,6 +7,7 @@ package tsdb
 
 import (
 	"context"
+	"math/rand"
 	"strings"
 
 	"github.com/go-kit/log"
@@ -46,6 +47,10 @@ func (s *UsersScanner) ScanUsers(ctx context.Context) (users, markedForDeletion
 		return nil, nil, err
 	}
 
+	rand.Shuffle(len(users), func(i, j int) {
+		users[i], users[j] = users[j], users[i]
+	})
+
 	// Check users for being owned by instance, and split users into non-deleted and deleted.
 	// We do these checks after listing all users, to improve cacheability of Iter (result is only cached at the end of Iter call).
 	for ix := 0; ix < len(users); {
diff --git a/pkg/storage/tsdb/users_scanner_test.go b/pkg/storage/tsdb/users_scanner_test.go
index 0af741d7778..7269e85f2f9 100644
--- a/pkg/storage/tsdb/users_scanner_test.go
+++ b/pkg/storage/tsdb/users_scanner_test.go
@@ -8,6 +8,7 @@ package tsdb
 import (
 	"context"
 	"errors"
+	"fmt"
 	"path"
 	"testing"
 
@@ -36,10 +37,10 @@ func TestUsersScanner_ScanUsers_ShouldReturnedOwnedUsersOnly(t *testing.T) {
 }
 
 func TestUsersScanner_ScanUsers_ShouldReturnUsersForWhichOwnerCheckOrTenantDeletionCheckFailed(t *testing.T) {
-	expected := []string{"user-1", "user-2"}
+	users := []string{"user-1", "user-2"}
 
 	bucketClient := &bucket.ClientMock{}
-	bucketClient.MockIter("", expected, nil)
+	bucketClient.MockIter("", users, nil)
 	bucketClient.MockExists(path.Join("user-1", TenantDeletionMarkPath), false, nil)
 	bucketClient.MockExists(path.Join("user-2", TenantDeletionMarkPath), false, errors.New("fail"))
 
@@ -50,11 +51,14 @@ func TestUsersScanner_ScanUsers_ShouldReturnUsersForWhichOwnerCheckOrTenantDelet
 	s := NewUsersScanner(bucketClient, isOwned, log.NewNopLogger())
 	actual, deleted, err := s.ScanUsers(context.Background())
 	require.NoError(t, err)
-	assert.Equal(t, expected, actual)
+	for _, userID := range users {
+		assert.Contains(t, actual, userID)
+	}
 	assert.Empty(t, deleted)
 }
 
 func TestUsersScanner_ScanUsers_ShouldNotReturnPrefixedUsedByMimirInternals(t *testing.T) {
+	users := []string{"user-1", "user-2"}
 	bucketClient := &bucket.ClientMock{}
 	bucketClient.MockIter("", []string{"user-1", "user-2", bucket.MimirInternalsPrefix}, nil)
 	bucketClient.MockExists(path.Join("user-1", TenantDeletionMarkPath), false, nil)
@@ -63,5 +67,32 @@ func TestUsersScanner_ScanUsers_ShouldNotReturnPrefixedUsedByMimirInternals(t *t
 	s := NewUsersScanner(bucketClient, AllUsers, log.NewNopLogger())
 	actual, _, err := s.ScanUsers(context.Background())
 	require.NoError(t, err)
-	assert.Equal(t, []string{"user-1", "user-2"}, actual)
+	for _, userID := range users {
+		assert.Contains(t, actual, userID)
+	}
+}
+
+func TestUsersScanner_ScanUsers_ShouldReturnRandomizedOrder(t *testing.T) {
+	var users = make([]string, 20)
+	for i := 0; i < len(users); i++ {
+		users[i] = fmt.Sprintf("user-%d", i)
+	}
+	bucketClient := &bucket.ClientMock{}
+	bucketClient.MockIter("", users, nil)
+	for i := 0; i < len(users); i++ {
+		bucketClient.MockExists(path.Join(users[i], TenantDeletionMarkPath), false, nil)
+	}
+
+	isOwned := func(_ string) (bool, error) {
+		return true, nil
+	}
+
+	s := NewUsersScanner(bucketClient, isOwned, log.NewNopLogger())
+	actual, deleted, err := s.ScanUsers(context.Background())
+	require.NoError(t, err)
+	for _, userID := range users {
+		assert.Contains(t, actual, userID)
+	}
+	assert.Empty(t, deleted)
+	assert.NotEqual(t, actual, users)
 }