Gate leader task assignment behind a dynamic ZooKeeper flag

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<String> liveInstances = Collections.emptyList();
     Map<String, Set<DatastreamTask>> previousAssignmentByInstance = Collections.emptyMap();
     Map<String, List<DatastreamTask>> 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);

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;
+  }
 }

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.
+   *
+   * <pre>Example: /{cluster}/assignmentEnabled</pre>
+   *
+   * The flag is interpreted as follows:
+   * <ul>
+   *   <li>Absent znode: assignment is enabled (this is the default).</li>
+   *   <li>Znode present with value {@code "false"} (case-insensitive): assignment is disabled.</li>
+   *   <li>Znode present with any other value: assignment is enabled.</li>
+   * </ul>
+   * @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
    *

datastream-server/src/main/java/com/linkedin/datastream/server/zk/ZkAdapter.java

@@ -467,6 +467,26 @@ 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() {
+    // 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;
+    }
+
+    return !Boolean.FALSE.toString().equalsIgnoreCase(data.trim());
+  }
+
   /**
    * Update ZooKeeper znode of a datastream
    * @param datastream Datastream name

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);
+  }
 }

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";