Replacing Coordinator Queue With Deque & Fixing Usage Of toMap Util

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

@@ -0,0 +1,20 @@
+/**
+ *  Copyright 2023 LinkedIn Corporation. All rights reserved.
+ *  Licensed under the BSD 2-Clause License. See the LICENSE file in the project root for license information.
+ *  See the NOTICE file in the project root for additional information regarding copyright ownership.
+ */
+package com.linkedin.datastream.server;
+
+import java.util.Properties;
+
+/**
+ * Callable Coordinator is used for overriding coordinator behaviors for tests
+ */
+public interface CallableCoordinatorForTest {
+  /**
+   * invoking constructor of coordinator with params,
+   * - datastreamCache to maintain all the datastreams in the cluster.
+   * - properties to use while creating coordinator.
+   * */
+  Coordinator invoke(CachedDatastreamReader cachedDatastreamReader, Properties properties);
+}

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

@@ -9,6 +9,7 @@
 import java.lang.reflect.Method;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -21,6 +22,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
+import java.util.Queue;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
@@ -136,7 +138,12 @@ private Coordinator createCoordinator(String zkAddr, String cluster) throws Exce
   }
 
   private Coordinator createCoordinator(String zkAddr, String cluster, Properties override) throws Exception {
-    return createCoordinator(zkAddr, cluster, override, new DummyTransportProviderAdminFactory());
+    return createCoordinator(zkAddr, cluster, override, new DummyTransportProviderAdminFactory(), Coordinator::new);
+  }
+
+  private Coordinator createCoordinator(String zkAddr, String cluster, Properties override,
+      TransportProviderAdminFactory transportProviderAdminFactory) throws Exception {
+    return createCoordinator(zkAddr, cluster, override, transportProviderAdminFactory, Coordinator::new);
   }
 
   private Coordinator createCoordinator(String zkAddr, String cluster, Properties override,
@@ -163,7 +170,7 @@ protected synchronized void handleEvent(CoordinatorEvent event) {
   }
 
   private Coordinator createCoordinator(String zkAddr, String cluster, Properties override,
-      TransportProviderAdminFactory transportProviderAdminFactory) throws Exception {
+      TransportProviderAdminFactory transportProviderAdminFactory, CallableCoordinatorForTest callableCoordinatorForTest) throws Exception {
     Properties props = new Properties();
     props.put(CoordinatorConfig.CONFIG_CLUSTER, cluster);
     props.put(CoordinatorConfig.CONFIG_ZK_ADDRESS, zkAddr);
@@ -172,7 +179,7 @@ private Coordinator createCoordinator(String zkAddr, String cluster, Properties
     props.putAll(override);
     ZkClient client = new ZkClient(zkAddr);
     _cachedDatastreamReader = new CachedDatastreamReader(client, cluster);
-    Coordinator coordinator = new Coordinator(_cachedDatastreamReader, props);
+    Coordinator coordinator = callableCoordinatorForTest.invoke(_cachedDatastreamReader, props);
     coordinator.addTransportProvider(DummyTransportProviderAdminFactory.PROVIDER_NAME,
         transportProviderAdminFactory.createTransportProviderAdmin(DummyTransportProviderAdminFactory.PROVIDER_NAME,
             new Properties()));
@@ -3945,6 +3952,60 @@ public void testThroughputViolatingTopicsHandlingForSingleDatastreamOnCreateWith
     coordinator.getDatastreamCache().getZkclient().close();
   }
 
+  @Test
+  public void testLeaderDoAssignmentForNewlyElectedLeaderFailurePath() throws Exception {
+    String testCluster = "testLeaderDoAssignmentForNewlyElectedLeaderFailurePath";
+    String connectorType = "connectorType";
+    String streamName = "testLeaderDoAssignmentForNewlyElectedLeaderFailurePath";
+
+    Queue<CoordinatorEvent> shadowCoordinatorQueue = new ArrayDeque<>();
+    Properties properties = new Properties();
+    Coordinator coordinator =
+        createCoordinator(_zkConnectionString, testCluster, properties, new DummyTransportProviderAdminFactory(),
+            (cachedDatastreamReader, props) -> new Coordinator(cachedDatastreamReader, props) {
+
+              // This override generates an exception while the newly elected leader performs pre assignment cleanup.
+              // The exception causes the handleLeaderDoAssignment handler to exit, along with inserting the same event
+              // in the queue for a reattempt.
+              @Override
+              protected void performPreAssignmentCleanup(List<DatastreamGroup> datastreamGroups) {
+                throw new RuntimeException("testing exception path in assignment cleanup routine");
+              }
+
+              // This override collects the coordinator queue events in a shadow queue for test purposes.
+              @Override
+              protected synchronized void handleEvent(CoordinatorEvent event) {
+                shadowCoordinatorQueue.add(event);
+                super.handleEvent(event);
+              }
+            });
+    TestHookConnector dummyConnector = new TestHookConnector("dummyConnector", connectorType);
+    coordinator.addConnector(connectorType, dummyConnector, new BroadcastStrategy(Optional.empty()), false,
+        new SourceBasedDeduper(), null);
+    coordinator.start();
+
+    ZkClient zkClient = new ZkClient(_zkConnectionString);
+
+    Datastream testDatastream =
+        DatastreamTestUtils.createAndStoreDatastreams(zkClient, testCluster, connectorType, streamName)[0];
+
+    coordinator.stop();
+    zkClient.close();
+    coordinator.getDatastreamCache().getZkclient().close();
+
+    // This is the event which should be added to the front of the queue once the handler exits on an exception.
+    CoordinatorEvent leaderDoAssignmentForNewlyElectedLeader =
+        new CoordinatorEvent(CoordinatorEvent.EventType.LEADER_DO_ASSIGNMENT, true);
+
+    // while-ing until the newly elected leader performs the handlerLeaderDoAssignment request for the first time.
+    while (!Objects.equals(shadowCoordinatorQueue.peek(), leaderDoAssignmentForNewlyElectedLeader)) {
+      shadowCoordinatorQueue.poll();
+    }
+
+    // As we expect the reattempt event to added to the front of the queue, the front of the queue should be the same.
+    Assert.assertEquals(shadowCoordinatorQueue.poll(), leaderDoAssignmentForNewlyElectedLeader);
+  }
+
   // This helper function helps compare the requesting topics with the topics reflected in the server.
   private BooleanSupplier validateIfViolatingTopicsAreReflectedInServer(Datastream testStream, Coordinator coordinator,
       Set<String> requestedThroughputViolatingTopics) {

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

@@ -834,7 +834,8 @@ private void handleAssignmentChange(boolean isDatastreamUpdate) throws TimeoutEx
     _assignedDatastreamTasks.putAll(currentAssignment.values()
         .stream()
         .flatMap(Collection::stream)
-        .collect(Collectors.toMap(DatastreamTask::getDatastreamTaskName, Function.identity())));
+        .collect(Collectors.toMap(DatastreamTask::getDatastreamTaskName, Function.identity(),
+            (existingTask, duplicateTask) -> existingTask)));
     List<DatastreamTask> newAssignment = new ArrayList<>(_assignedDatastreamTasks.values());
 
     if ((totalTasks - submittedTasks) > 0) {
@@ -1524,10 +1525,11 @@ private void scheduleLeaderDoAssignmentRetry(boolean isNewlyElectedLeader) {
     _log.info("Schedule retry for leader assigning tasks");
     _metrics.updateKeyedMeter(CoordinatorMetrics.KeyedMeter.HANDLE_LEADER_DO_ASSIGNMENT_NUM_RETRIES, 1);
     _leaderDoAssignmentScheduled.set(true);
+    // scheduling LEADER_DO_ASSIGNMENT event instantly to prevent any other event being handled before the reattempt.
     _leaderDoAssignmentScheduledFuture = _scheduledExecutor.schedule(() -> {
-      _eventQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(isNewlyElectedLeader));
+      _eventQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(isNewlyElectedLeader), false);
       _leaderDoAssignmentScheduled.set(false);
-    }, _config.getRetryIntervalMs(), TimeUnit.MILLISECONDS);
+    }, 0, TimeUnit.MILLISECONDS);
   }
 
   @VisibleForTesting
