support coordinator epoch2

zcoo · zcoo · commit 641614f9490d · 2025-11-13T15:27:45.000+08:00
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/MetadataManager.java
@@ -51,6 +51,7 @@
 import org.apache.fluss.server.zk.data.PartitionAssignment;
 import org.apache.fluss.server.zk.data.TableAssignment;
 import org.apache.fluss.server.zk.data.TableRegistration;
+import org.apache.fluss.server.zk.data.ZkVersion;
 import org.apache.fluss.shaded.zookeeper3.org.apache.zookeeper.KeeperException;
 import org.apache.fluss.utils.function.RunnableWithException;
 import org.apache.fluss.utils.function.ThrowingRunnable;
@@ -239,7 +240,9 @@ public void completeDeleteTable(long tableId) {
         // delete bucket assignments node, which will also delete the bucket state node,
         // so that all the zk nodes related to this table are deleted.
         rethrowIfIsNotNoNodeException(
-                () -> zookeeperClient.deleteTableAssignment(tableId),
+                () ->
+                        zookeeperClient.deleteTableAssignment(
+                                tableId, ZkVersion.MATCH_ANY_VERSION.getVersion()),
                 String.format("Delete tablet assignment meta fail for table %s.", tableId));
     }
 
@@ -248,7 +251,9 @@ public void completeDeletePartition(long partitionId) {
         // delete partition assignments node, which will also delete the bucket state node,
         // so that all the zk nodes related to this partition are deleted.
         rethrowIfIsNotNoNodeException(
-                () -> zookeeperClient.deletePartitionAssignment(partitionId),
+                () ->
+                        zookeeperClient.deletePartitionAssignment(
+                                partitionId, ZkVersion.MATCH_ANY_VERSION.getVersion()),
                 String.format("Delete tablet assignment meta fail for partition %s.", partitionId));
     }
 
