Add configurable periodic refresh of throughput-violating topics cache

datastream-server-restli/src/test/java/com/linkedin/datastream/server/TestCoordinator.java

@@ -4041,6 +4041,90 @@ public void testThroughputViolatingTopicsHandlingForSingleDatastreamOnCreate() t
     coordinator.getDatastreamCache().getZkclient().close();
   }
 
+  @Test
+  public void testThroughputViolatingTopicsPeriodicRefreshRebuildsStaleMap() throws Exception {
+    String testCluster = "testThroughputViolatingTopicsPeriodicRefreshRebuildsStaleMap";
+    String connectorType = "connectorType";
+    String streamName = "testThroughputViolatingTopicsPeriodicRefreshRebuildsStaleMap";
+
+    Properties properties = new Properties();
+    properties.put(CoordinatorConfig.CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_HANDLING, Boolean.TRUE.toString());
+    // Short refresh period so the scheduled rebuild fires within the test window.
+    properties.put(CoordinatorConfig.CONFIG_THROUGHPUT_VIOLATING_TOPICS_REFRESH_PERIOD_MS,
+        String.valueOf(Duration.ofMillis(500).toMillis()));
+    Coordinator coordinator = createCoordinator(_zkConnectionString, testCluster, properties);
+    TestHookConnector connector1 = new TestHookConnector("connector1", connectorType);
+    coordinator.addConnector(connectorType, connector1, new BroadcastStrategy(Optional.empty()), false,
+        new SourceBasedDeduper(), null);
+    coordinator.start();
+
+    ZkClient zkClient = new ZkClient(_zkConnectionString);
+    DatastreamStore store = new ZookeeperBackedDatastreamStore(_cachedDatastreamReader, zkClient, testCluster);
+    DatastreamResources resource = new DatastreamResources(store, coordinator);
+
+    Set<String> requestedThroughputViolatingTopics = new HashSet<>(Arrays.asList("OneTopic", "TwoTopic", "ThreeTopic"));
+    Datastream testStream = DatastreamTestUtils.createDatastreams(connectorType, streamName)[0];
+    Objects.requireNonNull(testStream.getMetadata())
+        .put(DatastreamMetadataConstants.THROUGHPUT_VIOLATING_TOPICS,
+            String.join(",", requestedThroughputViolatingTopics));
+    resource.create(testStream);
+
+    // Re-fetches the cache on every poll iteration (unlike validateIfViolatingTopicsAreReflectedInServer,
+    // which snapshots once) so it can observe the async rebuild after the cache is cleared.
+    BooleanSupplier cacheMatchesRequested = () -> {
+      Set<String> fetched = coordinator.getThroughputViolatingTopics(Collections.singletonList(testStream));
+      return fetched.size() == requestedThroughputViolatingTopics.size()
+          && fetched.containsAll(requestedThroughputViolatingTopics);
+    };
+
+    // Cache is populated on the create trigger.
+    Assert.assertTrue(PollUtils.poll(cacheMatchesRequested, Duration.ofMillis(200).toMillis(),
+        Duration.ofSeconds(5).toMillis()));
+
+    // Simulate a stale/emptied cache left behind by a missed or failed rebuild. No further assignment or
+    // datastream-update event occurs, so only the periodic refresh can repopulate it.
+    coordinator.clearThroughputViolatingTopicsMapForTesting();
+
+    // The periodic refresh must rebuild the cache from the current assignment.
+    Assert.assertTrue(PollUtils.poll(cacheMatchesRequested, Duration.ofMillis(200).toMillis(),
+        Duration.ofSeconds(10).toMillis()));
+
+    coordinator.stop();
+    zkClient.close();
+    coordinator.getDatastreamCache().getZkclient().close();
+  }
+
+  @Test
+  public void testThroughputViolatingTopicsPeriodicRefreshEnablementGating() throws Exception {
+    // Handling on, periodic refresh defaults to on.
+    Properties refreshDefault = new Properties();
+    refreshDefault.put(CoordinatorConfig.CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_HANDLING, Boolean.TRUE.toString());
+    Coordinator refreshOnByDefault = createCoordinator(_zkConnectionString,
+        "testThroughputViolatingTopicsPeriodicRefreshEnablementGatingDefault", refreshDefault);
+    Assert.assertTrue(refreshOnByDefault.isThroughputViolatingTopicsPeriodicRefreshEnabled());
+    refreshOnByDefault.getDatastreamCache().getZkclient().close();
+
+    // Handling on, periodic refresh explicitly off -> disabled.
+    Properties refreshOff = new Properties();
+    refreshOff.put(CoordinatorConfig.CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_HANDLING, Boolean.TRUE.toString());
+    refreshOff.put(CoordinatorConfig.CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_PERIODIC_REFRESH,
+        Boolean.FALSE.toString());
+    Coordinator refreshDisabled = createCoordinator(_zkConnectionString,
+        "testThroughputViolatingTopicsPeriodicRefreshEnablementGatingOff", refreshOff);
+    Assert.assertFalse(refreshDisabled.isThroughputViolatingTopicsPeriodicRefreshEnabled());
+    refreshDisabled.getDatastreamCache().getZkclient().close();
+
+    // Handling off -> periodic refresh is off regardless of the toggle.
+    Properties handlingOff = new Properties();
+    handlingOff.put(CoordinatorConfig.CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_HANDLING, Boolean.FALSE.toString());
+    handlingOff.put(CoordinatorConfig.CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_PERIODIC_REFRESH,
+        Boolean.TRUE.toString());
+    Coordinator handlingDisabled = createCoordinator(_zkConnectionString,
+        "testThroughputViolatingTopicsPeriodicRefreshEnablementGatingHandlingOff", handlingOff);
+    Assert.assertFalse(handlingDisabled.isThroughputViolatingTopicsPeriodicRefreshEnabled());
+    handlingDisabled.getDatastreamCache().getZkclient().close();
+  }
+
   @Test
   public void testThroughputViolatingTopicsHandlingForMultipleDatastreams() throws Exception {
     String testCluster = "testThroughputViolatingTopicsHandlingForMultipleDatastreams";

datastream-server/src/main/java/com/linkedin/datastream/server/Coordinator.java

@@ -192,6 +192,27 @@ public class Coordinator implements ZkAdapter.ZkAdapterListener, MetricsAware {
 
   private static final AtomicLong MAX_PARTITION_COUNT = new AtomicLong(0L);
 
+  // Trigger identifiers and log messages for (re)building the host-level throughput-violating topics cache.
+  private static final String TRIGGER_CREATE = "create";
+  private static final String TRIGGER_UPDATE = "update";
+  private static final String TRIGGER_PERIODIC_REFRESH = "periodic refresh";
+  private static final String POPULATE_VIOLATING_TOPICS_MSG =
+      "Populating the datastream violating topics to host level cache from the datastream objects on the {} trigger";
+  private static final String POPULATE_VIOLATING_TOPICS_EXCEPTION_MSG =
+      "Received an exception while populating the datastream violating topics to host level cache from the datastream "
+          + "objects on the {} trigger";
+  private static final String DROP_UNRESOLVED_TASK_MSG =
+      "Dropping task {} while rebuilding throughput-violating topics map on the {} trigger; its datastream task could "
+          + "not be resolved and will be omitted from this rebuild";
+  private static final String EMPTY_REBUILD_MSG =
+      "Throughput-violating topics map rebuild on the {} trigger produced an EMPTY result from a non-empty assignment "
+          + "of {} task(s); the host-level cache will be cleared. Possible partial assignment during rebalance.";
+  private static final String PARTIAL_REBUILD_MSG =
+      "Throughput-violating topics map rebuild on the {} trigger dropped {} of {} assigned task(s); the rebuild may "
+          + "be partial.";
+  private static final String SCHEDULED_PERIODIC_REFRESH_MSG =
+      "Scheduled periodic throughput-violating topics map refresh every {} ms";
+
   private final CachedDatastreamReader _datastreamCache;
   private final Properties _eventProducerConfig;
   private final CheckpointProvider _cpProvider;
@@ -339,6 +360,18 @@ public void start() {
     // Queue up one heartbeat per period with a initial delay of 3 periods
     _scheduledExecutor.scheduleAtFixedRate(() -> _eventQueue.put(CoordinatorEvent.HEARTBEAT_EVENT),
         _heartbeatPeriod.toMillis() * 3, _heartbeatPeriod.toMillis(), TimeUnit.MILLISECONDS);
+
+    // Periodically rebuild the throughput-violating topics host-level cache so its freshness is bounded by
+    // the refresh period rather than by rebalance timing. Gated by both the feature flag and the
+    // dedicated periodic-refresh toggle (enableThroughputViolatingTopicsPeriodicRefresh). The scheduled
+    // executor lives for the coordinator's lifetime (created here, shut down in stop()), so this is scheduled
+    // once and is not re-registered in onNewSession().
+    if (isThroughputViolatingTopicsPeriodicRefreshEnabled()) {
+      long refreshPeriodMs = _config.getThroughputViolatingTopicsRefreshPeriodMs();
+      _scheduledExecutor.scheduleAtFixedRate(this::runScheduledThroughputViolatingTopicsRefresh,
+          refreshPeriodMs, refreshPeriodMs, TimeUnit.MILLISECONDS);
+      _log.info(SCHEDULED_PERIODIC_REFRESH_MSG, refreshPeriodMs);
+    }
   }
 
   protected synchronized void createEventThread() {
@@ -569,14 +602,11 @@ public void onDatastreamUpdate() {
     }
 
     if (isThroughputViolatingTopicsHandlingEnabled()) {
-      _log.info(
-          "Populating the datastream violating topics to host level cache from the datastream objects on the update trigger");
+      _log.info(POPULATE_VIOLATING_TOPICS_MSG, TRIGGER_UPDATE);
       try {
         populateThroughputViolatingTopicsMap(datastreamGroups);
       } catch (Exception exception) {
-        _log.error(
-            "Received an exception while populating the datastream violating topics to host level cache from the "
-                + "datastream objects on the update trigger", exception);
+        _log.error(POPULATE_VIOLATING_TOPICS_EXCEPTION_MSG, TRIGGER_UPDATE, exception);
       }
     }
     queueHandleAssignmentOrDatastreamChangeEvent(CoordinatorEvent.createHandleDatastreamChangeEvent(), true);
@@ -631,13 +661,32 @@ Set<String> getThroughputViolatingTopics(List<Datastream> datastreams) {
     }
   }
 
+  // Test-only hook to simulate a stale/emptied host-level cache (e.g. after a missed or failed rebuild),
+  // used to verify that the periodic refresh rebuilds it from the current assignment.
+  @VisibleForTesting
+  void clearThroughputViolatingTopicsMapForTesting() {
+    _throughputViolatingTopicsMapWriteLock.lock();
+    try {
+      _throughputViolatingTopicsMap.clear();
+    } finally {
+      _throughputViolatingTopicsMapWriteLock.unlock();
+    }
+  }
+
   // This feature enables handling the management of throughput violating topics.
   // Latency metrics and SLAs would be reported separately for these topics if their
   // per partition throughput is not within brooklin's permissible bounds.
   public boolean isThroughputViolatingTopicsHandlingEnabled() {
     return _config.getEnableThroughputViolatingTopicsHandling();
   }
 
+  // Periodic rebuild of the throughput-violating topics cache is active only when the feature is enabled
+  // AND the dedicated periodic-refresh toggle has not been turned off via config. When off, the cache is
+  // rebuilt only on assignment / datastream-update events.
+  public boolean isThroughputViolatingTopicsPeriodicRefreshEnabled() {
+    return isThroughputViolatingTopicsHandlingEnabled() && _config.getEnableThroughputViolatingTopicsPeriodicRefresh();
+  }
+
   /**
    * onPartitionMovement is called when partition movement info has been put into zookeeper
    */
@@ -774,29 +823,66 @@ private void getAssignmentsFuture(List<Future<Boolean>> assignmentChangeFutures,
   @Override
   public void onAssignmentChange() {
     _log.info("Coordinator::onAssignmentChange is called");
-    queueHandleAssignmentOrDatastreamChangeEvent(CoordinatorEvent.createHandleAssignmentChangeEvent(), true);
 
+    // Rebuild the throughput-violating topics host-level cache BEFORE queuing the async task-start event.
+    // Queuing first opens a race: a newly started producer could emit before its topic is
+    // present in the map, misrouting its events to the normal-SLA path. This mirrors onDatastreamUpdate,
+    // which also populates before queuing.
     if (isThroughputViolatingTopicsHandlingEnabled()) {
       try {
-        // On creating a datastream if the metadata contains any throughput violating topics, we populate the host level cache
-        List<DatastreamGroup> datastreamGroups = _adapter.getInstanceAssignment(_adapter.getInstanceName())
-            .stream()
-            .map(this::getDatastreamTask)
-            .filter(Objects::nonNull)
-            .map(task -> new DatastreamGroup(task.getDatastreams()))
-            .collect(Collectors.toList());
-        _log.info(
-            "Populating the datastream violating topics to host level cache from the datastream objects on the create trigger");
-        populateThroughputViolatingTopicsMap(datastreamGroups);
+        refreshThroughputViolatingTopicsMap(TRIGGER_CREATE);
       } catch (Exception exception) {
-        _log.error(
-            "Received an exception while populating the datastream violating topics to host level cache from the "
-                + "datastream objects on the create trigger", exception);
+        _log.error(POPULATE_VIOLATING_TOPICS_EXCEPTION_MSG, TRIGGER_CREATE, exception);
       }
     }
+
+    queueHandleAssignmentOrDatastreamChangeEvent(CoordinatorEvent.createHandleAssignmentChangeEvent(), true);
+
     _log.info("Coordinator::onAssignmentChange completed successfully");
   }
 
+  // Rebuilds the throughput-violating topics host-level cache from the current instance assignment.
+  // Shared by onAssignmentChange (create trigger) and the periodic refresh, so map freshness is
+  // decoupled from rebalance timing. Rebuild uses replace-all semantics via populateThroughputViolatingTopicsMap.
+  //
+  // Observability: a task that cannot be resolved (e.g. its ZNode disappeared mid-rebalance so
+  // getDatastreamTask returns null) is silently omitted from the replace-all rebuild; we WARN and count
+  // it so partial rebuilds are detectable. A non-empty assignment that yields no datastream groups
+  // would clear the cache to empty, so we WARN on that too.
+  private void refreshThroughputViolatingTopicsMap(String trigger) {
+    List<String> assignment = _adapter.getInstanceAssignment(_adapter.getInstanceName());
+    List<DatastreamGroup> datastreamGroups = new ArrayList<>();
+    int droppedTasks = 0;
+    for (String taskName : assignment) {
+      DatastreamTask task = getDatastreamTask(taskName);
+      if (task == null) {
+        droppedTasks++;
+        _log.warn(DROP_UNRESOLVED_TASK_MSG, taskName, trigger);
+        continue;
+      }
+      datastreamGroups.add(new DatastreamGroup(task.getDatastreams()));
+    }
+
+    if (!assignment.isEmpty() && datastreamGroups.isEmpty()) {
+      _log.warn(EMPTY_REBUILD_MSG, trigger, assignment.size());
+    } else if (droppedTasks > 0) {
+      _log.warn(PARTIAL_REBUILD_MSG, trigger, droppedTasks, assignment.size());
+    }
+
+    _log.info(POPULATE_VIOLATING_TOPICS_MSG, trigger);
+    populateThroughputViolatingTopicsMap(datastreamGroups);
+  }
+
+  // Wraps refreshThroughputViolatingTopicsMap for the scheduled executor: any thrown exception must be
+  // swallowed here, otherwise scheduleAtFixedRate would silently cancel all future refreshes.
+  private void runScheduledThroughputViolatingTopicsRefresh() {
+    try {
+      refreshThroughputViolatingTopicsMap(TRIGGER_PERIODIC_REFRESH);
+    } catch (Exception exception) {
+      _log.error(POPULATE_VIOLATING_TOPICS_EXCEPTION_MSG, TRIGGER_PERIODIC_REFRESH, exception);
+    }
+  }
+
   private int getAssignmentTaskCount(Map<String, List<DatastreamTask>> assignment) {
     return assignment.values().stream().mapToInt(List::size).sum();
   }

datastream-server/src/main/java/com/linkedin/datastream/server/CoordinatorConfig.java

@@ -47,6 +47,12 @@ public final class CoordinatorConfig {
   public static final String CONFIG_MARK_DATASTREAMS_STOPPED_RETRY_PERIOD_MS = PREFIX + "markDatastreamsStoppedRetryPeriodMs";
 
   public static final String CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_HANDLING = PREFIX + "enableThroughputViolatingTopicsHandling";
+  // how often the host-level throughput-violating topics cache is rebuilt from the current assignment,
+  // independent of rebalance timing. Bounds any staleness window to this period.
+  public static final String CONFIG_THROUGHPUT_VIOLATING_TOPICS_REFRESH_PERIOD_MS = PREFIX + "throughputViolatingTopicsRefreshPeriodMs";
+  // whether the periodic rebuild of the host-level throughput-violating topics cache is enabled. When
+  // disabled, the cache is only rebuilt on assignment / datastream-update events.
+  public static final String CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_PERIODIC_REFRESH = PREFIX + "enableThroughputViolatingTopicsPeriodicRefresh";
   public static final String CONFIG_LOG_SIZE_LIMIT_IN_BYTES = PREFIX + "logSizeLimitInBytes";
   // 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
@@ -60,6 +66,8 @@ public final class CoordinatorConfig {
   public static final int DEFAULT_MARK_DATASTREMS_STOPPED_RETRY_PERIOD_MS = 10 * 1000;
   public static final int DEFAULT_LOG_SIZE_LIMIT_IN_BYTES = 1024 * 1024;
   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;
 
   private final String _cluster;
   private final String _zkAddress;
@@ -85,6 +93,8 @@ public final class CoordinatorConfig {
   private final long _markDatastreamsStoppedTimeoutMs;
   private final long _markDatastreamsStoppedRetryPeriodMs;
   private final boolean _enableThroughputViolatingTopicsHandling;
+  private final long _throughputViolatingTopicsRefreshPeriodMs;
+  private final boolean _enableThroughputViolatingTopicsPeriodicRefresh;
   private final double _logSizeLimitInBytes;
   private final long _provisioningSlaThresholdMs;
 
@@ -125,6 +135,12 @@ public CoordinatorConfig(Properties config) {
         DEFAULT_MARK_DATASTREMS_STOPPED_RETRY_PERIOD_MS);
     _enableThroughputViolatingTopicsHandling = _properties.getBoolean(
         CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_HANDLING, false);
+    _throughputViolatingTopicsRefreshPeriodMs = _properties.getLong(
+        CONFIG_THROUGHPUT_VIOLATING_TOPICS_REFRESH_PERIOD_MS,
+        DEFAULT_THROUGHPUT_VIOLATING_TOPICS_REFRESH_PERIOD_MS);
+    _enableThroughputViolatingTopicsPeriodicRefresh = _properties.getBoolean(
+        CONFIG_ENABLE_THROUGHPUT_VIOLATING_TOPICS_PERIODIC_REFRESH,
+        DEFAULT_ENABLE_THROUGHPUT_VIOLATING_TOPICS_PERIODIC_REFRESH);
     _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);
@@ -194,6 +210,14 @@ public boolean getEnableThroughputViolatingTopicsHandling() {
     return _enableThroughputViolatingTopicsHandling;
   }
 
+  public long getThroughputViolatingTopicsRefreshPeriodMs() {
+    return _throughputViolatingTopicsRefreshPeriodMs;
+  }
+
+  public boolean getEnableThroughputViolatingTopicsPeriodicRefresh() {
+    return _enableThroughputViolatingTopicsPeriodicRefresh;
+  }
+
   // Configuration properties for Assignment Tokens Feature
 
   public boolean getEnableAssignmentTokens() {