From 5f841068a18cde14456bd2486610a8a96ef29bb3 Mon Sep 17 00:00:00 2001
From: Armando Santos <armando@well-typed.com>
Date: Fri, 29 Sep 2023 16:58:33 +0100
Subject: [PATCH] Fix NotReleasedConnections in Server2 tests

---
 ouroboros-network-framework/CHANGELOG.md      |  2 ++
 .../Test/Ouroboros/Network/Server2/Sim.hs     | 30 +++++++++----------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/ouroboros-network-framework/CHANGELOG.md b/ouroboros-network-framework/CHANGELOG.md
index 1d0a2600368..5350d843eac 100644
--- a/ouroboros-network-framework/CHANGELOG.md
+++ b/ouroboros-network-framework/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 * Split `test` component into `io-tests` and `sim-tests`.
 
+* Fix Server2 [sim test](https://github.com/input-output-hk/ouroboros-network/issues/4607) by synchronizing connection/disconnection events.
+
 ## 0.9.0.0 -- 2023-08-21
 
 ### Breaking changes
diff --git a/ouroboros-network-framework/sim-tests/Test/Ouroboros/Network/Server2/Sim.hs b/ouroboros-network-framework/sim-tests/Test/Ouroboros/Network/Server2/Sim.hs
index 56be217f203..61400534f73 100644
--- a/ouroboros-network-framework/sim-tests/Test/Ouroboros/Network/Server2/Sim.hs
+++ b/ouroboros-network-framework/sim-tests/Test/Ouroboros/Network/Server2/Sim.hs
@@ -24,7 +24,7 @@ import           Control.Applicative (Alternative)
 import qualified Control.Concurrent.Class.MonadSTM as LazySTM
 import           Control.Concurrent.Class.MonadSTM.Strict
 import           Control.Exception (SomeAsyncException (..))
-import           Control.Monad (replicateM, when)
+import           Control.Monad (replicateM)
 import           Control.Monad.Class.MonadAsync
 import           Control.Monad.Class.MonadFork
 import           Control.Monad.Class.MonadSay
@@ -777,8 +777,7 @@ multinodeExperiment inboundTrTracer trTracer inboundTracer cmTracer
                           (mkNextRequests connVar)
                           timeLimitsHandshake
                           acceptedConnLimit
-                          ( \ connectionManager _ serverAsync -> do
-                            linkOnly (const True) serverAsync
+                          ( \ connectionManager _ _serverAsync -> do
                             connectionLoop SingInitiatorResponderMode localAddr cc connectionManager Map.empty connVar
                             return Nothing
                           )
@@ -829,12 +828,11 @@ multinodeExperiment inboundTrTracer trTracer inboundTracer cmTracer
          -> StrictTVar m (Map.Map (ConnectionId peerAddr) (TemperatureBundle (StrictTQueue m [req])))
          -- ^ mini protocol queues
          -> m ()
-    connectionLoop muxMode localAddr cc cm connMap0 connVar = go True connMap0
+    connectionLoop muxMode localAddr cc cm connMap0 connVar = go connMap0
       where
-        go :: Bool -- if false do not run 'unregisterOutboundConnection'
-           -> Map.Map peerAddr (HandleWithExpandedCtx muxMode peerAddr DataFlowProtocolData ByteString m [resp] a) -- active connections
+        go :: Map.Map peerAddr (HandleWithExpandedCtx muxMode peerAddr DataFlowProtocolData ByteString m [resp] a) -- active connections
            -> m ()
-        go !unregister !connMap = atomically (readTQueue cc) >>= \ case
+        go !connMap = atomically (readTQueue cc) >>= \ case
           NewConnection remoteAddr -> do
             let mkQueue :: forall pt. SingProtocolTemperature pt
                         -> STM m (StrictTQueue m [req])
@@ -852,18 +850,20 @@ multinodeExperiment inboundTrTracer trTracer inboundTracer cmTracer
                           $ requestOutboundConnection cm remoteAddr
             case connHandle of
               Left _ ->
-                go False connMap
+                go connMap
               Right (Connected _ _ h) -> do
                 qs <- atomically $ traverse id $ makeBundle mkQueue
                 atomically $ modifyTVar connVar
                            $ Map.insert (connId remoteAddr) qs
-                go True (Map.insert remoteAddr h connMap)
+                go (Map.insert remoteAddr h connMap)
               Right Disconnected {} -> return ()
           Disconnect remoteAddr -> do
-            atomically $ modifyTVar connVar $ Map.delete (connId remoteAddr)
-            when unregister $
-              void (unregisterOutboundConnection cm remoteAddr)
-            go False (Map.delete remoteAddr connMap)
+            atomically $ do
+              m <- readTVar connVar
+              check (Map.member (connId remoteAddr) m)
+              writeTVar connVar (Map.delete (connId remoteAddr) m)
+            void (unregisterOutboundConnection cm remoteAddr)
+            go (Map.delete remoteAddr connMap)
           RunMiniProtocols remoteAddr reqs -> do
             atomically $ do
               mqs <- Map.lookup (connId remoteAddr) <$> readTVar connVar
@@ -897,8 +897,8 @@ multinodeExperiment inboundTrTracer trTracer inboundTracer cmTracer
                   Left  {} -> do
                     atomically
                       $ modifyTVar connVar (Map.delete (connId remoteAddr))
-                    go unregister (Map.delete remoteAddr connMap)
-                  Right {} -> go unregister connMap
+                    go (Map.delete remoteAddr connMap)
+                  Right {} -> go connMap
           Shutdown -> return ()
           where
             connId remoteAddr = ConnectionId { localAddress  = localAddr