From 7b8eda25a1556f4f874abffb446893240d2728af Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <tomas.dvorak@graylog.com>
Date: Tue, 20 Aug 2024 11:41:16 +0200
Subject: [PATCH] prevent possible race conditions in skipping index sets
 during rotation

---
 .../periodical/IndexRotationThread.java       |  7 ++-----
 .../periodical/IndexRotationThreadTest.java   | 20 +++++++++++++------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/graylog2-server/src/main/java/org/graylog2/periodical/IndexRotationThread.java b/graylog2-server/src/main/java/org/graylog2/periodical/IndexRotationThread.java
index 3d1cc3a00dfa..3047eb1c4175 100644
--- a/graylog2-server/src/main/java/org/graylog2/periodical/IndexRotationThread.java
+++ b/graylog2-server/src/main/java/org/graylog2/periodical/IndexRotationThread.java
@@ -91,12 +91,9 @@ public void initialize() {
     public void doRun() {
         // Point deflector to a new index if required.
         if (cluster.isConnected()) {
-            indexSetRegistry
-                    .getAll().stream()
-                    .filter(indexSet -> !isCurrentlyMigrated(indexSet))
-                    .forEach((indexSet) -> {
+            indexSetRegistry.forEach((indexSet) -> {
                 try {
-                    if (indexSet.getConfig().isWritable()) {
+                    if (indexSet.getConfig().isWritable() && !isCurrentlyMigrated(indexSet)) {
                         checkAndRepair(indexSet);
                         checkForRotation(indexSet);
                     } else {
diff --git a/graylog2-server/src/test/java/org/graylog2/periodical/IndexRotationThreadTest.java b/graylog2-server/src/test/java/org/graylog2/periodical/IndexRotationThreadTest.java
index 89696ce45337..75d5eaac5583 100644
--- a/graylog2-server/src/test/java/org/graylog2/periodical/IndexRotationThreadTest.java
+++ b/graylog2-server/src/test/java/org/graylog2/periodical/IndexRotationThreadTest.java
@@ -18,6 +18,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import jakarta.annotation.Nonnull;
+import jakarta.inject.Provider;
 import org.assertj.core.api.Assertions;
 import org.graylog2.datatiering.DataTieringOrchestrator;
 import org.graylog2.indexer.IndexSet;
@@ -41,10 +42,9 @@
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
-import jakarta.inject.Provider;
-
 import java.util.Arrays;
-import java.util.HashSet;
+import java.util.List;
+import java.util.function.Consumer;
 
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
@@ -174,9 +174,17 @@ private Cluster mockCluster(boolean connected) {
 
     @Nonnull
     private IndexSetRegistry mockIndexSetRegistry(IndexSet... indexSets) {
-        final IndexSetRegistry indexSetRegistry = Mockito.mock(IndexSetRegistry.class);
-        Mockito.when(indexSetRegistry.getAll()).thenReturn(new HashSet<>(Arrays.asList(indexSets)));
-        return indexSetRegistry;
+        final IndexSetRegistry registry = Mockito.mock(IndexSetRegistry.class);
+        final List<IndexSet> sets = Arrays.asList(indexSets);
+
+        // mock the Iterable forEach implementation :-/
+        Mockito.doAnswer(invocation -> {
+            Consumer<IndexSet> action = invocation.getArgument(0);
+            sets.forEach(action);
+            return null;
+        }).when(registry).forEach(Mockito.any());
+
+        return registry;
     }
 
     @Test