Fixing Deadlock Bug Between ZK Callback & Event Thread On Acquiring Coordinator Object

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

@@ -4077,6 +4077,163 @@ protected synchronized void handleEvent(CoordinatorEvent event) {
     });
   }
 
+  @Test
+  public void testSessionExpiryCallbackThreadAttemptingToAcquireCoordinatorObjectBeforeHandlingEvent() throws Exception {
+    String testCluster = "dummyCluster";
+    long testHeartbeatPeriod = Duration.ofSeconds(2).toMillis();
+
+    Properties properties = new Properties();
+    properties.put(CoordinatorConfig.CONFIG_CLUSTER, testCluster);
+    properties.put(CoordinatorConfig.CONFIG_ZK_ADDRESS, _zkConnectionString);
+    // custom heartbeat period of 2 second.
+    properties.put(CoordinatorConfig.CONFIG_HEARTBEAT_PERIOD_MS, String.valueOf(testHeartbeatPeriod));
+
+    final Coordinator.CoordinatorEventProcessor[] testCoordinatorEventProcessor = {null};
+    Coordinator coordinator =
+        createCoordinator(_zkConnectionString, testCluster, properties, new DummyTransportProviderAdminFactory(),
+            (cachedDatastreamReader, props) -> new Coordinator(cachedDatastreamReader, props) {
+              @Override
+              protected synchronized void createEventThread() {
+                testCoordinatorEventProcessor[0] = new CoordinatorEventProcessor() {
+                  // Mimicking the coordinator's event thread's runnable method.
+                  // 1. Sleeping before calling handleEvent to let zk session expiry
+                  //    thread acquire coordinator object before event thread enters
+                  //    handleEvent.
+                  // 2. Handling a No-Op Event.
+                  // 3. Notifying the zk session expiry threads to attempt acquiring the
+                  //    coordinator object.
+                  @Override
+                  public void run() {
+                    // This flag will be enabled when an interrupt was called
+                    // on the event thread while the event thread was sleeping.
+                    boolean isInterruptedInSleep = false;
+                    while (!isInterrupted()) {
+                      try {
+                        // Step 1
+                        // Making sure we sleep for more than heartbeat period to
+                        // mock the scenario where the zk session expiry thread
+                        // acquires the coordinator object before the event thread does.
+                        Thread.sleep(testHeartbeatPeriod + Duration.ofMillis(500).toMillis());
+                      } catch (InterruptedException e) {
+                        isInterruptedInSleep = true;
+                      }
+                      // Step 2
+                      // Handling an event requires acquiring the coordinator's object.
+                      handleEvent(new CoordinatorEvent(CoordinatorEvent.EventType.NO_OP, null));
+                      // Step 3
+                      notifyThreadsWaitingForCoordinatorObjectSynchronization();
+                      if (isInterruptedInSleep) {
+                        break;
+                      }
+                    }
+                  }
+
+                  @Override
+                  protected void notifyThreadsWaitingForCoordinatorObjectSynchronization() {
+                    super.notifyThreadsWaitingForCoordinatorObjectSynchronization();
+                  }
+                };
+                testCoordinatorEventProcessor[0].setDaemon(true);
+              }
+
+              @Override
+              CoordinatorEventProcessor getEventThread() {
+                return testCoordinatorEventProcessor[0];
+              }
+            });
+
+    coordinator.start();
+    ZkClient zkClient = new ZkClient(_zkConnectionString);
+
+    coordinator.onSessionExpired();
+    Assert.assertTrue(PollUtils.poll(coordinator::isZkSessionExpired, 100, testHeartbeatPeriod));
+    // Making sure we don't run into a deadlock scenario.
+    Assert.assertTrue(PollUtils.poll(() -> !coordinator.getEventThread().isAlive(), 100, testHeartbeatPeriod));
+
+    coordinator.stop();
+    zkClient.close();
+    coordinator.getDatastreamCache().getZkclient().close();
+  }
+
+  @Test
+  public void testSessionExpiryCallbackThreadAttemptingToAcquireCoordinatorObjectAfterHandlingEvent() throws Exception {
+    String testCluster = "dummyCluster";
+    long testHeartbeatPeriod = Duration.ofSeconds(2).toMillis();
+
+    Properties properties = new Properties();
+    properties.put(CoordinatorConfig.CONFIG_CLUSTER, testCluster);
+    properties.put(CoordinatorConfig.CONFIG_ZK_ADDRESS, _zkConnectionString);
+    // custom heartbeat period of 2 second.
+    properties.put(CoordinatorConfig.CONFIG_HEARTBEAT_PERIOD_MS, String.valueOf(testHeartbeatPeriod));
+
+    final Coordinator.CoordinatorEventProcessor[] testCoordinatorEventProcessor = {null};
+    Coordinator coordinator =
+        createCoordinator(_zkConnectionString, testCluster, properties, new DummyTransportProviderAdminFactory(),
+            (cachedDatastreamReader, props) -> new Coordinator(cachedDatastreamReader, props) {
+              @Override
+              protected synchronized void createEventThread() {
+                testCoordinatorEventProcessor[0] = new CoordinatorEventProcessor() {
+                  // Mimicking the coordinator's event thread's runnable method.
+                  // 1. Handling a No-Op Event.
+                  // 2. Sleeping after calling handleEvent to let zk session expiry
+                  //    thread wait for notification from event thread to access coordinator
+                  //    object.
+                  // 3. Notifying the zk session expiry threads to attempt acquiring the
+                  //    coordinator object.
+                  @Override
+                  public void run() {
+                    // This flag will be enabled when an interrupt was called
+                    // on the event thread while the event thread was sleeping.
+                    boolean isInterruptedInSleep = false;
+                    // Step 1
+                    // Handling an event requires acquiring the coordinator's object.
+                    handleEvent(new CoordinatorEvent(CoordinatorEvent.EventType.NO_OP, null));
+                    while (!isInterrupted()) {
+                      try {
+                        // Step 2
+                        // Making sure we sleep for less than heartbeat period to
+                        // mock the scenario where the zk session expiry thread
+                        // is waiting for notification from the event thread.
+                        Thread.sleep(testHeartbeatPeriod - Duration.ofMillis(500).toMillis());
+                      } catch (InterruptedException e) {
+                        isInterruptedInSleep = true;
+                      }
+                      // Step 3
+                      notifyThreadsWaitingForCoordinatorObjectSynchronization();
+                      if (isInterruptedInSleep) {
+                        break;
+                      }
+                    }
+                  }
+
+                  @Override
+                  protected void notifyThreadsWaitingForCoordinatorObjectSynchronization() {
+                    super.notifyThreadsWaitingForCoordinatorObjectSynchronization();
+                  }
+                };
+                testCoordinatorEventProcessor[0].setDaemon(true);
+              }
+
+              @Override
+              CoordinatorEventProcessor getEventThread() {
+                return testCoordinatorEventProcessor[0];
+              }
+            });
+
+
+    coordinator.start();
+    ZkClient zkClient = new ZkClient(_zkConnectionString);
+
+    coordinator.onSessionExpired();
+    Assert.assertTrue(PollUtils.poll(coordinator::isZkSessionExpired, 100, testHeartbeatPeriod));
+    // Making sure we don't run into a deadlock scenario.
+    Assert.assertTrue(PollUtils.poll(() -> !coordinator.getEventThread().isAlive(), 100, testHeartbeatPeriod));
+
+    coordinator.stop();
+    zkClient.close();
+    coordinator.getDatastreamCache().getZkclient().close();
+  }
+
   // 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

