[server] Use async zk operations to improve performance of CoordinatorServer initCoordinatorContext() (#1381)

Author: zcoo

fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java

@@ -463,6 +463,12 @@ public class ConfigOptions {
     // ------------------------------------------------------------------------
     //  ZooKeeper Client Settings
     // ------------------------------------------------------------------------
+    public static final ConfigOption<Integer> ZOOKEEPER_MAX_INFLIGHT_REQUESTS =
+            key("zookeeper.client.max-inflight-requests")
+                    .intType()
+                    .defaultValue(100)
+                    .withDescription(
+                            "The maximum number of unacknowledged requests the client will send to ZooKeeper before blocking.");
 
     public static final ConfigOption<Duration> ZOOKEEPER_SESSION_TIMEOUT =
             key("zookeeper.client.session-timeout")

fluss-server/src/main/java/org/apache/fluss/server/authorizer/DefaultAuthorizer.java

@@ -388,7 +388,7 @@ private void updateResourceAcl(
 
                 } else {
                     LOG.trace("Deleting path for {} because it had no ACLs remaining", resource);
-                    zooKeeperClient.contitionalDeleteResourceAcl(
+                    zooKeeperClient.conditionalDeleteResourceAcl(
                             resource, currentVersionedAcls.zkVersion);
                 }
                 writeComplete = true;

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

@@ -272,8 +272,12 @@ private void initCoordinatorContext() throws Exception {
         int[] currentServers = zooKeeperClient.getSortedTabletServerList();
         List<ServerInfo> tabletServerInfos = new ArrayList<>();
         List<ServerNode> internalServerNodes = new ArrayList<>();
+
+        long start4loadTabletServer = System.currentTimeMillis();
+        Map<Integer, TabletServerRegistration> tabletServerRegistrations =
+                zooKeeperClient.getTabletServers(currentServers);
         for (int server : currentServers) {
-            TabletServerRegistration registration = zooKeeperClient.getTabletServer(server).get();
+            TabletServerRegistration registration = tabletServerRegistrations.get(server);
             ServerInfo serverInfo =
                     new ServerInfo(
                             server,
@@ -299,48 +303,74 @@ private void initCoordinatorContext() throws Exception {
         }
 
         coordinatorContext.setLiveTabletServers(tabletServerInfos);
+        LOG.info(
+                "Load tablet servers success in {}ms when initializing coordinator context.",
+                System.currentTimeMillis() - start4loadTabletServer);
+
         // init tablet server channels
         coordinatorChannelManager.startup(internalServerNodes);
 
         // load all tables
+        long start4loadTables = System.currentTimeMillis();
         List<TableInfo> autoPartitionTables = new ArrayList<>();
         List<Tuple2<TableInfo, Long>> lakeTables = new ArrayList<>();
+        Set<TablePath> tablePathSet = new HashSet<>();
         for (String database : metadataManager.listDatabases()) {
             for (String tableName : metadataManager.listTables(database)) {
-                TablePath tablePath = TablePath.of(database, tableName);
-                TableInfo tableInfo = metadataManager.getTable(tablePath);
-                coordinatorContext.putTablePath(tableInfo.getTableId(), tablePath);
-                coordinatorContext.putTableInfo(tableInfo);
-                if (tableInfo.getTableConfig().isDataLakeEnabled()) {
-                    // always set to current time,
-                    // todo: should get from the last lake snapshot
-                    lakeTables.add(Tuple2.of(tableInfo, System.currentTimeMillis()));
-                }
-                if (tableInfo.isPartitioned()) {
-                    Map<String, Long> partitions =
-                            zooKeeperClient.getPartitionNameAndIds(tablePath);
+                tablePathSet.add(TablePath.of(database, tableName));
+            }
+        }
+        Map<TablePath, TableInfo> tablePath2TableInfoMap = metadataManager.getTables(tablePathSet);
+        List<TablePath> partitionedTablePathList =
+                tablePath2TableInfoMap.entrySet().stream()
+                        .filter(entry -> entry.getValue().isPartitioned())
+                        .map(Map.Entry::getKey)
+                        .collect(Collectors.toList());
+        Map<TablePath, Map<String, Long>> tablePathMap =
+                zooKeeperClient.getPartitionNameAndIdsForTables(partitionedTablePathList);
+        for (TablePath tablePath : tablePathSet) {
+            TableInfo tableInfo = tablePath2TableInfoMap.get(tablePath);
+            coordinatorContext.putTablePath(tableInfo.getTableId(), tablePath);
+            coordinatorContext.putTableInfo(tableInfo);
+            if (tableInfo.getTableConfig().isDataLakeEnabled()) {
+                // always set to current time,
+                // todo: should get from the last lake snapshot
+                lakeTables.add(Tuple2.of(tableInfo, System.currentTimeMillis()));
+            }
+            if (tableInfo.isPartitioned()) {
+                Map<String, Long> partitions = tablePathMap.get(tablePath);
+                if (partitions != null) {
                     for (Map.Entry<String, Long> partition : partitions.entrySet()) {
                         // put partition info to coordinator context
                         coordinatorContext.putPartition(
                                 partition.getValue(),
                                 PhysicalTablePath.of(tableInfo.getTablePath(), partition.getKey()));
                     }
-                    // if the table is auto partition, put the partitions info
-                    if (tableInfo
-                            .getTableConfig()
-                            .getAutoPartitionStrategy()
-                            .isAutoPartitionEnabled()) {
-                        autoPartitionTables.add(tableInfo);
-                    }
+                }
+                // if the table is auto partition, put the partitions info
+                if (tableInfo
+                        .getTableConfig()
+                        .getAutoPartitionStrategy()
+                        .isAutoPartitionEnabled()) {
+                    autoPartitionTables.add(tableInfo);
                 }
             }
         }
+        LOG.info(
+                "Load tables success in {}ms when initializing coordinator context.",
+                System.currentTimeMillis() - start4loadTables);
+
         autoPartitionManager.initAutoPartitionTables(autoPartitionTables);
         lakeTableTieringManager.initWithLakeTables(lakeTables);
 
         // load all assignment
+        long start4loadAssignment = System.currentTimeMillis();
         loadTableAssignment();
         loadPartitionAssignment();
+        LOG.info(
+                "Load table and partition assignment success in {}ms when initializing coordinator context.",
+                System.currentTimeMillis() - start4loadAssignment);
+
         long end = System.currentTimeMillis();
         LOG.info("Current total {} tables in the cluster.", coordinatorContext.allTables().size());
         LOG.info(
@@ -355,17 +385,19 @@ private void initCoordinatorContext() throws Exception {
     private void loadTableAssignment() throws Exception {
         List<String> assignmentTables = zooKeeperClient.getChildren(TableIdsZNode.path());
         Set<Long> deletedTables = new HashSet<>();
-        for (String tableIdStr : assignmentTables) {
-            long tableId = Long.parseLong(tableIdStr);
+        List<Long> tableIds =
+                assignmentTables.stream().map(Long::parseLong).collect(Collectors.toList());
+        Map<Long, TableAssignment> tableId2tableAssignmentMap =
+                zooKeeperClient.getTablesAssignments(tableIds);
+        for (Long tableId : tableIds) {
             // if table id not in current coordinator context,
             // we'll consider it as deleted
             if (!coordinatorContext.containsTableId(tableId)) {
                 deletedTables.add(tableId);
             }
-            Optional<TableAssignment> optAssignment = zooKeeperClient.getTableAssignment(tableId);
-            if (optAssignment.isPresent()) {
-                TableAssignment tableAssignment = optAssignment.get();
-                loadAssignment(tableId, tableAssignment, null);
+            TableAssignment assignment = tableId2tableAssignmentMap.get(tableId);
+            if (assignment != null) {
+                loadAssignment(tableId, assignment, null);
             } else {
                 LOG.warn(
                         "Can't get the assignment for table {} with id {}.",
@@ -378,44 +410,51 @@ private void loadTableAssignment() throws Exception {
 
     private void loadPartitionAssignment() throws Exception {
         // load all assignment
-        List<String> partitionAssignmentNodes =
-                zooKeeperClient.getChildren(PartitionIdsZNode.path());
+        List<Long> partitionAssignmentNodes =
+                zooKeeperClient.getChildren(PartitionIdsZNode.path()).stream()
+                        .map(Long::parseLong)
+                        .collect(Collectors.toList());
         Set<TablePartition> deletedPartitions = new HashSet<>();
-        for (String partitionIdStr : partitionAssignmentNodes) {
-            long partitionId = Long.parseLong(partitionIdStr);
-            Optional<PartitionAssignment> optAssignment =
-                    zooKeeperClient.getPartitionAssignment(partitionId);
-            if (!optAssignment.isPresent()) {
+        Map<Long, PartitionAssignment> partitionId2partitionAssignmentMap =
+                zooKeeperClient.getPartitionsAssignments(partitionAssignmentNodes);
+        for (Long partitionId : partitionAssignmentNodes) {
+            PartitionAssignment assignment = partitionId2partitionAssignmentMap.get(partitionId);
+            if (assignment == null) {
                 LOG.warn("Can't get the assignment for table partition {}.", partitionId);
                 continue;
             }
-            PartitionAssignment partitionAssignment = optAssignment.get();
-            long tableId = partitionAssignment.getTableId();
+            long tableId = assignment.getTableId();
             // partition id doesn't exist in coordinator context, consider it as deleted
             if (!coordinatorContext.containsPartitionId(partitionId)) {
                 deletedPartitions.add(new TablePartition(tableId, partitionId));
             }
-            loadAssignment(tableId, optAssignment.get(), partitionId);
+            loadAssignment(tableId, assignment, partitionId);
         }
         coordinatorContext.queuePartitionDeletion(deletedPartitions);
     }
 
     private void loadAssignment(
             long tableId, TableAssignment tableAssignment, @Nullable Long partitionId)
             throws Exception {
+        Set<TableBucket> tableBucketSet = new HashSet<>();
         for (Map.Entry<Integer, BucketAssignment> entry :
                 tableAssignment.getBucketAssignments().entrySet()) {
             int bucketId = entry.getKey();
             BucketAssignment bucketAssignment = entry.getValue();
             // put the assignment information to context
             TableBucket tableBucket = new TableBucket(tableId, partitionId, bucketId);
+            tableBucketSet.add(tableBucket);
             coordinatorContext.updateBucketReplicaAssignment(
                     tableBucket, bucketAssignment.getReplicas());
-            Optional<LeaderAndIsr> optLeaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket);
+        }
+        Map<TableBucket, LeaderAndIsr> leaderAndIsrMap =
+                zooKeeperClient.getLeaderAndIsrs(tableBucketSet);
+        for (TableBucket tableBucket : tableBucketSet) {
+            LeaderAndIsr leaderAndIsr = leaderAndIsrMap.get(tableBucket);
             // update bucket LeaderAndIsr info
-            optLeaderAndIsr.ifPresent(
-                    leaderAndIsr ->
-                            coordinatorContext.putBucketLeaderAndIsr(tableBucket, leaderAndIsr));
+            if (leaderAndIsr != null) {
+                coordinatorContext.putBucketLeaderAndIsr(tableBucket, leaderAndIsr);
+            }
         }
     }
 

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

@@ -55,6 +55,8 @@
 
 import javax.annotation.Nullable;
 
+import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -316,6 +318,29 @@ public TableInfo getTable(TablePath tablePath) throws TableNotExistException {
         return tableReg.toTableInfo(tablePath, schemaInfo, defaultTableLakeOptions);
     }
 
+    public Map<TablePath, TableInfo> getTables(Collection<TablePath> tablePaths)
+            throws TableNotExistException {
+        Map<TablePath, TableInfo> result = new HashMap<>();
+        try {
+            Map<TablePath, TableRegistration> tablePath2TableRegistrations =
+                    zookeeperClient.getTables(tablePaths);
+            for (TablePath tablePath : tablePaths) {
+                if (!tablePath2TableRegistrations.containsKey(tablePath)) {
+                    throw new TableNotExistException("Table '" + tablePath + "' does not exist.");
+                }
+                TableRegistration tableReg = tablePath2TableRegistrations.get(tablePath);
+                SchemaInfo schemaInfo = getLatestSchema(tablePath);
+                result.put(
+                        tablePath,
+                        tableReg.toTableInfo(tablePath, schemaInfo, defaultTableLakeOptions));
+            }
+        } catch (Exception e) {
+            throw new FlussRuntimeException(
+                    String.format("Failed to get tables '%s'.", tablePaths), e);
+        }
+        return result;
+    }
+
     public TableRegistration getTableRegistration(TablePath tablePath) {
         Optional<TableRegistration> optionalTable;
         try {

fluss-server/src/main/java/org/apache/fluss/server/zk/ZkAsyncRequest.java

@@ -0,0 +1,48 @@
+/*
+ * 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;
+
+/** The base class for ZooKeeper async operation request. */
+public abstract class ZkAsyncRequest {
+
+    private final String path;
+
+    protected ZkAsyncRequest(String path) {
+        this.path = path;
+    }
+
+    public String getPath() {
+        return path;
+    }
+
+    // -------------------------------------------------------------------------------------------
+
+    /** The request for ZooKeeper getData async operation. */
+    public static class ZkGetDataRequest extends ZkAsyncRequest {
+        protected ZkGetDataRequest(String path) {
+            super(path);
+        }
+    }
+
+    /** The request for ZooKeeper getChildren async operation. */
+    public static class ZkGetChildrenRequest extends ZkAsyncRequest {
+        protected ZkGetChildrenRequest(String path) {
+            super(path);
+        }
+    }
+}