Fix for RoundRobinStorageLocationSelectorStrategy not to pick the same storage location over and again (#8634)

* Fix for 8614: iterators returned by strategy don't share iteration state.

* Adding comments

server/src/main/java/org/apache/druid/segment/loading/RoundRobinStorageLocationSelectorStrategy.java

@@ -48,6 +48,9 @@ public class RoundRobinStorageLocationSelectorStrategy implements StorageLocatio
 
       private final int numStorageLocations = storageLocations.size();
       private int remainingIterations = numStorageLocations;
+      // Each call to this methods starts with a different startIndex to avoid the same location being picked up over
+      // again. See https://github.com/apache/incubator-druid/issues/8614.
+      private int i = startIndex.getAndUpdate(n -> (n + 1) % numStorageLocations);
 
       @Override
       public boolean hasNext()
@@ -62,8 +65,10 @@ public class RoundRobinStorageLocationSelectorStrategy implements StorageLocatio
           throw new NoSuchElementException();
         }
         remainingIterations--;
-        final StorageLocation nextLocation =
+        final StorageLocation nextLocation = storageLocations.get(i++);
-            storageLocations.get(startIndex.getAndUpdate(n -> (n + 1) % numStorageLocations));
+        if (i == numStorageLocations) {
+          i = 0;
+        }
         return nextLocation;
       }
     };

server/src/test/java/org/apache/druid/segment/loading/StorageLocationSelectorStrategyTest.java

@@ -110,16 +110,7 @@ public class StorageLocationSelectorStrategyTest
 
     StorageLocationSelectorStrategy roundRobinStrategy = new RoundRobinStorageLocationSelectorStrategy(storageLocations);
 
-    iterateLocs(localStorageFolder1, localStorageFolder2, localStorageFolder3, roundRobinStrategy);
+    // First call to getLocations()
-    iterateLocs(localStorageFolder1, localStorageFolder2, localStorageFolder3, roundRobinStrategy);
-    iterateLocs(localStorageFolder1, localStorageFolder2, localStorageFolder3, roundRobinStrategy);
-    iterateLocs(localStorageFolder1, localStorageFolder2, localStorageFolder3, roundRobinStrategy);
-    iterateLocs(localStorageFolder1, localStorageFolder2, localStorageFolder3, roundRobinStrategy);
-  }
-
-  private void iterateLocs(File localStorageFolder1, File localStorageFolder2, File localStorageFolder3,
-                           StorageLocationSelectorStrategy roundRobinStrategy)
-  {
     Iterator<StorageLocation> locations = roundRobinStrategy.getLocations();
 
     StorageLocation loc1 = locations.next();
@@ -133,6 +124,60 @@ public class StorageLocationSelectorStrategyTest
     StorageLocation loc3 = locations.next();
     Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_3",
         localStorageFolder3, loc3.getPath());
+
+
+    // Second call to getLocations()
+    locations = roundRobinStrategy.getLocations();
+
+    loc2 = locations.next();
+    Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_2",
+        localStorageFolder2, loc2.getPath());
+
+    loc3 = locations.next();
+    Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_3",
+        localStorageFolder3, loc3.getPath());
+
+    loc1 = locations.next();
+    Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_1",
+        localStorageFolder1, loc1.getPath());
+  }
+
+  @Test
+  public void testRoundRobinLocationSelectorStrategyMultipleCallsToGetLocations() throws Exception
+  {
+    List<StorageLocation> storageLocations = new ArrayList<>();
+
+    final File localStorageFolder1 = tmpFolder.newFolder("local_storage_folder_1");
+    final File localStorageFolder2 = tmpFolder.newFolder("local_storage_folder_2");
+    final File localStorageFolder3 = tmpFolder.newFolder("local_storage_folder_3");
+
+    storageLocations.add(new StorageLocation(localStorageFolder1, 10000000000L, null));
+    storageLocations.add(new StorageLocation(localStorageFolder2, 10000000000L, null));
+    storageLocations.add(new StorageLocation(localStorageFolder3, 10000000000L, null));
+
+    StorageLocationSelectorStrategy roundRobinStrategy = new RoundRobinStorageLocationSelectorStrategy(storageLocations);
+
+    Iterator<StorageLocation> locations = roundRobinStrategy.getLocations();
+
+    StorageLocation loc1 = locations.next();
+    Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_1",
+        localStorageFolder1, loc1.getPath());
+
+    locations = roundRobinStrategy.getLocations();
+
+    StorageLocation loc2 = locations.next();
+    Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_2",
+        localStorageFolder2, loc2.getPath());
+
+    locations = roundRobinStrategy.getLocations();
+
+    StorageLocation loc3 = locations.next();
+    Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_3",
+        localStorageFolder3, loc3.getPath());
+
+    loc1 = locations.next();
+    Assert.assertEquals("The next element of the iterator should point to path local_storage_folder_1",
+        localStorageFolder1, loc1.getPath());
   }
 
 }