MINOR: Add read/write all operation

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java

@@ -80,7 +80,6 @@
 import java.util.Map;
 import java.util.OptionalInt;
 import java.util.Properties;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutionException;
@@ -498,29 +497,24 @@ public CompletableFuture<ListGroupsResponseData> listGroups(
             );
         }
 
-        final Set<TopicPartition> existingPartitionSet = runtime.partitions();
-
-        if (existingPartitionSet.isEmpty()) {
-            return CompletableFuture.completedFuture(new ListGroupsResponseData());
-        }
-
-        final List<CompletableFuture<List<ListGroupsResponseData.ListedGroup>>> futures =
-            new ArrayList<>();
-
-        for (TopicPartition tp : existingPartitionSet) {
-            futures.add(runtime.scheduleReadOperation(
+        final List<CompletableFuture<List<ListGroupsResponseData.ListedGroup>>> futures = FutureUtils.mapExceptionally(
+            runtime.scheduleReadAllOperation(
                 "list-groups",
-                tp,
-                (coordinator, lastCommittedOffset) -> coordinator.listGroups(request.statesFilter(), request.typesFilter(), lastCommittedOffset)
-            ).exceptionally(exception -> {
+                (coordinator, lastCommittedOffset) -> coordinator.listGroups(
+                    request.statesFilter(),
+                    request.typesFilter(),
+                    lastCommittedOffset
+                )
+            ),
+            exception -> {
                 exception = Errors.maybeUnwrapException(exception);
                 if (exception instanceof NotCoordinatorException) {
                     return Collections.emptyList();
                 } else {
                     throw new CompletionException(exception);
                 }
-            }));
-        }
+            }
+        );
 
         return FutureUtils
             .combineFutures(futures, ArrayList::new, List::addAll)
@@ -963,23 +957,21 @@ public void onPartitionsDeleted(
     ) throws ExecutionException, InterruptedException {
         throwIfNotActive();
 
-        final Set<TopicPartition> existingPartitionSet = runtime.partitions();
-        final List<CompletableFuture<Void>> futures = new ArrayList<>(existingPartitionSet.size());
-
-        existingPartitionSet.forEach(partition -> futures.add(
-            runtime.scheduleWriteOperation(
-                "on-partition-deleted",
-                partition,
-                Duration.ofMillis(config.offsetCommitTimeoutMs),
-                coordinator -> coordinator.onPartitionsDeleted(topicPartitions)
-            ).exceptionally(exception -> {
-                log.error("Could not delete offsets for deleted partitions {} in coordinator {} due to: {}.",
-                    partition, partition, exception.getMessage(), exception);
-                return null;
-            })
-        ));
-
-        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get();
+        CompletableFuture.allOf(
+            FutureUtils.mapExceptionally(
+                runtime.scheduleWriteAllOperation(
+                    "on-partition-deleted",
+                    Duration.ofMillis(config.offsetCommitTimeoutMs),
+                    coordinator -> coordinator.onPartitionsDeleted(topicPartitions)
+                ),
+                exception -> {
+                    log.error("Could not delete offsets for deleted partitions {} due to: {}.",
+                        topicPartitions, exception.getMessage(), exception
+                    );
+                    return null;
+                }
+            ).toArray(new CompletableFuture[0])
+        ).get();
     }
 
     /**

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java

@@ -41,19 +41,19 @@
 
 import java.time.Duration;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.OptionalInt;
 import java.util.OptionalLong;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 /**
  * The CoordinatorRuntime provides a framework to implement coordinators such as the group coordinator
@@ -1446,6 +1446,32 @@ public <T> CompletableFuture<T> scheduleWriteOperation(
         return event.future;
     }
 
+    /**
+     * Schedule a write operation for each coordinator.
+     *
+     * @param name      The name of the write operation.
+     * @param timeout   The write operation timeout.
+     * @param op        The write operation.
+     *
+     * @return A list of futures where each future will be completed with the result of the write operation
+     * when the operation is completed or an exception if the write operation failed.
+     *
+     * @param <T> The type of the result.
+     */
+    public <T> List<CompletableFuture<T>> scheduleWriteAllOperation(
+        String name,
+        Duration timeout,
+        CoordinatorWriteOperation<S, T, U> op
+    ) {
+        throwIfNotRunning();
+        log.debug("Scheduled execution of write all operation {}.", name);
+        return coordinators
+            .keySet()
+            .stream()
+            .map(tp -> scheduleWriteOperation(name, tp, timeout, op))
+            .collect(Collectors.toList());
+    }
+
     /**
      * Schedules a transactional write operation.
      *
@@ -1535,12 +1561,12 @@ public CompletableFuture<Void> scheduleTransactionCompletion(
     /**
      * Schedules a read operation.
      *
-     * @param name  The name of the write operation.
+     * @param name  The name of the read operation.
      * @param tp    The address of the coordinator (aka its topic-partitions).
      * @param op    The read operation.
      *
      * @return A future that will be completed with the result of the read operation
-     * when the operation is completed or an exception if the write operation failed.
+     * when the operation is completed or an exception if the read operation failed.
      *
      * @param <T> The type of the result.
      */
@@ -1556,6 +1582,30 @@ public <T> CompletableFuture<T> scheduleReadOperation(
         return event.future;
     }
 
+    /**
+     * Schedules a read operation for each coordinator.
+     *
+     * @param name  The name of the read operation.
+     * @param op    The read operation.
+     *
+     * @return A list of futures where each future will be completed with the result of the read operation
+     * when the operation is completed or an exception if the read operation failed.
+     *
+     * @param <T> The type of the result.
+     */
+    public <T> List<CompletableFuture<T>> scheduleReadAllOperation(
+        String name,
+        CoordinatorReadOperation<S, T> op
+    ) {
+        throwIfNotRunning();
+        log.debug("Scheduled execution of read all operation {}.", name);
+        return coordinators
+            .keySet()
+            .stream()
+            .map(tp -> scheduleReadOperation(name, tp, op))
+            .collect(Collectors.toList());
+    }
+
     /**
      * Schedules an internal event.
      *
@@ -1572,15 +1622,6 @@ private void scheduleInternalOperation(
         enqueue(new CoordinatorInternalEvent(name, tp, op));
     }
 
-    /**
-     * @return The topic partitions of the coordinators currently registered in the
-     * runtime.
-     */
-    public Set<TopicPartition> partitions() {
-        throwIfNotRunning();
-        return new HashSet<>(coordinators.keySet());
-    }
-
     /**
      * Schedules the loading of a coordinator. This is called when the broker is elected as
      * the leader for a partition.