KAFKA-17593; [7/N] Introduce CoordinatorExecutor

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

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.common.runtime;
+
+import org.apache.kafka.common.KafkaException;
+
+/**
+ * An interface to schedule and cancel asynchronous tasks. The TaskRunnable
+ * interface defines the tasks to be executed in the executor and the
+ * TaskOperation defines the operation scheduled to the runtime to
+ * process the output of the executed task.
+ *
+ * @param <T> The record type.
+ */
+public interface CoordinatorExecutor<T> {
+    /**
+     * The task's runnable.
+     *
+     * @param <R> The return type.
+     */
+    interface TaskRunnable<R> {
+        R run() throws Throwable;
+    }
+
+    /**
+     * The task's write operation to handle the output
+     * of the task.
+     *
+     * @param <T> The record type.
+     * @param <R> The return type of the task.
+     */
+    interface TaskOperation<T, R> {
+        CoordinatorResult<Void, T> onComplete(
+            R result,
+            Throwable exception
+        ) throws KafkaException;
+    }
+
+    /**
+     * Schedule an asynchronous tasks. Note that only one task for a given key can
+     * be executed at the time.
+     *
+     * @param key       The key to identify the task.
+     * @param task      The task itself.
+     * @param operation The runtime operation to handle the output of the task.
+     * @return True if the task was scheduled; False otherwise.
+     *
+     * @param <R> The return type of the task.
+     */
+    <R> boolean schedule(
+        String key,
+        TaskRunnable<R> task,
+        TaskOperation<T, R> operation
+    );
+
+    /**
+     * Return true if the key is associated to a task; false otherwise.
+     *
+     * @param key The key to identify the task.
+     * @return A boolean indicating whether the task is scheduled or not.
+     */
+    boolean isScheduled(String key);
+
+    /**
+     * Cancel the given task
+     *
+     * @param key The key to identify the task.
+     */
+    void cancel(String key);
+}

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

@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.common.runtime;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.CoordinatorLoadInProgressException;
+import org.apache.kafka.common.errors.NotCoordinatorException;
+import org.apache.kafka.common.utils.LogContext;
+
+import org.slf4j.Logger;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+
+public class CoordinatorExecutorImpl<S extends CoordinatorShard<U>, U> implements CoordinatorExecutor<U> {
+    private static class TaskResult<R> {
+        final R result;
+        final Throwable exception;
+
+        TaskResult(
+            R result,
+            Throwable exception
+        ) {
+            this.result = result;
+            this.exception = exception;
+        }
+    }
+
+    private final Logger log;
+    private final TopicPartition shard;
+    private final CoordinatorRuntime<S, U> runtime;
+    private final ExecutorService executor;
+    private final Map<String, TaskRunnable<?>> tasks = new ConcurrentHashMap<>();
+
+    public CoordinatorExecutorImpl(
+        LogContext logContext,
+        TopicPartition shard,
+        CoordinatorRuntime<S, U> runtime,
+        ExecutorService executor
+    ) {
+        this.log = logContext.logger(CoordinatorExecutorImpl.class);
+        this.shard = shard;
+        this.runtime = runtime;
+        this.executor = executor;
+    }
+
+    private <R> TaskResult<R> executeTask(TaskRunnable<R> task) {
+        try {
+            return new TaskResult<>(task.run(), null);
+        } catch (Throwable ex) {
+            return new TaskResult<>(null, ex);
+        }
+    }
+
+    @Override
+    public <R> boolean schedule(
+        String key,
+        TaskRunnable<R> task,
+        TaskOperation<U, R> operation
+    ) {
+        // If there is already a task in-flight, we reject adding a new one.
+        if (tasks.containsKey(key)) return false;
+
+        // We use the task as a lock for the key.
+        tasks.put(key, task);
+
+        // Submit the task.
+        executor.submit(() -> {
+            // If the task associated with the task is not us, it means
+            // that the task was either replaced or cancelled. We stop.
+            if (tasks.get(key) != task) return;
+
+            // Executor the task.
+            final TaskResult<R> result = executeTask(task);
+
+            // Schedule the operation.
+            runtime.scheduleWriteOperation(
+                key,
+                shard,
+                Duration.ofMillis(Long.MAX_VALUE),
+                coordinator -> {
+                    // If the task associated with the task is not us, it means
+                    // that the task was either replaced or cancelled. We stop.
+                    if (!tasks.remove(key, task)) {
+                        return new CoordinatorResult<>(Collections.emptyList(), null);
+                    }
+
+                    // Call the underlying write operation with the result of the task.
+                    return operation.onComplete(result.result, result.exception);
+                }
+            ).exceptionally(exception -> {
+                if (exception instanceof NotCoordinatorException || exception instanceof CoordinatorLoadInProgressException) {
+                    log.debug("The write event for the task {} failed due to {}. Ignoring it because " +
+                        "the coordinator is not active.", key, exception.getMessage());
+                } else {
+                    log.error("The write event for the task {} failed due to {}. Ignoring it. ",
+                        key, exception.getMessage());
+                }
+                return null;
+            });
+        });
+
+        return true;
+    }
+
+    @Override
+    public boolean isScheduled(String key) {
+        return tasks.containsKey(key);
+    }
+
+    @Override
+    public void cancel(String key) {
+        tasks.remove(key);
+    }
+
+    public void cancelAll() {
+        Iterator<String> iterator = tasks.keySet().iterator();
+        while (iterator.hasNext()) {
+            iterator.remove();
+        }
+    }
+}

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

