From 6f506180404dc581589e607418f2499a199766da Mon Sep 17 00:00:00 2001
From: Robinson <robinson@dorkbox.com>
Date: Mon, 4 Dec 2023 13:36:47 +0100
Subject: [PATCH] Now properly waits for event dispatcher to shutdown in unit
 tests

---
 src/dorkbox/network/connection/EndPoint.kt    | 12 +++++++
 .../network/connection/EventDispatcher.kt     | 33 ++++++++++++++++---
 test/dorkboxTest/network/BaseTest.kt          |  3 +-
 3 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/src/dorkbox/network/connection/EndPoint.kt b/src/dorkbox/network/connection/EndPoint.kt
index 3a0a1dd6..55ba3411 100644
--- a/src/dorkbox/network/connection/EndPoint.kt
+++ b/src/dorkbox/network/connection/EndPoint.kt
@@ -1079,6 +1079,18 @@ abstract class EndPoint<CONNECTION : Connection> private constructor(val type: C
         return networkEventPoller.isDispatch()
     }
 
+    /**
+     * Shuts-down each event dispatcher executor, and waits for it to gracefully shutdown.
+     *
+     * Once shutdown, it cannot be restarted and the application MUST recreate the endpoint
+     *
+     * @param timeout how long to wait, must be > 0
+     * @param timeoutUnit what the unit count is
+     */
+    fun shutdownEventDispatcher(timeout: Long = 15, timeoutUnit: TimeUnit = TimeUnit.SECONDS) {
+        logger.info("Waiting for Event Dispatcher to shutdown...")
+        eventDispatch.shutdownAndWait(timeout, timeoutUnit)
+    }
 
     /**
      * Reset the running state when there's an error starting up
diff --git a/src/dorkbox/network/connection/EventDispatcher.kt b/src/dorkbox/network/connection/EventDispatcher.kt
index 0501cf2d..8c50874f 100644
--- a/src/dorkbox/network/connection/EventDispatcher.kt
+++ b/src/dorkbox/network/connection/EventDispatcher.kt
@@ -42,6 +42,10 @@ internal class EventDispatcher(val type: String) {
         fun isDispatch(): Boolean {
             return dispatcher.isDispatch(type)
         }
+
+        fun shutdownAndWait(timeout: Long, timeoutUnit: TimeUnit) {
+            dispatcher.shutdownAndWait(type, timeout, timeoutUnit)
+        }
     }
 
     companion object {
@@ -75,8 +79,6 @@ internal class EventDispatcher(val type: String) {
         )
     }.toTypedArray()
 
-
-
     val HANDSHAKE: ED
     val CONNECT: ED
     val ERROR: ED
@@ -97,6 +99,21 @@ internal class EventDispatcher(val type: String) {
     }
 
 
+    /**
+     * Shuts-down each event dispatcher executor, and waits for it to gracefully shutdown. Once shutdown, it cannot be restarted.
+     *
+     * @param timeout how long to wait
+     * @param timeoutUnit what the unit count is
+     */
+    fun shutdownAndWait(timeout: Long, timeoutUnit: TimeUnit) {
+        require(timeout > 0) { logger.error("The EventDispatcher shutdown timeout must be > 0!") }
+
+        HANDSHAKE.shutdownAndWait(timeout, timeoutUnit)
+        CONNECT.shutdownAndWait(timeout, timeoutUnit)
+        ERROR.shutdownAndWait(timeout, timeoutUnit)
+        CLOSE.shutdownAndWait(timeout, timeoutUnit)
+    }
+
     /**
      *  Checks if the current execution thread is running inside one of the event dispatchers.
      */
@@ -122,9 +139,17 @@ internal class EventDispatcher(val type: String) {
     }
 
     /**
-     * Each event type runs inside its own coroutine dispatcher.
+     * shuts-down the current execution thread and waits for it complete.
+     */
+    private fun shutdownAndWait(type: EDType, timeout: Long, timeoutUnit: TimeUnit) {
+        executors[type.ordinal].shutdown()
+        executors[type.ordinal].awaitTermination(timeout, timeoutUnit)
+    }
+
+    /**
+     * Each event type runs inside its own thread executor.
      *
-     * We want EACH event type to run in its own dispatcher... on its OWN thread, in order to prevent deadlocks
+     * We want EACH event type to run in its own executor... on its OWN thread, in order to prevent deadlocks
      * This is because there are blocking dependencies: DISCONNECT -> CONNECT.
      *
      * If an event is RE-ENTRANT, then it will immediately execute!
diff --git a/test/dorkboxTest/network/BaseTest.kt b/test/dorkboxTest/network/BaseTest.kt
index f26c3ef4..ad89aad1 100644
--- a/test/dorkboxTest/network/BaseTest.kt
+++ b/test/dorkboxTest/network/BaseTest.kt
@@ -299,7 +299,6 @@ abstract class BaseTest {
             endPoint.stopDriver()
         }
 
-
         // run actions before we actually shutdown, but after we wait
         if (!successClients || !successServers) {
             Assert.fail("Shutdown latch not triggered ($successClients|$successServers)!")
@@ -308,6 +307,7 @@ abstract class BaseTest {
         // we must always make sure that aeron is shut-down before starting again.
         clients.forEach { endPoint ->
             endPoint.ensureStopped()
+            endPoint.shutdownEventDispatcher() // once shutdown, it cannot be restarted!
 
             if (!Client.ensureStopped(endPoint.config.copy())) {
                 throw IllegalStateException("Unable to continue, AERON client was unable to stop.")
@@ -316,6 +316,7 @@ abstract class BaseTest {
 
         servers.forEach { endPoint ->
             endPoint.ensureStopped()
+            endPoint.shutdownEventDispatcher() // once shutdown, it cannot be restarted!
 
             if (!Client.ensureStopped(endPoint.config.copy())) {
                 throw IllegalStateException("Unable to continue, AERON server was unable to stop.")