MINOR: Simplify timeout handling in CoordinatorRuntime (#21334)

Remove the timeout parameter from all CoordinatorRuntime schedule
methods and use the writeTimeout field (configured via the builder)
instead. This simplifies the API since the timeout is now configured
once at runtime construction time.

- Rename defaultWriteTimeout to writeTimeout
- Rename withDefaultWriteTimeOut to withWriteTimeout
- Remove timeout parameter from scheduleWriteOperation,
scheduleWriteAllOperation, scheduleTransactionalWriteOperation, and
scheduleTransactionCompletion methods
- Update GroupCoordinator.completeTransaction to remove timeout
parameter
- Update all callers and tests accordingly

Reviewers: Sean Quah <squah@confluent.io>, Lianet Magrans
 <lmagrans@confluent.io>

Author: dajac

coordinator-common/src/main/java/org/apache/kafka/coordinator/common/runtime/CoordinatorExecutorImpl.java

@@ -23,7 +23,6 @@
 
 import org.slf4j.Logger;
 
-import java.time.Duration;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
@@ -36,21 +35,18 @@ private record TaskResult<R>(R result, Throwable exception) { }
     private final TopicPartition shard;
     private final CoordinatorRuntime<S, U> runtime;
     private final ExecutorService executor;
-    private final Duration writeTimeout;
     private final Map<String, TaskRunnable<?>> tasks = new ConcurrentHashMap<>();
 
     public CoordinatorExecutorImpl(
         LogContext logContext,
         TopicPartition shard,
         CoordinatorRuntime<S, U> runtime,
-        ExecutorService executor,
-        Duration writeTimeout
+        ExecutorService executor
     ) {
         this.log = logContext.logger(CoordinatorExecutorImpl.class);
         this.shard = shard;
         this.runtime = runtime;
         this.executor = executor;
-        this.writeTimeout = writeTimeout;
     }
 
     private <R> TaskResult<R> executeTask(TaskRunnable<R> task) {
@@ -83,7 +79,6 @@ public <R> boolean schedule(
             runtime.scheduleWriteOperation(
                 key,
                 shard,
-                writeTimeout,
                 coordinator -> {
                     // If the task associated with the key is not us, it means
                     // that the task was either replaced or cancelled. We stop.

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

@@ -113,7 +113,7 @@ public static class Builder<S extends CoordinatorShard<U>, U> {
         private CoordinatorShardBuilderSupplier<S, U> coordinatorShardBuilderSupplier;
         private Time time = Time.SYSTEM;
         private Timer timer;
-        private Duration defaultWriteTimeout;
+        private Duration writeTimeout;
         private CoordinatorRuntimeMetrics runtimeMetrics;
         private CoordinatorMetrics coordinatorMetrics;
         private Serializer<U> serializer;
@@ -162,8 +162,8 @@ public Builder<S, U> withTimer(Timer timer) {
             return this;
         }
 
-        public Builder<S, U> withDefaultWriteTimeOut(Duration defaultWriteTimeout) {
-            this.defaultWriteTimeout = defaultWriteTimeout;
+        public Builder<S, U> withWriteTimeout(Duration writeTimeout) {
+            this.writeTimeout = writeTimeout;
             return this;
         }
 
@@ -246,7 +246,7 @@ public CoordinatorRuntime<S, U> build() {
                 coordinatorShardBuilderSupplier,
                 time,
                 timer,
-                defaultWriteTimeout,
+                writeTimeout,
                 runtimeMetrics,
                 coordinatorMetrics,
                 serializer,
@@ -378,7 +378,7 @@ public void schedule(
                 @Override
                 public void run() {
                     String eventName = "Timeout(tp=" + tp + ", key=" + key + ")";
-                    CoordinatorWriteEvent<Void> event = new CoordinatorWriteEvent<>(eventName, tp, defaultWriteTimeout, coordinator -> {
+                    CoordinatorWriteEvent<Void> event = new CoordinatorWriteEvent<>(eventName, tp, writeTimeout, coordinator -> {
                         log.debug("Executing write event {} for timer {}.", eventName, key);
 
                         // If the task is different, it means that the timer has been
@@ -644,8 +644,7 @@ private CoordinatorContext(
                 logContext,
                 tp,
                 CoordinatorRuntime.this,
-                executorService,
-                defaultWriteTimeout
+                executorService
             );
             this.bufferSupplier = new BufferSupplier.GrowableBufferSupplier();
             this.cachedBufferSize = new AtomicLong(0);
@@ -1964,7 +1963,7 @@ public void onHighWatermarkUpdated(
     /**
      * The write operation timeout
      */
-    private final Duration defaultWriteTimeout;
+    private final Duration writeTimeout;
 
     /**
      * The coordinators keyed by topic partition.
@@ -2051,7 +2050,7 @@ public void onHighWatermarkUpdated(
      * @param coordinatorShardBuilderSupplier   The coordinator builder.
      * @param time                              The system time.
      * @param timer                             The system timer.
-     * @param defaultWriteTimeout               The write operation timeout.
+     * @param writeTimeout                      The write operation timeout.
      * @param runtimeMetrics                    The runtime metrics.
      * @param coordinatorMetrics                The coordinator metrics.
      * @param serializer                        The serializer.
@@ -2070,7 +2069,7 @@ private CoordinatorRuntime(
         CoordinatorShardBuilderSupplier<S, U> coordinatorShardBuilderSupplier,
         Time time,
         Timer timer,
-        Duration defaultWriteTimeout,
+        Duration writeTimeout,
         CoordinatorRuntimeMetrics runtimeMetrics,
         CoordinatorMetrics coordinatorMetrics,
         Serializer<U> serializer,
@@ -2083,7 +2082,7 @@ private CoordinatorRuntime(
         this.log = logContext.logger(CoordinatorRuntime.class);
         this.time = time;
         this.timer = timer;
-        this.defaultWriteTimeout = defaultWriteTimeout;
+        this.writeTimeout = writeTimeout;
         this.coordinators = new ConcurrentHashMap<>();
         this.processor = processor;
         this.partitionWriter = partitionWriter;
@@ -2198,7 +2197,6 @@ private void withActiveContextOrThrow(
      *
      * @param name      The name of the write operation.
      * @param tp        The address of the coordinator (aka its topic-partitions).
-     * @param timeout   The write operation timeout.
      * @param op        The write operation.
      *
      * @return A future that will be completed with the result of the write operation
@@ -2209,12 +2207,11 @@ private void withActiveContextOrThrow(
     public <T> CompletableFuture<T> scheduleWriteOperation(
         String name,
         TopicPartition tp,
-        Duration timeout,
         CoordinatorWriteOperation<S, T, U> op
     ) {
         throwIfNotRunning();
         log.debug("Scheduled execution of write operation {}.", name);
-        CoordinatorWriteEvent<T> event = new CoordinatorWriteEvent<>(name, tp, timeout, op);
+        CoordinatorWriteEvent<T> event = new CoordinatorWriteEvent<>(name, tp, writeTimeout, op);
         enqueueLast(event);
         return event.future;
     }
@@ -2223,7 +2220,6 @@ public <T> CompletableFuture<T> scheduleWriteOperation(
      * 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
@@ -2233,15 +2229,14 @@ public <T> CompletableFuture<T> scheduleWriteOperation(
      */
     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))
+            .map(tp -> scheduleWriteOperation(name, tp, op))
             .collect(Collectors.toList());
     }
 
@@ -2253,7 +2248,6 @@ public <T> List<CompletableFuture<T>> scheduleWriteAllOperation(
      * @param transactionalId   The transactional id.
      * @param producerId        The producer id.
      * @param producerEpoch     The producer epoch.
-     * @param timeout           The write operation timeout.
      * @param op                The write operation.
      * @param apiVersion        The Version of the Txn_Offset_Commit request
      *
@@ -2268,7 +2262,6 @@ public <T> CompletableFuture<T> scheduleTransactionalWriteOperation(
         String transactionalId,
         long producerId,
         short producerEpoch,
-        Duration timeout,
         CoordinatorWriteOperation<S, T, U> op,
         int apiVersion
     ) {
@@ -2288,7 +2281,7 @@ public <T> CompletableFuture<T> scheduleTransactionalWriteOperation(
                 producerId,
                 producerEpoch,
                 verificationGuard,
-                timeout,
+                writeTimeout,
                 op
             );
             enqueueLast(event);
@@ -2317,8 +2310,7 @@ public CompletableFuture<Void> scheduleTransactionCompletion(
         short producerEpoch,
         int coordinatorEpoch,
         TransactionResult result,
-        short transactionVersion,
-        Duration timeout
+        short transactionVersion
     ) {
         throwIfNotRunning();
         log.debug("Scheduled execution of transaction completion for {} with producer id={}, producer epoch={}, " +
@@ -2331,7 +2323,7 @@ public CompletableFuture<Void> scheduleTransactionCompletion(
             coordinatorEpoch,
             result,
             transactionVersion,
-            timeout
+            writeTimeout
         );
         enqueueLast(event);
         return event.future;

coordinator-common/src/test/java/org/apache/kafka/coordinator/common/runtime/CoordinatorExecutorImplTest.java

@@ -60,19 +60,17 @@ public void testTaskSuccessfulLifecycle() {
             LOG_CONTEXT,
             SHARD_PARTITION,
             runtime,
-            executorService,
-            WRITE_TIMEOUT
+            executorService
         );
 
         when(runtime.scheduleWriteOperation(
             eq(TASK_KEY),
             eq(SHARD_PARTITION),
-            eq(WRITE_TIMEOUT),
             any()
         )).thenAnswer(args -> {
             assertTrue(executor.isScheduled(TASK_KEY));
             CoordinatorRuntime.CoordinatorWriteOperation<CoordinatorShard<String>, Void, String> op =
-                args.getArgument(3);
+                args.getArgument(2);
             assertEquals(
                 new CoordinatorResult<>(List.of("record"), null),
                 op.generateRecordsAndResult(coordinatorShard)
@@ -120,18 +118,16 @@ public void testTaskFailedLifecycle() {
             LOG_CONTEXT,
             SHARD_PARTITION,
             runtime,
-            executorService,
-            WRITE_TIMEOUT
+            executorService
         );
 
         when(runtime.scheduleWriteOperation(
             eq(TASK_KEY),
             eq(SHARD_PARTITION),
-            eq(WRITE_TIMEOUT),
             any()
         )).thenAnswer(args -> {
             CoordinatorRuntime.CoordinatorWriteOperation<CoordinatorShard<String>, Void, String> op =
-                args.getArgument(3);
+                args.getArgument(2);
             assertEquals(
                 new CoordinatorResult<>(List.of(), null),
                 op.generateRecordsAndResult(coordinatorShard)
@@ -178,8 +174,7 @@ public void testTaskCancelledBeforeBeingExecuted() {
             LOG_CONTEXT,
             SHARD_PARTITION,
             runtime,
-            executorService,
-            WRITE_TIMEOUT
+            executorService
         );
 
         when(executorService.submit(any(Runnable.class))).thenAnswer(args -> {
@@ -223,21 +218,19 @@ public void testTaskCancelledAfterBeingExecutedButBeforeWriteOperationIsExecuted
             LOG_CONTEXT,
             SHARD_PARTITION,
             runtime,
-            executorService,
-            WRITE_TIMEOUT
+            executorService
         );
 
         when(runtime.scheduleWriteOperation(
             eq(TASK_KEY),
             eq(SHARD_PARTITION),
-            eq(WRITE_TIMEOUT),
             any()
         )).thenAnswer(args -> {
             // Cancel the task before running the write operation.
             executor.cancel(TASK_KEY);
 
             CoordinatorRuntime.CoordinatorWriteOperation<CoordinatorShard<String>, Void, String> op =
-                args.getArgument(3);
+                args.getArgument(2);
             Throwable ex = assertThrows(RejectedExecutionException.class, () -> op.generateRecordsAndResult(coordinatorShard));
             return FutureUtils.failedFuture(ex);
         });
@@ -278,14 +271,12 @@ public void testTaskSchedulingWriteOperationFailed() {
             LOG_CONTEXT,
             SHARD_PARTITION,
             runtime,
-            executorService,
-            WRITE_TIMEOUT
+            executorService
         );
 
         when(runtime.scheduleWriteOperation(
             eq(TASK_KEY),
             eq(SHARD_PARTITION),
-            eq(WRITE_TIMEOUT),
             any()
         )).thenReturn(FutureUtils.failedFuture(new Throwable("Oh no!")));
 
@@ -327,19 +318,17 @@ public void testCancelAllTasks() {
             LOG_CONTEXT,
             SHARD_PARTITION,
             runtime,
-            executorService,
-            WRITE_TIMEOUT
+            executorService
         );
 
         List<CoordinatorRuntime.CoordinatorWriteOperation<CoordinatorShard<String>, Void, String>> writeOperations = new ArrayList<>();
         List<CompletableFuture<Void>> writeFutures = new ArrayList<>();
         when(runtime.scheduleWriteOperation(
             anyString(),
             eq(SHARD_PARTITION),
-            eq(WRITE_TIMEOUT),
             any()
         )).thenAnswer(args -> {
-            writeOperations.add(args.getArgument(3));
+            writeOperations.add(args.getArgument(2));
             CompletableFuture<Void> writeFuture = new CompletableFuture<>();
             writeFutures.add(writeFuture);
             return writeFuture;