From 1122c04c8d2a52f4ae8ac868c99eea1ee35438bd Mon Sep 17 00:00:00 2001 From: Manish Date: Thu, 9 Jul 2026 17:50:47 +0530 Subject: [PATCH 1/2] Gate leader task assignment behind a dynamic ZooKeeper flag Add a dynamic ZooKeeper-backed kill-switch for datastream task assignment. The leader Coordinator now checks the /{cluster}/assignmentEnabled znode before performing assignment: - Absent znode -> assignment enabled (default) - Value "false" -> assignment disabled - Any other value / empty -> assignment enabled When disabled, the leader blocks in handleLeaderDoAssignment and re-checks the flag every 5s (configurable via brooklin.server.coordinator.assignmentEnabledCheckPeriodMs) until it is re-enabled or the znode is deleted. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../datastream/server/Coordinator.java | 36 +++++++++++++++ .../datastream/server/CoordinatorConfig.java | 10 +++++ .../datastream/server/zk/KeyBuilder.java | 19 ++++++++ .../datastream/server/zk/ZkAdapter.java | 26 +++++++++++ .../server/TestCoordinatorConfig.java | 13 ++++++ .../datastream/server/zk/TestZkAdapter.java | 45 +++++++++++++++++++ 6 files changed, 149 insertions(+) diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java b/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java index c0232ffd8..e1dbdfa89 100644 --- a/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java +++ b/datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java @@ -1680,6 +1680,12 @@ && hasValidDestination(datastream) && !isDeletingOrExpired(datastream)) private void handleLeaderDoAssignment(boolean isNewlyElectedLeader) { boolean succeeded = true; _log.info("START: Coordinator::handleLeaderDoAssignment."); + + // Honor the dynamic ZooKeeper assignment-enabled flag. If assignment is disabled, block here and keep + // re-checking at a fixed interval until it is re-enabled (or the flag znode is deleted), we lose leadership, + // or the coordinator is shutting down. + blockUntilAssignmentEnabled(); + List liveInstances = Collections.emptyList(); Map> previousAssignmentByInstance = Collections.emptyMap(); Map> newAssignmentsByInstance = Collections.emptyMap(); @@ -1751,6 +1757,36 @@ private void handleLeaderDoAssignment(boolean isNewlyElectedLeader) { _log.info("END: Coordinator::handleLeaderDoAssignment."); } + /** + * Block the leader Coordinator while datastream task assignment is disabled via the dynamic ZooKeeper flag + * (see {@link com.linkedin.datastream.server.zk.KeyBuilder#assignmentEnabled(String)}). + * + * The flag defaults to enabled, so this method returns immediately when the flag znode is absent or does not + * hold the value {@code "false"}. When assignment is disabled, this method re-checks the flag at a fixed + * interval ({@link CoordinatorConfig#getAssignmentEnabledCheckPeriodMs()}) and returns once assignment is + * re-enabled, the flag znode is deleted, this instance is no longer the leader, or the thread is interrupted. + */ + private void blockUntilAssignmentEnabled() { + boolean wasDisabled = false; + while (_adapter.isLeader() && !_adapter.isAssignmentEnabled()) { + wasDisabled = true; + _log.warn("Datastream task assignment is disabled via the ZooKeeper assignment-enabled flag. " + + "Skipping assignment and re-checking in {} ms.", _config.getAssignmentEnabledCheckPeriodMs()); + try { + Thread.sleep(_config.getAssignmentEnabledCheckPeriodMs()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + _log.warn("Interrupted while waiting for datastream task assignment to be re-enabled. " + + "Proceeding without further waiting."); + return; + } + } + + if (wasDisabled) { + _log.info("Datastream task assignment is enabled again. Resuming assignment."); + } + } + private void scheduleLeaderDoAssignmentRetry(boolean isNewlyElectedLeader) { _log.info("Schedule retry for leader assigning tasks"); _metrics.updateKeyedMeter(CoordinatorMetrics.KeyedMeter.HANDLE_LEADER_DO_ASSIGNMENT_NUM_RETRIES, 1); diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/CoordinatorConfig.java b/datastream-server/src/main/java/com/linkedin/datastream/server/CoordinatorConfig.java index 20c615331..18faf2e79 100644 --- a/datastream-server/src/main/java/com/linkedin/datastream/server/CoordinatorConfig.java +++ b/datastream-server/src/main/java/com/linkedin/datastream/server/CoordinatorConfig.java @@ -57,6 +57,8 @@ public final class CoordinatorConfig { // SLA threshold for stream provisioning: a datastream's INITIALIZING -> READY duration at or below this // is counted as within SLA, above it as outside SLA public static final String CONFIG_PROVISIONING_SLA_THRESHOLD_MS = PREFIX + "provisioningSlaThresholdMs"; + // how long the leader waits between checks of the ZooKeeper assignment-enabled flag while assignment is disabled + public static final String CONFIG_ASSIGNMENT_ENABLED_CHECK_PERIOD_MS = PREFIX + "assignmentEnabledCheckPeriodMs"; public static final int DEFAULT_MAX_ASSIGNMENT_RETRY_COUNT = 100; public static final long DEFAULT_STOP_PROPAGATION_TIMEOUT_MS = 60 * 1000; @@ -68,6 +70,7 @@ public final class CoordinatorConfig { public static final long DEFAULT_PROVISIONING_SLA_THRESHOLD_MS = Duration.ofMinutes(10).toMillis(); public static final long DEFAULT_THROUGHPUT_VIOLATING_TOPICS_REFRESH_PERIOD_MS = Duration.ofMinutes(5).toMillis(); public static final boolean DEFAULT_ENABLE_THROUGHPUT_VIOLATING_TOPICS_PERIODIC_REFRESH = true; + public static final long DEFAULT_ASSIGNMENT_ENABLED_CHECK_PERIOD_MS = Duration.ofSeconds(5).toMillis(); private final String _cluster; private final String _zkAddress; @@ -97,6 +100,7 @@ public final class CoordinatorConfig { private final boolean _enableThroughputViolatingTopicsPeriodicRefresh; private final double _logSizeLimitInBytes; private final long _provisioningSlaThresholdMs; + private final long _assignmentEnabledCheckPeriodMs; /** @@ -144,6 +148,8 @@ public CoordinatorConfig(Properties config) { _logSizeLimitInBytes = _properties.getDouble(CONFIG_LOG_SIZE_LIMIT_IN_BYTES, DEFAULT_LOG_SIZE_LIMIT_IN_BYTES); _provisioningSlaThresholdMs = _properties.getLong(CONFIG_PROVISIONING_SLA_THRESHOLD_MS, DEFAULT_PROVISIONING_SLA_THRESHOLD_MS); + _assignmentEnabledCheckPeriodMs = _properties.getLong(CONFIG_ASSIGNMENT_ENABLED_CHECK_PERIOD_MS, + DEFAULT_ASSIGNMENT_ENABLED_CHECK_PERIOD_MS); } public Properties getConfigProperties() { @@ -255,4 +261,8 @@ public double getLogSizeLimitInBytes() { public long getProvisioningSlaThresholdMs() { return _provisioningSlaThresholdMs; } + + public long getAssignmentEnabledCheckPeriodMs() { + return _assignmentEnabledCheckPeriodMs; + } } diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/zk/KeyBuilder.java b/datastream-server/src/main/java/com/linkedin/datastream/server/zk/KeyBuilder.java index 8c79c6a7f..ee521c652 100644 --- a/datastream-server/src/main/java/com/linkedin/datastream/server/zk/KeyBuilder.java +++ b/datastream-server/src/main/java/com/linkedin/datastream/server/zk/KeyBuilder.java @@ -26,6 +26,7 @@ public final class KeyBuilder { private static final String DATASTREAM_ASSIGNMENT_TOKEN_FOR_INSTANCE = "/%s/dms/%s/assignmentTokens/%s"; private static final String CONNECTORS = "/%s/connectors"; private static final String CONNECTOR = "/%s/connectors/%s"; + private static final String ASSIGNMENT_ENABLED = "/%s/assignmentEnabled"; // Suppresses default constructor, ensuring non-instantiability. private KeyBuilder() { @@ -251,6 +252,24 @@ public static String connectors(String cluster) { return String.format(CONNECTORS, cluster); } + /** + * Get the ZooKeeper znode that acts as a dynamic flag controlling whether the leader Coordinator is + * allowed to perform datastream task assignment. + * + * 
Example: /{cluster}/assignmentEnabled
+ * + * The flag is interpreted as follows: + *
    + *
  • Absent znode: assignment is enabled (this is the default).
    • + *
  • Znode present with value {@code "false"} (case-insensitive): assignment is disabled.
    • + *
  • Znode present with any other value: assignment is enabled.
    • + *
+ * @param cluster Brooklin cluster name + */ + public static String assignmentEnabled(String cluster) { + return String.format(ASSIGNMENT_ENABLED, cluster); + } + /** * Get the ZooKeeper znode for a specific datastream task under a connector znode * diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java b/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java index a868dc6dc..6f01aa7d5 100644 --- a/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java +++ b/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java @@ -467,6 +467,32 @@ private void joinLeaderElection() { } } + /** + * Check whether the leader Coordinator is currently allowed to perform datastream task assignment, + * based on a dynamic flag stored in ZooKeeper at {@link KeyBuilder#assignmentEnabled(String)}. + * + * The flag defaults to enabled. Assignment is considered disabled only when the znode exists and holds + * the value {@code "false"} (case-insensitive). An absent znode, an empty/null value, or any other value + * is treated as enabled. + * @return true if assignment is enabled, false if it is explicitly disabled + */ + public boolean isAssignmentEnabled() { + String path = KeyBuilder.assignmentEnabled(_cluster); + if (!_zkclient.exists(path)) { + // No znode means assignment is enabled (the default). + return true; + } + + // Use returnNullIfPathNotExists=true to gracefully handle the node being deleted between the + // exists() check and this read. + String data = _zkclient.readData(path, true); + if (data == null || data.trim().isEmpty()) { + return true; + } + + return !Boolean.FALSE.toString().equalsIgnoreCase(data.trim()); + } + /** * Update ZooKeeper znode of a datastream * @param datastream Datastream name diff --git a/datastream-server/src/test/java/com/linkedin/datastream/server/TestCoordinatorConfig.java b/datastream-server/src/test/java/com/linkedin/datastream/server/TestCoordinatorConfig.java index 4e1def8c0..722b0cb0e 100644 --- a/datastream-server/src/test/java/com/linkedin/datastream/server/TestCoordinatorConfig.java +++ b/datastream-server/src/test/java/com/linkedin/datastream/server/TestCoordinatorConfig.java @@ -110,4 +110,17 @@ public void testEnableThroughputViolatingTopicsPeriodicRefreshConfig() { CoordinatorConfig overridden = createCoordinatorConfig(props); Assert.assertFalse(overridden.getEnableThroughputViolatingTopicsPeriodicRefresh()); } + + @Test + public void testAssignmentEnabledCheckPeriodConfig() { + Properties props = new Properties(); + CoordinatorConfig config = createCoordinatorConfig(props); + Assert.assertEquals(config.getAssignmentEnabledCheckPeriodMs(), + CoordinatorConfig.DEFAULT_ASSIGNMENT_ENABLED_CHECK_PERIOD_MS); + Assert.assertEquals(config.getAssignmentEnabledCheckPeriodMs(), Duration.ofSeconds(5).toMillis()); + + props.put(CoordinatorConfig.CONFIG_ASSIGNMENT_ENABLED_CHECK_PERIOD_MS, "1000"); + CoordinatorConfig overridden = createCoordinatorConfig(props); + Assert.assertEquals(overridden.getAssignmentEnabledCheckPeriodMs(), 1000L); + } } diff --git a/datastream-server/src/test/java/com/linkedin/datastream/server/zk/TestZkAdapter.java b/datastream-server/src/test/java/com/linkedin/datastream/server/zk/TestZkAdapter.java index a6a4aacfa..d1d1bbad6 100644 --- a/datastream-server/src/test/java/com/linkedin/datastream/server/zk/TestZkAdapter.java +++ b/datastream-server/src/test/java/com/linkedin/datastream/server/zk/TestZkAdapter.java @@ -134,6 +134,51 @@ public void testSmoke() { client.close(); } + @Test + public void testIsAssignmentEnabled() { + String testCluster = "test_adapter_assignment_enabled"; + + ZkAdapter adapter = createZkAdapter(testCluster); + adapter.connect(); + + String flagPath = KeyBuilder.assignmentEnabled(testCluster); + ZkClient client = new ZkClient(_zkConnectionString); + + // No znode means assignment is enabled (the default). + Assert.assertFalse(client.exists(flagPath)); + Assert.assertTrue(adapter.isAssignmentEnabled()); + + // Explicit "false" disables assignment. + client.ensurePath(flagPath); + client.writeData(flagPath, "false"); + Assert.assertFalse(adapter.isAssignmentEnabled()); + + // Case-insensitive "false" also disables assignment. + client.writeData(flagPath, "FALSE"); + Assert.assertFalse(adapter.isAssignmentEnabled()); + + // "true" re-enables assignment. + client.writeData(flagPath, "true"); + Assert.assertTrue(adapter.isAssignmentEnabled()); + + // Any other/unrecognized value is treated as enabled. + client.writeData(flagPath, "yes"); + Assert.assertTrue(adapter.isAssignmentEnabled()); + + // Empty value is treated as enabled. + client.writeData(flagPath, ""); + Assert.assertTrue(adapter.isAssignmentEnabled()); + + // Deleting the znode restores the enabled default. + client.writeData(flagPath, "false"); + Assert.assertFalse(adapter.isAssignmentEnabled()); + client.delete(flagPath); + Assert.assertTrue(adapter.isAssignmentEnabled()); + + adapter.disconnect(); + client.close(); + } + @Test public void testLeaderElection() { String testCluster = "test_adapter_leader"; From 80c49fd0f9dab2ee089bf5df0e8a3db6b671d4f3 Mon Sep 17 00:00:00 2001 From: manishGoyalCode Date: Thu, 9 Jul 2026 19:59:49 +0530 Subject: [PATCH 2/2] Remove redundant exists() check in isAssignmentEnabled readData(path, true) already returns null for an absent znode (it catches ZkNoNodeException internally), so the preceding exists() call was an extra ZooKeeper round-trip on every assignment. Collapse to a single read. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../com/linkedin/datastream/server/zk/ZkAdapter.java | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java b/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java index 6f01aa7d5..1691763ed 100644 --- a/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java +++ b/datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java @@ -477,15 +477,9 @@ private void joinLeaderElection() { * @return true if assignment is enabled, false if it is explicitly disabled */ public boolean isAssignmentEnabled() { - String path = KeyBuilder.assignmentEnabled(_cluster); - if (!_zkclient.exists(path)) { - // No znode means assignment is enabled (the default). - return true; - } - - // Use returnNullIfPathNotExists=true to gracefully handle the node being deleted between the - // exists() check and this read. - String data = _zkclient.readData(path, true); + // Pass returnNullIfPathNotExists=true so an absent znode returns null instead of throwing; a null or + // empty value means assignment is enabled (the default). + String data = _zkclient.readData(KeyBuilder.assignmentEnabled(_cluster), true); if (data == null || data.trim().isEmpty()) { return true; }