@@ -1614,7 +1616,7 @@ private void revokeUnclaimedAssignmentTokens(Map<String, List<AssignmentToken>>
     }
   }
 
-  private void performPreAssignmentCleanup(List<DatastreamGroup> datastreamGroups) {
+  protected void performPreAssignmentCleanup(List<DatastreamGroup> datastreamGroups) {
 
     // Map between instance to tasks assigned to the instance.
     Map<String, Set<DatastreamTask>> previousAssignmentByInstance = _adapter.getAllAssignedDatastreamTasks();

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

@@ -52,7 +52,7 @@ private CoordinatorEvent(EventType eventType) {
     _eventMetadata = null;
   }
 
-  private CoordinatorEvent(EventType eventType, Object eventMetadata) {
+  protected CoordinatorEvent(EventType eventType, Object eventMetadata) {
     _eventType = eventType;
     _eventMetadata = eventMetadata;
   }

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

@@ -7,12 +7,12 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Deque;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.LinkedBlockingDeque;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -44,7 +44,7 @@ class CoordinatorEventBlockingQueue implements MetricsAware {
   static final String GAUGE_KEY = "queuedEvents";
 
   private final Set<CoordinatorEvent> _eventSet;
-  private final Queue<CoordinatorEvent> _eventQueue;
+  private final Deque<CoordinatorEvent> _eventQueue;
   private final DynamicMetricsManager _dynamicMetricsManager;
   private final Gauge<Integer> _gauge;
   private final Counter _counter;
@@ -59,7 +59,7 @@ class CoordinatorEventBlockingQueue implements MetricsAware {
    */
   CoordinatorEventBlockingQueue(String key) {
     _eventSet = new HashSet<>();
-    _eventQueue = new LinkedBlockingQueue<>();
+    _eventQueue = new LinkedBlockingDeque<>();
     _dynamicMetricsManager = DynamicMetricsManager.getInstance();
 
     String prefix = buildMetricName(key);
@@ -73,16 +73,30 @@ class CoordinatorEventBlockingQueue implements MetricsAware {
 
 
   /**
-   * Add a single event to the queue, overwriting events with the same name and same metadata.
+   * Add a single event to the queue. Defaults to adding the event at the end of the queue.
    * @param event CoordinatorEvent event to add to the queue
    */
   public synchronized void put(CoordinatorEvent event) {
-    LOG.info("Queuing event {} to event queue", event.getType());
+    put(event, true);
+  }
+
+  /**
+   * Add a single event to the queue, de-duping events with the same name and same metadata.
+   * @param event CoordinatorEvent event to add to the queue
+   * @param insertInTheEnd if true, indicates to add the event to the end of the queue and front, otherwise.
+   */
+  public synchronized void put(CoordinatorEvent event, boolean insertInTheEnd) {
+    LOG.info("Queuing event {} at the " + (insertInTheEnd ? "end" : "front") + " of the event queue", event.getType());
     if (_eventSet.contains(event)) {
       _counter.inc(); // count duplicate event
     } else {
       // only insert if there isn't an event present in the queue with the same name and same metadata.
-      boolean result = _eventQueue.offer(event);
+      boolean result;
+      if (insertInTheEnd) {
+        result = _eventQueue.offer(event);
+      } else {
+        result = _eventQueue.offerFirst(event);
+      }
       if (!result) {
         return;
       }

datastream-server/src/main/java/com/linkedin/datastream/server/assignment/StickyPartitionAssignmentStrategy.java

@@ -487,7 +487,8 @@ public Map<String, List<DatastreamTask>> getTasksToCleanUp(List<DatastreamGroup>
     Map<String, DatastreamTask> assignmentsMap = currentAssignment.values()
         .stream()
         .flatMap(Collection::stream)
-        .collect(Collectors.toMap(DatastreamTask::getDatastreamTaskName, Function.identity()));
+        .collect(Collectors.toMap(DatastreamTask::getDatastreamTaskName, Function.identity(),
+            (existingTask, duplicateTask) -> existingTask));
 
     for (String instance : currentAssignment.keySet()) {
       // find the dependency tasks which also exist in the assignmentsMap.

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

@@ -759,10 +759,11 @@ public void updateAllAssignmentsAndIssueTokens(Map<String, List<DatastreamTask>>
   private Map<DatastreamGroup, Set<String>> getStoppingDatastreamGroupInstances(
       List<DatastreamGroup> stoppingDatastreamGroups) {
     Map<String, Set<DatastreamTask>> currentAssignment = getAllAssignedDatastreamTasks();
-    Set<String> stoppingDatastreamTaskPrefixes = stoppingDatastreamGroups.stream().
-        map(DatastreamGroup::getTaskPrefix).collect(toSet());
-    Map<String, DatastreamGroup> taskPrefixDatastreamGroups = stoppingDatastreamGroups.stream().
-        collect(Collectors.toMap(DatastreamGroup::getTaskPrefix, Function.identity()));
+    Set<String> stoppingDatastreamTaskPrefixes =
+        stoppingDatastreamGroups.stream().map(DatastreamGroup::getTaskPrefix).collect(toSet());
+    Map<String, DatastreamGroup> taskPrefixDatastreamGroups = stoppingDatastreamGroups.stream()
+        .collect(Collectors.toMap(DatastreamGroup::getTaskPrefix, Function.identity(),
+            (existingDatastreamGroup, duplicateDatastreamGroup) -> existingDatastreamGroup));
 
     Map<DatastreamGroup, Set<String>> stoppingDgInstances = new HashMap<>();
     currentAssignment.keySet()

datastream-server/src/test/java/com/linkedin/datastream/server/TestCoordinatorEventBlockingQueue.java

@@ -42,20 +42,20 @@ public void resetMetrics() {
   public void testHappyPath() throws Exception {
     CoordinatorEventBlockingQueue eventBlockingQueue = new CoordinatorEventBlockingQueue(SIMPLE_NAME);
     eventBlockingQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(false));
-    eventBlockingQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(true));
-    eventBlockingQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(false));
-    eventBlockingQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(true));
+    eventBlockingQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(true), false);
+    eventBlockingQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(false), true);
+    eventBlockingQueue.put(CoordinatorEvent.createLeaderDoAssignmentEvent(true), false);
     eventBlockingQueue.put(CoordinatorEvent.createLeaderPartitionAssignmentEvent("test1"));
     eventBlockingQueue.put(CoordinatorEvent.createLeaderPartitionAssignmentEvent("test1"));
-    eventBlockingQueue.put(CoordinatorEvent.createLeaderPartitionAssignmentEvent("test2"));
+    eventBlockingQueue.put(CoordinatorEvent.createLeaderPartitionAssignmentEvent("test2"), false);
     eventBlockingQueue.put(CoordinatorEvent.HANDLE_ASSIGNMENT_CHANGE_EVENT);
     eventBlockingQueue.put(CoordinatorEvent.HANDLE_ASSIGNMENT_CHANGE_EVENT);
     eventBlockingQueue.put(CoordinatorEvent.HANDLE_ASSIGNMENT_CHANGE_EVENT);
     Assert.assertEquals(eventBlockingQueue.size(), 5);
-    Assert.assertEquals(eventBlockingQueue.take(), CoordinatorEvent.createLeaderDoAssignmentEvent(false));
+    Assert.assertEquals(eventBlockingQueue.take(), CoordinatorEvent.createLeaderPartitionAssignmentEvent("test2"));
     Assert.assertEquals(eventBlockingQueue.take(), CoordinatorEvent.createLeaderDoAssignmentEvent(true));
+    Assert.assertEquals(eventBlockingQueue.take(), CoordinatorEvent.createLeaderDoAssignmentEvent(false));
     Assert.assertEquals(eventBlockingQueue.take(), CoordinatorEvent.createLeaderPartitionAssignmentEvent("test1"));
-    Assert.assertEquals(eventBlockingQueue.take(), CoordinatorEvent.createLeaderPartitionAssignmentEvent("test2"));
     Assert.assertEquals(eventBlockingQueue.take(), CoordinatorEvent.HANDLE_ASSIGNMENT_CHANGE_EVENT);
   }
 

datastream-server/src/test/java/com/linkedin/datastream/server/assignment/TestStickyPartitionAssignmentStrategy.java

@@ -758,6 +758,39 @@ public void testExpectedNumberOfTasks() {
     Assert.assertEquals(-1, getNumTasksForDatastreamFromZK(ds.get(0).getName()));
   }
 
+  @Test
+  public void testTasksCleanUpWithDuplicatesAcrossInstances() {
+    StickyPartitionAssignmentStrategy strategy =
+        createStickyPartitionAssignmentStrategy(3, 90, true, getZkClient(true), _clusterName);
+    List<DatastreamGroup> datastreams = generateDatastreams("testTasksCleanUpWithDuplicatesAcrossInstances", 1, 3);
+
+    Map<String, Set<DatastreamTask>> assignment = generateEmptyAssignment(datastreams, 2, 3, true);
+
+    List<String> newPartitions = ImmutableList.of("t-0", "t-1", "tt-0", "tt-1", "ttt-0", "ttt-1", "ttt-2");
+    assignment =
+        strategy.assignPartitions(assignment, new DatastreamGroupPartitionsMetadata(datastreams.get(0), newPartitions));
+
+    // This following snippet demonstrates a previous leader performing some task movements but got interrupted, OOMed
+    // or hit session expiry.
+    // The previous leader was able to add some tasks to the newer instance's assignment, but couldn't remove them from
+    // previous instance's assignment.
+    // The next leader should be able to identify and cleanup, even though there'll be duplicate tasks across instances.
+    DatastreamTask previousTask = null;
+    for (String instance : assignment.keySet()) {
+      if (previousTask != null) {
+        assignment.get(instance).add(previousTask);
+      }
+      previousTask = assignment.get(instance).iterator().next();
+    }
+
+    try {
+      Map<String, List<DatastreamTask>> taskToCleanup = strategy.getTasksToCleanUp(datastreams, assignment);
+      Assert.assertEquals(taskToCleanup.size(), 0);
+    } catch (Exception exception) {
+      Assert.fail("Received exception while finding tasks to be cleaned up", exception.getCause());
+    }
+  }
+
   private int getNumTasksForDatastreamFromZK(String taskPrefix) {
     String numTasksPath = KeyBuilder.datastreamNumTasks(_clusterName, taskPrefix);
     if (!_zkClient.exists(numTasksPath)) {