@@ -301,7 +306,8 @@ public long createTable(
                     long tableId = zookeeperClient.getTableIdAndIncrement();
                     if (tableAssignment != null) {
                         // register table assignment
-                        zookeeperClient.registerTableAssignment(tableId, tableAssignment);
+                        zookeeperClient.registerTableAssignment(
+                                tableId, tableAssignment, ZkVersion.MATCH_ANY_VERSION.getVersion());
                     }
                     // register the table
                     zookeeperClient.registerTable(
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/TableBucketStateMachine.java b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/statemachine/TableBucketStateMachine.java
@@ -373,7 +373,8 @@ public void batchHandleOnlineChangeAndInitLeader(Set<TableBucket> tableBuckets)
         if (!tableBucketLeadAndIsrInfos.isEmpty()) {
             try {
                 zooKeeperClient.batchRegisterLeaderAndIsrForTablePartition(
-                        tableBucketLeadAndIsrInfos);
+                        tableBucketLeadAndIsrInfos,
+                        coordinatorContext.getCoordinatorEpochZkVersion());
                 registerSuccessList.addAll(tableBucketLeadAndIsrInfos);
             } catch (Exception e) {
                 LOG.error(
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZkEpoch.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZkEpoch.java
@@ -0,0 +1,38 @@
+/*
+ * 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.fluss.server.zk;
+
+/** Class for coordinator epoch and coordinator epoch zk version. */
+public class ZkEpoch {
+    private final int coordinatorEpoch;
+    private final int coordinatorEpochZkVersion;
+
+    public ZkEpoch(int coordinatorEpoch, int coordinatorEpochZkVersion) {
+        this.coordinatorEpoch = coordinatorEpoch;
+        this.coordinatorEpochZkVersion = coordinatorEpochZkVersion;
+    }
+
+    public int getCoordinatorEpoch() {
+        return coordinatorEpoch;
+    }
+
+    public int getCoordinatorEpochZkVersion() {
+        return coordinatorEpochZkVersion;
+    }
+}
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperClient.java
@@ -118,6 +118,13 @@
 /**
  * This class includes methods for write/read various metadata (leader address, tablet server
  * registration, table assignment, table, schema) in Zookeeper.
+ *
+ * <p>In some method, 'expectedZkVersion' is used to execute an epoch Zookeeper version check. We
+ * have the following principals to judge if it's necessary to execute epoch Zookeeper version
+ * check. If all condition met, we need to execute epoch Zookeeper version check. 1. The method
+ * create/modify/delete Zk node. 2. It's executed by coordinator server. 3. It is about
+ * metadata(table/partition/leaderAndIsr) rather than server info or ACL info. 4. The Zk node is
+ * persistent rather than ephemeral.
  */
 @Internal
 public class ZooKeeperClient implements AutoCloseable {
@@ -182,9 +189,9 @@ public Optional<Integer> fenceBecomeCoordinatorLeader(int coordinatorId) throws
             ensureEpochZnodeExists();
 
             try {
-                Tuple2<Integer, Integer> getEpoch = getCurrentEpoch();
-                int currentEpoch = getEpoch.f0;
-                int currentVersion = getEpoch.f1;
+                ZkEpoch getEpoch = getCurrentEpoch();
+                int currentEpoch = getEpoch.getCoordinatorEpoch();
+                int currentVersion = getEpoch.getCoordinatorEpochZkVersion();
                 int newEpoch = currentEpoch + 1;
                 LOG.info(
                         "Coordinator leader {} tries to update epoch. Current epoch={}, Zookeeper version={}, new epoch={}",
@@ -254,15 +261,15 @@ public void ensureEpochZnodeExists() throws Exception {
     }
 
     /** Get epoch now in ZK. */
-    public Tuple2<Integer, Integer> getCurrentEpoch() throws Exception {
+    public ZkEpoch getCurrentEpoch() throws Exception {
         Stat currentStat = new Stat();
         byte[] bytes =
                 zkClient.getData()
                         .storingStatIn(currentStat)
                         .forPath(ZkData.CoordinatorEpochZNode.path());
         int currentEpoch = ZkData.CoordinatorEpochZNode.decode(bytes);
         int currentVersion = currentStat.getVersion();
-        return new Tuple2<>(currentEpoch, currentVersion);
+        return new ZkEpoch(currentEpoch, currentVersion);
     }
 
     // --------------------------------------------------------------------------------------------
@@ -318,13 +325,13 @@ public int[] getSortedTabletServerList() throws Exception {
     // --------------------------------------------------------------------------------------------
 
     /** Register table assignment to ZK. */
-    public void registerTableAssignment(long tableId, TableAssignment tableAssignment)
-            throws Exception {
+    public void registerTableAssignment(
+            long tableId, TableAssignment tableAssignment, int expectedZkVersion) throws Exception {
         String path = TableIdZNode.path(tableId);
-        zkClient.create()
-                .creatingParentsIfNeeded()
-                .withMode(CreateMode.PERSISTENT)
-                .forPath(path, TableIdZNode.encode(tableAssignment));
+        byte[] data = TableIdZNode.encode(tableAssignment);
+
+        createRecursiveWithEpochCheck(path, data, expectedZkVersion, false);
+
         LOG.info("Registered table assignment {} for table id {}.", tableAssignment, tableId);
     }
 
@@ -373,22 +380,27 @@ public Map<Long, PartitionAssignment> getPartitionsAssignments(Collection<Long>
                 "partition assignment");
     }
 
-    public void updateTableAssignment(long tableId, TableAssignment tableAssignment)
-            throws Exception {
+    public void updateTableAssignment(
+            long tableId, TableAssignment tableAssignment, int expectedZkVersion) throws Exception {
         String path = TableIdZNode.path(tableId);
-        zkClient.setData().forPath(path, TableIdZNode.encode(tableAssignment));
+        byte[] data = TableIdZNode.encode(tableAssignment);
+        CuratorOp updateOp = zkOp.updateOp(path, data);
+        List<CuratorOp> ops = wrapRequestWithEpochCheck(updateOp, expectedZkVersion);
+
+        zkClient.transaction().forOperations(ops);
         LOG.info("Updated table assignment {} for table id {}.", tableAssignment, tableId);
     }
 
-    public void deleteTableAssignment(long tableId) throws Exception {
+    public void deleteTableAssignment(long tableId, int expectedZkVersion) throws Exception {
         String path = TableIdZNode.path(tableId);
-        zkClient.delete().deletingChildrenIfNeeded().forPath(path);
+        deleteRecursiveWithEpochCheck(path, expectedZkVersion, false);
         LOG.info("Deleted table assignment for table id {}.", tableId);
     }
 
-    public void deletePartitionAssignment(long partitionId) throws Exception {
+    public void deletePartitionAssignment(long partitionId, int expectedZkVersion)
+            throws Exception {
         String path = PartitionIdZNode.path(partitionId);
-        zkClient.delete().deletingChildrenIfNeeded().forPath(path);
+        deleteRecursiveWithEpochCheck(path, expectedZkVersion, false);
         LOG.info("Deleted table assignment for partition id {}.", partitionId);
     }
 
@@ -404,12 +416,13 @@ public void registerLeaderAndIsr(
         String path = LeaderAndIsrZNode.path(tableBucket);
         byte[] data = LeaderAndIsrZNode.encode(leaderAndIsr);
 
-        createRecursive(path, data, expectedZkVersion, false);
+        createRecursiveWithEpochCheck(path, data, expectedZkVersion, false);
         LOG.info("Registered {} for bucket {} in Zookeeper.", leaderAndIsr, tableBucket);
     }
 
     public void batchRegisterLeaderAndIsrForTablePartition(
-            List<RegisterTableBucketLeadAndIsrInfo> registerList) throws Exception {
+            List<RegisterTableBucketLeadAndIsrInfo> registerList, int expectedZkVersion)
+            throws Exception {
         if (registerList.isEmpty()) {
             return;
         }
@@ -445,12 +458,14 @@ public void batchRegisterLeaderAndIsrForTablePartition(
             ops.add(parentNodeCreate);
             ops.add(currentNodeCreate);
             if (ops.size() == MAX_BATCH_SIZE) {
-                zkClient.transaction().forOperations(ops);
+                List<CuratorOp> wrapOps = wrapRequestsWithEpochCheck(ops, expectedZkVersion);
+                zkClient.transaction().forOperations(wrapOps);
                 ops.clear();
             }
         }
         if (!ops.isEmpty()) {
-            zkClient.transaction().forOperations(ops);
+            List<CuratorOp> wrapOps = wrapRequestsWithEpochCheck(ops, expectedZkVersion);
+            zkClient.transaction().forOperations(wrapOps);
         }
         LOG.info(
                 "Batch registered leadAndIsr for tableId: {}, partitionId: {}, partitionName: {}  in Zookeeper.",
@@ -489,7 +504,7 @@ public void updateLeaderAndIsr(
         byte[] data = LeaderAndIsrZNode.encode(leaderAndIsr);
 
         CuratorOp updateOp = zkOp.updateOp(path, data);
-        List<CuratorOp> ops = wrapRequestWithCoordinatorEpochCheck(updateOp, expectedZkVersion);
+        List<CuratorOp> ops = wrapRequestWithEpochCheck(updateOp, expectedZkVersion);
 
         zkClient.transaction().forOperations(ops);
         LOG.info("Updated {} for bucket {} in Zookeeper.", leaderAndIsr, tableBucket);
@@ -511,22 +526,22 @@ public void batchUpdateLeaderAndIsr(
             byte[] data = LeaderAndIsrZNode.encode(leaderAndIsr);
             CuratorOp updateOp = zkClient.transactionOp().setData().forPath(path, data);
             ops.add(updateOp);
-            if (ops.size() == MAX_BATCH_SIZE - 1) {
-                List<CuratorOp> wrapOps =
-                        wrapRequestsWithCoordinatorEpochCheck(ops, expectedZkVersion);
+            if (ops.size() == MAX_BATCH_SIZE) {
+                List<CuratorOp> wrapOps = wrapRequestsWithEpochCheck(ops, expectedZkVersion);
                 zkClient.transaction().forOperations(wrapOps);
                 ops.clear();
             }
         }
         if (!ops.isEmpty()) {
-            List<CuratorOp> wrapOps = wrapRequestsWithCoordinatorEpochCheck(ops, expectedZkVersion);
+            List<CuratorOp> wrapOps = wrapRequestsWithEpochCheck(ops, expectedZkVersion);
             zkClient.transaction().forOperations(wrapOps);
         }
     }
 
-    public void deleteLeaderAndIsr(TableBucket tableBucket) throws Exception {
+    public void deleteLeaderAndIsr(TableBucket tableBucket, int expectedZkVersion)
+            throws Exception {
         String path = LeaderAndIsrZNode.path(tableBucket);
-        zkClient.delete().forPath(path);
+        deleteRecursiveWithEpochCheck(path, expectedZkVersion, false);
         LOG.info("Deleted LeaderAndIsr for bucket {} in Zookeeper.", tableBucket);
     }
 
@@ -1601,11 +1616,19 @@ public static <K> Map<K, List<String>> processGetChildrenResponses(
         return result;
     }
 
-    public void createRecursive(
+    /**
+     * create a node (recursively if parent path not exists) with Zk epoch version check.
+     *
+     * @param path the path to create
+     * @param data the data to write
+     * @param throwIfPathExists whether to throw exception if path exist
+     * @throws Exception if any error occurs
+     */
+    public void createRecursiveWithEpochCheck(
             String path, byte[] data, int expectedZkVersion, boolean throwIfPathExists)
             throws Exception {
         CuratorOp createOp = zkOp.createOp(path, data, CreateMode.PERSISTENT);
-        List<CuratorOp> ops = wrapRequestWithCoordinatorEpochCheck(createOp, expectedZkVersion);
+        List<CuratorOp> ops = wrapRequestWithEpochCheck(createOp, expectedZkVersion);
 
         try {
             // try to directly create
@@ -1622,19 +1645,48 @@ public void createRecursive(
                 throw new IllegalArgumentException("Invalid path {}" + path);
             }
             String parentPath = path.substring(0, indexOfLastSlash);
-            createRecursive(parentPath, null, expectedZkVersion, throwIfPathExists);
+            createRecursiveWithEpochCheck(parentPath, null, expectedZkVersion, throwIfPathExists);
             // After creating parent, retry creating the original path
             zkClient.transaction().forOperations(ops);
         }
     }
 
-    public List<CuratorOp> wrapRequestWithCoordinatorEpochCheck(
-            CuratorOp request, int expectedZkVersion) throws Exception {
-        return wrapRequestsWithCoordinatorEpochCheck(
-                Collections.singletonList(request), expectedZkVersion);
+    /**
+     * Delete a node (and recursively delete children) with Zk epoch version check.
+     *
+     * @param path the path to delete
+     * @param expectedZkVersion the expected coordinator epoch zk version
+     * @param throwIfPathNotExists whether to throw exception if path does not exist
+     * @throws Exception if any error occurs
+     */
+    public void deleteRecursiveWithEpochCheck(
+            String path, int expectedZkVersion, boolean throwIfPathNotExists) throws Exception {
+        // delete children recursively
+        List<String> children = getChildren(path);
+        for (String child : children) {
+            deleteRecursiveWithEpochCheck(path + "/" + child, expectedZkVersion, false);
+        }
+
+        CuratorOp deleteOp = zkOp.deleteOp(path);
+        List<CuratorOp> ops = wrapRequestWithEpochCheck(deleteOp, expectedZkVersion);
+
+        try {
+            // delete itself
+            zkClient.transaction().forOperations(ops);
+        } catch (KeeperException.NoNodeException e) {
+            // should exist
+            if (throwIfPathNotExists) {
+                throw e;
+            }
+        }
+    }
+
+    public List<CuratorOp> wrapRequestWithEpochCheck(CuratorOp request, int expectedZkVersion)
+            throws Exception {
+        return wrapRequestsWithEpochCheck(Collections.singletonList(request), expectedZkVersion);
     }
 
-    public List<CuratorOp> wrapRequestsWithCoordinatorEpochCheck(
+    public List<CuratorOp> wrapRequestsWithEpochCheck(
             List<CuratorOp> requestList, int expectedZkVersion) throws Exception {
         if (ZkVersion.MATCH_ANY_VERSION.getVersion() == expectedZkVersion) {
             return requestList;
diff --git a/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperOp.java b/fluss-server/src/main/java/org/apache/fluss/server/zk/ZooKeeperOp.java
@@ -44,6 +44,10 @@ public CuratorOp updateOp(String path, byte[] data) throws Exception {
         return zkClient.transactionOp().setData().forPath(path, data);
     }
 
+    public CuratorOp deleteOp(String path) throws Exception {
+        return zkClient.transactionOp().delete().forPath(path);
+    }
+
     public static List<CuratorOp> multiRequest(CuratorOp op1, CuratorOp op2) {
         List<CuratorOp> ops = new ArrayList<>();
         ops.add(op1);
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorServerElectionTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/CoordinatorServerElectionTest.java
@@ -91,7 +91,7 @@ void testCoordinatorServerElection() throws Exception {
             }
         }
         assertThat(elected).isNotNull();
-        assertThat(zookeeperClient.getCurrentEpoch().f0)
+        assertThat(zookeeperClient.getCurrentEpoch().getCoordinatorEpoch())
                 .isEqualTo(CoordinatorContext.INITIAL_COORDINATOR_EPOCH);
         elected.close();
         elected.start();
@@ -101,7 +101,7 @@ void testCoordinatorServerElection() throws Exception {
         CoordinatorAddress secondLeaderAddress =
                 zookeeperClient.getCoordinatorLeaderAddress().get();
         assertThat(secondLeaderAddress).isNotEqualTo(firstLeaderAddress);
-        assertThat(zookeeperClient.getCurrentEpoch().f0)
+        assertThat(zookeeperClient.getCurrentEpoch().getCoordinatorEpoch())
                 .isEqualTo(CoordinatorContext.INITIAL_COORDINATOR_EPOCH + 1);
 
         // kill other 2 coordinator servers except the first one
@@ -115,7 +115,7 @@ void testCoordinatorServerElection() throws Exception {
         CoordinatorAddress thirdLeaderAddress = zookeeperClient.getCoordinatorLeaderAddress().get();
 
         assertThat(thirdLeaderAddress.getId()).isEqualTo(firstLeaderAddress.getId());
-        assertThat(zookeeperClient.getCurrentEpoch().f0)
+        assertThat(zookeeperClient.getCurrentEpoch().getCoordinatorEpoch())
                 .isEqualTo(CoordinatorContext.INITIAL_COORDINATOR_EPOCH + 2);
     }
 
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TableManagerTest.java b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/TableManagerTest.java
diff --git a/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java b/fluss-server/src/test/java/org/apache/fluss/server/zk/ZooKeeperClientTest.java
