From 8a511d27651ec39e14fee637c7bd30c44d879d6d Mon Sep 17 00:00:00 2001 From: zipg Date: Fri, 10 Jul 2026 00:32:04 +0800 Subject: [PATCH] fix(zookeeper): optimize node list loading --- .../dbx/agent/zookeeper/ZooKeeperAgent.java | 106 ++++++++++++++++-- .../agent/zookeeper/ZooKeeperAgentTest.java | 39 +++++++ 2 files changed, 135 insertions(+), 10 deletions(-) diff --git a/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java b/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java index 309960944c..72b4329266 100644 --- a/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java +++ b/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java @@ -17,7 +17,13 @@ import java.nio.charset.CharsetDecoder; import java.nio.charset.CodingErrorAction; import java.nio.charset.StandardCharsets; import java.util.*; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; public final class ZooKeeperAgent { private static final Gson GSON = new GsonBuilder().serializeNulls().create(); @@ -26,6 +32,16 @@ public final class ZooKeeperAgent { private static final int DEFAULT_CONNECTION_TIMEOUT_MS = 15000; private static final int DEFAULT_BASE_SLEEP_TIME_MS = 1000; private static final int DEFAULT_MAX_RETRIES = 3; + private static final String STAT_LOOKUP_CONCURRENCY_PROPERTY = "dbx.zookeeper.statLookupConcurrency"; + private static final String STAT_LOOKUP_CONCURRENCY_ENV = "DBX_ZOOKEEPER_STAT_LOOKUP_CONCURRENCY"; + private static final int DEFAULT_STAT_LOOKUP_CONCURRENCY = 16; + private static final int MIN_STAT_LOOKUP_CONCURRENCY = 1; + private static final int MAX_STAT_LOOKUP_CONCURRENCY = 64; + private static final int STAT_LOOKUP_CONCURRENCY = configuredStatLookupConcurrency(); + private static final ExecutorService STAT_LOOKUP_EXECUTOR = Executors.newFixedThreadPool( + STAT_LOOKUP_CONCURRENCY, + daemonThreadFactory("dbx-zookeeper-stat-lookup-") + ); private static final List CAPABILITIES = Collections.unmodifiableList(Arrays.asList( AgentProtocol.CAPABILITY_CONNECT, AgentProtocol.CAPABILITY_TEST_CONNECTION, @@ -143,16 +159,7 @@ public final class ZooKeeperAgent { Collections.sort(paths); int offset = cursor == null ? 0 : Math.max(0, cursor.offset); int end = Math.min(paths.size(), offset + limit); - for (int index = offset; index < end; index++) { - Stat stat = active.checkExists().forPath(paths.get(index)); - if (stat == null) { - continue; - } - Map row = new LinkedHashMap<>(); - row.put("key", paths.get(index)); - row.putAll(statMetadata(stat)); - keys.add(row); - } + keys.addAll(listRowsWithMetadata(active, paths.subList(offset, end))); result.put("keys", keys); result.put("continuation", end < paths.size() ? encodeCursor(new Cursor(root, recursive, end)) : null); @@ -278,6 +285,7 @@ public final class ZooKeeperAgent { } case AgentProtocol.METHOD_SHUTDOWN -> { closeClient(); + STAT_LOOKUP_EXECUTOR.shutdownNow(); System.exit(0); yield Collections.singletonMap("ok", true); } @@ -449,6 +457,55 @@ public final class ZooKeeperAgent { return result; } + private static List> listRowsWithMetadata(CuratorFramework active, List paths) throws Exception { + if (paths.isEmpty()) { + return Collections.emptyList(); + } + + try { + List>> futures = new ArrayList<>(); + for (String path : paths) { + futures.add(STAT_LOOKUP_EXECUTOR.submit(() -> rowWithMetadata(active, path))); + } + + List> rows = new ArrayList<>(); + for (Future> future : futures) { + Map row = future.get(); + if (row != null) { + rows.add(row); + } + } + return rows; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw e; + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception exception) { + throw exception; + } + if (cause instanceof Error error) { + throw error; + } + throw new RuntimeException(cause); + } + } + + private static Map rowWithMetadata(CuratorFramework active, String path) throws Exception { + try { + Stat stat = active.checkExists().forPath(path); + if (stat == null) { + return null; + } + Map row = new LinkedHashMap<>(); + row.put("key", path); + row.putAll(statMetadata(stat)); + return row; + } catch (KeeperException.NoNodeException e) { + return null; + } + } + private static void collectRecursive(CuratorFramework active, String path, List result) throws Exception { List children; try { @@ -525,6 +582,35 @@ public final class ZooKeeperAgent { return element == null || element.isJsonNull() ? fallback : element.getAsBoolean(); } + static int configuredStatLookupConcurrency(String propertyValue, String envValue) { + String configured = firstNonBlank(propertyValue, envValue); + if (configured == null) { + return DEFAULT_STAT_LOOKUP_CONCURRENCY; + } + try { + int parsed = Integer.parseInt(configured.trim()); + return Math.max(MIN_STAT_LOOKUP_CONCURRENCY, Math.min(MAX_STAT_LOOKUP_CONCURRENCY, parsed)); + } catch (NumberFormatException e) { + return DEFAULT_STAT_LOOKUP_CONCURRENCY; + } + } + + private static int configuredStatLookupConcurrency() { + return configuredStatLookupConcurrency( + System.getProperty(STAT_LOOKUP_CONCURRENCY_PROPERTY), + System.getenv(STAT_LOOKUP_CONCURRENCY_ENV) + ); + } + + private static ThreadFactory daemonThreadFactory(String prefix) { + AtomicInteger sequence = new AtomicInteger(1); + return runnable -> { + Thread thread = new Thread(runnable, prefix + sequence.getAndIncrement()); + thread.setDaemon(true); + return thread; + }; + } + private static String firstNonBlank(String... values) { for (String value : values) { if (value != null && !value.isBlank()) { diff --git a/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java b/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java index 9e30769e57..c50f3518a7 100644 --- a/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java +++ b/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java @@ -103,6 +103,16 @@ final class ZooKeeperAgentTest { Assertions.assertEquals("/app/name", ZooKeeperAgent.normalizePath("/app/name")); } + @Test + void statLookupConcurrencyUsesConservativeDefaultsAndOverrides() { + Assertions.assertEquals(16, ZooKeeperAgent.configuredStatLookupConcurrency(null, null)); + Assertions.assertEquals(20, ZooKeeperAgent.configuredStatLookupConcurrency("20", "12")); + Assertions.assertEquals(12, ZooKeeperAgent.configuredStatLookupConcurrency(null, "12")); + Assertions.assertEquals(1, ZooKeeperAgent.configuredStatLookupConcurrency("0", null)); + Assertions.assertEquals(64, ZooKeeperAgent.configuredStatLookupConcurrency("128", null)); + Assertions.assertEquals(16, ZooKeeperAgent.configuredStatLookupConcurrency("not-a-number", null)); + } + @Test void connectAndTestConnectionWorkAgainstTestingServer() throws Exception { try (TestingServer server = new TestingServer()) { @@ -403,6 +413,23 @@ final class ZooKeeperAgentTest { } } + @Test + void listPrefixKeepsDirectChildMetadataForLazyExpansion() throws Exception { + try (TestingServer server = new TestingServer()) { + connect(server); + result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); + result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); + + JsonObject list = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\",\"recursive\":false}")); + JsonObject app = listedRow(list, "/app"); + JsonObject config = listedRow(list, "/config"); + + Assertions.assertEquals(1, app.get("numChildren").getAsInt()); + Assertions.assertEquals(0, config.get("numChildren").getAsInt()); + Assertions.assertEquals(3, config.get("dataLength").getAsInt()); + } + } + @Test void listPrefixDefaultsToRecursiveForDbxKvBrowser() throws Exception { try (TestingServer server = new TestingServer()) { @@ -554,4 +581,16 @@ final class ZooKeeperAgentTest { } return keys; } + + private static JsonObject listedRow(JsonObject listResult, String key) { + JsonArray rows = listResult.getAsJsonArray("keys"); + for (JsonElement row : rows) { + JsonObject object = row.getAsJsonObject(); + if (key.equals(object.get("key").getAsString())) { + return object; + } + } + Assertions.fail("expected listed row for " + key); + return new JsonObject(); + } }