BMM Restart Improvements Part 2. Followers claim their tokens after handling stop assignments

datastream-server-api/src/main/java/com/linkedin/datastream/server/api/connector/Connector.java

@@ -72,7 +72,7 @@ public interface Connector extends MetricsAware, DatastreamChangeListener {
    * @throws UnsupportedOperationException if not implemented by Connector classes.
    */
   default List<String> getActiveTasks() {
-    throw new UnsupportedOperationException("Running tasks API is not supported unless implemented by coordinators");
+    throw new UnsupportedOperationException("Running tasks API is not supported unless implemented by connectors");
   }
 
   /**

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

@@ -716,11 +716,11 @@ private void handleAssignmentChange(boolean isDatastreamUpdate) throws TimeoutEx
 
     if (_config.getEnableAssignmentTokens()) {
       try {
+        ExecutorService executor = Executors.newSingleThreadExecutor();
         // Queue assignment token claim task
-        _executor.schedule(() -> maybeClaimAssignmentTokensForStoppingStreams(newAssignment, oldAssignment),
-            0, TimeUnit.MILLISECONDS);
+        executor.submit(() -> maybeClaimAssignmentTokensForStoppingStreams(newAssignment, oldAssignment));
       } catch (RejectedExecutionException ex) {
-        _log.warn("Failed to schedule the task for claiming assignment tokens", ex);
+        _log.warn("Failed to submit the task for claiming assignment tokens", ex);
       }
     }
 
@@ -828,6 +828,7 @@ void maybeClaimAssignmentTokensForStoppingStreams(List<DatastreamTask> newAssign
             filter(t -> stoppingStreamNames.contains(t.getTaskPrefix())).
             map(DatastreamTask::getId).collect(Collectors.toSet());
 
+        // TODO Evaluate whether we need to optimize here and make this call for each datastream
         if (PollUtils.poll(() -> connectorTasksHaveStopped(connector, stoppingDatastreamTasks),
             _config.getTaskStopCheckRetryPeriodMs(), _config.getTaskStopCheckTimeoutMs())) {
           _adapter.claimAssignmentTokensForDatastreams(stoppingStreams, _adapter.getInstanceName());