@@ -60,6 +60,7 @@
 import java.util.OptionalInt;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -118,6 +119,7 @@ public static class Builder<S extends CoordinatorShard<U>, U> {
         private Serializer<U> serializer;
         private Compression compression;
         private int appendLingerMs;
+        private ExecutorService executorService;
 
         public Builder<S, U> withLogPrefix(String logPrefix) {
             this.logPrefix = logPrefix;
@@ -189,6 +191,11 @@ public Builder<S, U> withAppendLingerMs(int appendLingerMs) {
             return this;
         }
 
+        public Builder<S, U> withExecutorService(ExecutorService executorService) {
+            this.executorService = executorService;
+            return this;
+        }
+
         public CoordinatorRuntime<S, U> build() {
             if (logPrefix == null)
                 logPrefix = "";
@@ -216,6 +223,8 @@ public CoordinatorRuntime<S, U> build() {
                 compression = Compression.NONE;
             if (appendLingerMs < 0)
                 throw new IllegalArgumentException("AppendLinger must be >= 0");
+            if (executorService == null)
+                throw new IllegalArgumentException("ExecutorService must be set.");
 
             return new CoordinatorRuntime<>(
                 logPrefix,
@@ -231,7 +240,8 @@ public CoordinatorRuntime<S, U> build() {
                 coordinatorMetrics,
                 serializer,
                 compression,
-                appendLingerMs
+                appendLingerMs,
+                executorService
             );
         }
     }
@@ -551,6 +561,11 @@ class CoordinatorContext {
          */
         final EventBasedCoordinatorTimer timer;
 
+        /**
+         * The coordinator executor.
+         */
+        final CoordinatorExecutorImpl<S, U> executor;
+
         /**
          * The current state.
          */
@@ -603,6 +618,7 @@ private CoordinatorContext(
             this.epoch = -1;
             this.deferredEventQueue = new DeferredEventQueue(logContext);
             this.timer = new EventBasedCoordinatorTimer(tp, logContext);
+            this.executor = new CoordinatorExecutorImpl<>(logContext, tp, CoordinatorRuntime.this, executorService);
             this.bufferSupplier = new BufferSupplier.GrowableBufferSupplier();
         }
 
@@ -633,6 +649,7 @@ private void transitionTo(
                             .withSnapshotRegistry(snapshotRegistry)
                             .withTime(time)
                             .withTimer(timer)
+                            .withExecutor(executor)
                             .withCoordinatorMetrics(coordinatorMetrics)
                             .withTopicPartition(tp)
                             .build(),
@@ -714,6 +731,7 @@ private void unload() {
                 highWatermarklistener = null;
             }
             timer.cancelAll();
+            executor.cancelAll();
             deferredEventQueue.failAll(Errors.NOT_COORDINATOR.exception());
             failCurrentBatch(Errors.NOT_COORDINATOR.exception());
             if (coordinator != null) {
@@ -1899,6 +1917,12 @@ public void onHighWatermarkUpdated(
      */
     private final int appendLingerMs;
 
+    /**
+     * The executor service used by the coordinator runtime to schedule
+     * asynchronous tasks.
+     */
+    private final ExecutorService executorService;
+
     /**
      * Atomic boolean indicating whether the runtime is running.
      */
@@ -1926,6 +1950,7 @@ public void onHighWatermarkUpdated(
      * @param serializer                        The serializer.
      * @param compression                       The compression codec.
      * @param appendLingerMs                    The append linger time in ms.
+     * @param executorService                   The executor service.
      */
     @SuppressWarnings("checkstyle:ParameterNumber")
     private CoordinatorRuntime(
@@ -1942,7 +1967,8 @@ private CoordinatorRuntime(
         CoordinatorMetrics coordinatorMetrics,
         Serializer<U> serializer,
         Compression compression,
-        int appendLingerMs
+        int appendLingerMs,
+        ExecutorService executorService
     ) {
         this.logPrefix = logPrefix;
         this.logContext = logContext;
@@ -1960,6 +1986,7 @@ private CoordinatorRuntime(
         this.serializer = serializer;
         this.compression = compression;
         this.appendLingerMs = appendLingerMs;
+        this.executorService = executorService;
     }
 
     /**
@@ -2423,6 +2450,7 @@ public void close() throws Exception {
             }
         });
         coordinators.clear();
+        executorService.shutdown();
         Utils.closeQuietly(runtimeMetrics, "runtime metrics");
         log.info("Coordinator runtime closed.");
     }

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

@@ -75,6 +75,17 @@ CoordinatorShardBuilder<S, U> withTimer(
         CoordinatorTimer<Void, U> timer
     );
 
+    /**
+     * Sets the coordinator executor.
+     *
+     * @param executor The coordinator executor.
+     *
+     * @return The builder.
+     */
+    CoordinatorShardBuilder<S, U> withExecutor(
+        CoordinatorExecutor<U> executor
+    );
+
     /**
      * Sets the coordinator metrics.
      *