[server] Improve the implementation of Coordinator HA

Author: wuchong

fluss-common/src/main/java/org/apache/fluss/exception/NotCoordinatorLeaderException.java

@@ -17,7 +17,11 @@
 
 package org.apache.fluss.exception;
 
-/** Exception thrown when a request is sent to a stand by coordinator server. since: 0.9 */
+/**
+ * Exception thrown when a request is sent to a stand by coordinator server.
+ *
+ * @since 1.0
+ */
 public class NotCoordinatorLeaderException extends ApiException {
 
     private static final long serialVersionUID = 1L;

fluss-rpc/src/main/java/org/apache/fluss/rpc/gateway/CoordinatorGateway.java

@@ -40,6 +40,9 @@
 /** The entry point of RPC gateway interface for coordinator server. */
 public interface CoordinatorGateway extends RpcGateway, AdminGateway {
 
+    /** Checks if the current server is the leader. */
+    boolean isLeader();
+
     /**
      * AdjustIsr request to adjust (expend or shrink) the ISR set for request table bucket.
      *

fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/server/FlussRequestHandler.java

@@ -17,8 +17,11 @@
 
 package org.apache.fluss.rpc.netty.server;
 
+import org.apache.fluss.exception.NotCoordinatorLeaderException;
 import org.apache.fluss.rpc.RpcGatewayService;
+import org.apache.fluss.rpc.gateway.CoordinatorGateway;
 import org.apache.fluss.rpc.messages.ApiMessage;
+import org.apache.fluss.rpc.protocol.ApiKeys;
 import org.apache.fluss.rpc.protocol.ApiMethod;
 import org.apache.fluss.rpc.protocol.RequestType;
 
@@ -35,9 +38,11 @@ public class FlussRequestHandler implements RequestHandler<FlussRequest> {
     private static final Logger LOG = LoggerFactory.getLogger(FlussRequestHandler.class);
 
     private final RpcGatewayService service;
+    private final boolean isCoordinator;
 
     public FlussRequestHandler(RpcGatewayService service) {
         this.service = service;
+        this.isCoordinator = service instanceof CoordinatorGateway;
     }
 
     @Override
@@ -58,6 +63,16 @@ public void processRequest(FlussRequest request) {
                             request.isInternal(),
                             request.getAddress(),
                             request.getPrincipal()));
+            // check if the coordinator server is the current leader if the API is a coordinator
+            // TODO: we should only check coordinator APIs instead of all APIs
+            if (isCoordinator && api.getApiKey() != ApiKeys.API_VERSIONS) {
+                if (!((CoordinatorGateway) service).isLeader()) {
+                    request.fail(
+                            new NotCoordinatorLeaderException(
+                                    "This coordinator server is not the current leader."));
+                    return;
+                }
+            }
             // invoke the corresponding method on RpcGateway instance.
             CompletableFuture<?> responseFuture =
                     (CompletableFuture<?>) api.getMethod().invoke(service, message);

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorContext.java

@@ -710,10 +710,10 @@ public void resetContext() {
         tablesToBeDeleted.clear();
         coordinatorEpoch = 0;
         clearTablesState();
-        // clear the live tablet servers
         liveTabletServers.clear();
         shuttingDownTabletServers.clear();
         serverTags.clear();
+        liveCoordinatorServers.clear();
     }
 
     public int getTotalPartitionCount() {

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorLeaderElection.java

@@ -53,17 +53,15 @@ public class CoordinatorLeaderElection implements AutoCloseable {
     private final String serverId;
     private final LeaderLatch leaderLatch;
     private final AtomicBoolean isLeader = new AtomicBoolean(false);
-    private final CompletableFuture<Void> leaderReadyFuture = new CompletableFuture<>();
-    // Single-threaded executor to run leader init callbacks outside Curator's EventThread.
+    // Cached thread pool to run leader init/cleanup callbacks outside Curator's EventThread.
     // Curator's LeaderLatchListener callbacks run on its internal EventThread; performing
     // synchronous ZK operations there causes deadlock because ZK response dispatch also
-    // needs that same thread.
+    // needs that same thread. A cached pool is used because these callbacks are transient
+    // (only run during leadership transitions), so idle threads can be reclaimed.
     private final ExecutorService leaderCallbackExecutor;
     // Tracks the pending cleanup task so that init can wait for it to complete.
     private final AtomicReference<CompletableFuture<Void>> pendingCleanup =
             new AtomicReference<>(CompletableFuture.completedFuture(null));
-    private volatile Runnable initLeaderServices;
-    private volatile Consumer<Throwable> cleanupLeaderServices;
 
     public CoordinatorLeaderElection(ZooKeeperClient zkClient, String serverId) {
         this.serverId = serverId;
@@ -73,34 +71,35 @@ public CoordinatorLeaderElection(ZooKeeperClient zkClient, String serverId) {
                         ZkData.CoordinatorElectionZNode.path(),
                         String.valueOf(serverId));
         this.leaderCallbackExecutor =
-                Executors.newSingleThreadExecutor(
+                Executors.newCachedThreadPool(
                         r -> {
                             Thread t = new Thread(r, "coordinator-leader-callback-" + serverId);
+                            // Daemon threads ensure the JVM can exit even if close() is not
+                            // called. Orderly shutdown is handled by close() -> shutdownNow().
                             t.setDaemon(true);
                             return t;
                         });
     }
 
     /**
-     * Starts the leader election process asynchronously. The returned future completes when this
-     * server becomes the leader for the first time and initializes the leader services.
+     * Starts the leader election process asynchronously.
      *
      * <p>After the first election, the server will continue to participate in future elections.
      * When re-elected as leader, the initLeaderServices callback will be invoked again.
      *
      * @param initLeaderServices the callback to initialize leader services once elected
      * @param cleanupLeaderServices the callback to clean up leader services when losing leadership
-     * @return a CompletableFuture that completes when this server becomes leader for the first time
      */
-    public CompletableFuture<Void> startElectLeaderAsync(
+    public void startElectLeaderAsync(
             Runnable initLeaderServices, Consumer<Throwable> cleanupLeaderServices) {
-        this.initLeaderServices = initLeaderServices;
-        this.cleanupLeaderServices = cleanupLeaderServices;
-
         leaderLatch.addListener(
                 new LeaderLatchListener() {
                     @Override
                     public void isLeader() {
+                        // return if already marked as leader to avoid duplicate init calls
+                        if (isLeader.get()) {
+                            return;
+                        }
                         LOG.info("Coordinator server {} has become the leader.", serverId);
                         // Capture the pending cleanup future at this point so that
                         // init waits for it before proceeding.
@@ -124,15 +123,15 @@ public void isLeader() {
                                     }
                                     try {
                                         initLeaderServices.run();
-                                        leaderReadyFuture.complete(null);
                                     } catch (Exception e) {
                                         LOG.error(
                                                 "Failed to initialize leader services for server {}",
                                                 serverId,
                                                 e);
-                                        leaderReadyFuture.completeExceptionally(e);
                                     }
                                 });
+                        // Set leader flag before init completes, so when zk found this leader, the
+                        // leader can accept requests
                         isLeader.set(true);
                     }
 
@@ -142,31 +141,27 @@ public void notLeader() {
                             LOG.warn(
                                     "Coordinator server {} has lost the leadership, cleaning up leader services.",
                                     serverId);
-                            // Run cleanup on a separate daemon thread (NOT on the
-                            // leaderCallbackExecutor) to avoid blocking init tasks.
+                            // Submit cleanup to leaderCallbackExecutor. The cached thread
+                            // pool can spawn a new thread even if init is still running.
                             // The cleanup completion is tracked via pendingCleanup so
                             // that subsequent init waits for it.
                             CompletableFuture<Void> cleanupFuture = new CompletableFuture<>();
                             pendingCleanup.set(cleanupFuture);
-                            Thread cleanupThread =
-                                    new Thread(
-                                            () -> {
-                                                try {
-                                                    if (cleanupLeaderServices != null) {
-                                                        cleanupLeaderServices.accept(null);
-                                                    }
-                                                } catch (Exception e) {
-                                                    LOG.error(
-                                                            "Failed to cleanup leader services for server {}",
-                                                            serverId,
-                                                            e);
-                                                } finally {
-                                                    cleanupFuture.complete(null);
-                                                }
-                                            },
-                                            "coordinator-leader-cleanup-" + serverId);
-                            cleanupThread.setDaemon(true);
-                            cleanupThread.start();
+                            leaderCallbackExecutor.execute(
+                                    () -> {
+                                        try {
+                                            if (cleanupLeaderServices != null) {
+                                                cleanupLeaderServices.accept(null);
+                                            }
+                                        } catch (Exception e) {
+                                            LOG.error(
+                                                    "Failed to cleanup leader services for server {}",
+                                                    serverId,
+                                                    e);
+                                        } finally {
+                                            cleanupFuture.complete(null);
+                                        }
+                                    });
                         }
                     }
                 });
@@ -176,11 +171,7 @@ public void notLeader() {
             LOG.info("Coordinator server {} started leader election.", serverId);
         } catch (Exception e) {
             LOG.error("Failed to start LeaderLatch for server {}", serverId, e);
-            leaderReadyFuture.completeExceptionally(
-                    new RuntimeException("Leader election start failed", e));
         }
-
-        return leaderReadyFuture;
     }
 
     @Override
@@ -196,12 +187,6 @@ public void close() {
         }
 
         leaderCallbackExecutor.shutdownNow();
-
-        // Complete the future exceptionally if it hasn't been completed yet
-        if (!leaderReadyFuture.isDone()) {
-            leaderReadyFuture.completeExceptionally(
-                    new RuntimeException("Leader election closed for server " + serverId));
-        }
     }
 
     public boolean isLeader() {