@@ -321,23 +321,30 @@ public void start() {
         _heartbeatPeriod.toMillis() * 3, _heartbeatPeriod.toMillis(), TimeUnit.MILLISECONDS);
   }
 
-  private synchronized void createEventThread() {
+  protected synchronized void createEventThread() {
     _eventThread = new CoordinatorEventProcessor();
     _eventThread.setDaemon(true);
   }
 
   private synchronized void startEventThread() {
     if (!_shutdown) {
-      _eventThread.start();
+      CoordinatorEventProcessor eventThread = getEventThread();
+      eventThread.start();
     }
   }
 
   private synchronized boolean stopEventThread() {
     // interrupt the thread if it's not gracefully shutdown
-    while (_eventThread.isAlive()) {
+    CoordinatorEventProcessor eventThread = getEventThread();
+    while (eventThread.isAlive()) {
+      // Waits to acquire the Coordinator object for a maximum of _heartbeat period.
+      // The time bound waiting prevents the caller thread to not infinitely wait if
+      // the event thread is already shutdown.
+      waitForNotificationFromEventThread(_heartbeatPeriod);
       try {
-        _eventThread.interrupt();
-        _eventThread.join(EVENT_THREAD_SHORT_JOIN_TIMEOUT);
+        _log.info("Attempting to interrupt the event thread.");
+        eventThread.interrupt();
+        eventThread.join(EVENT_THREAD_SHORT_JOIN_TIMEOUT);
       } catch (InterruptedException e) {
         _log.warn("Exception caught while interrupting the event thread", e);
         return true;
@@ -346,16 +353,38 @@ private synchronized boolean stopEventThread() {
     return false;
   }
 
+  // Waiting for the event thread to die.
   private synchronized boolean waitForEventThreadToJoin() {
+    CoordinatorEventProcessor eventThread = getEventThread();
+    if (!eventThread.isAlive()) {
+      return false;
+    }
+    // Waits to acquire the Coordinator object for a maximum of _heartbeat period.
+    // The time bound waiting prevents the caller thread to not infinitely wait if
+    // the event thread is already shutdown.
+    waitForNotificationFromEventThread(_heartbeatPeriod);
     try {
-      _eventThread.join(EVENT_THREAD_LONG_JOIN_TIMEOUT);
+      _log.info("Waiting for {} milliseconds for the event thread to die.", EVENT_THREAD_LONG_JOIN_TIMEOUT);
+      eventThread.join(EVENT_THREAD_LONG_JOIN_TIMEOUT);
     } catch (InterruptedException e) {
       _log.warn("Exception caught while waiting the event thread to stop", e);
       return true;
     }
     return false;
   }
 
+  // Waits for a notification for specified duration from the event thread before acquiring the Coordinator object.
+  private void waitForNotificationFromEventThread(Duration duration) {
+    try {
+      // This intrinsic conditional variable helps to halt threads (zk callback threads, main server thread) before
+      // attempting to acquire the Coordinator object. We never halt the event thread (coordinator thread)
+      // explicitly via this CV.
+      wait(duration.toMillis());
+    } catch (InterruptedException e) {
+      _log.warn("Exception caught while waiting for the notification from the event thread", e);
+    }
+  }
+
   /**
    * Stop coordinator (and all connectors)
    */
@@ -2296,7 +2325,7 @@ CachedDatastreamReader getDatastreamCache() {
     return _datastreamCache;
   }
 
-  private class CoordinatorEventProcessor extends Thread {
+  protected class CoordinatorEventProcessor extends Thread {
     @Override
     public void run() {
       _log.info("START CoordinatorEventProcessor thread");
@@ -2305,6 +2334,7 @@ public void run() {
           CoordinatorEvent event = _eventQueue.take();
           if (event != null) {
             handleEvent(event);
+            notifyThreadsWaitingForCoordinatorObjectSynchronization();
           }
         } catch (InterruptedException e) {
           _log.warn("CoordinatorEventProcessor interrupted", e);
@@ -2315,6 +2345,12 @@ public void run() {
       }
       _log.info("END CoordinatorEventProcessor");
     }
+
+    // Via the intrinsic conditional variable, notify other threads that might
+    // be waiting on acquiring access on the Coordinator object.
+    protected synchronized void notifyThreadsWaitingForCoordinatorObjectSynchronization() {
+      notifyAll();
+    }
   }
 
   @VisibleForTesting