Send message batches in parallel in AwsStateStoreCommitterDriver

gchq · Aug 12, 2024 · e3d92d7 · e3d92d7
1 parent 09698a9
commit e3d92d7
Showing 1 changed file with 3 additions and 12 deletions.
diff --git a/...ers/src/main/java/sleeper/systemtest/drivers/statestore/AwsStateStoreCommitterDriver.java b/...ers/src/main/java/sleeper/systemtest/drivers/statestore/AwsStateStoreCommitterDriver.java
@@ -19,11 +19,11 @@
 import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
 import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
 
+import sleeper.core.util.SplitIntoBatches;
 import sleeper.systemtest.dsl.instance.SystemTestInstanceContext;
 import sleeper.systemtest.dsl.statestore.StateStoreCommitMessage;
 import sleeper.systemtest.dsl.statestore.StateStoreCommitterDriver;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.UUID;
 import java.util.stream.Stream;
@@ -43,17 +43,8 @@ public AwsStateStoreCommitterDriver(SystemTestInstanceContext instance, AmazonSQ
 
     @Override
     public void sendCommitMessages(Stream<StateStoreCommitMessage> messages) {
-        List<StateStoreCommitMessage> batch = new ArrayList<>();
-        messages.forEach(message -> {
-            batch.add(message);
-            if (batch.size() == 10) {
-                sendMessageBatch(batch);
-                batch.clear();
-            }
-        });
-        if (!batch.isEmpty()) {
-            sendMessageBatch(batch);
-        }
+        SplitIntoBatches.toListsOfSize(10, messages)
+                .parallel().forEach(this::sendMessageBatch);
     }
 
     private void sendMessageBatch(List<StateStoreCommitMessage> batch) {