inbound-governor: RemoteHot → RemoteWarm transition

Be careful when overwriting the state.
IntersectMBO · Sep 23, 2021 · 04fee6f · 04fee6f
1 parent 0b1f356
commit 04fee6f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/ouroboros-network-framework/src/Ouroboros/Network/InboundGovernor.hs b/ouroboros-network-framework/src/Ouroboros/Network/InboundGovernor.hs
@@ -286,7 +286,11 @@ inboundGovernor trTracer tracer serverControlChannel inboundIdleTimeout
 
                   let isHot = mpdMiniProtocolTemp == Hot
                       state' = ( if isHot
-                                 then updateRemoteState tConnId RemoteWarm
+                                 then mapRemoteState tConnId
+                                        $ \remoteState ->
+                                          case remoteState of
+                                            RemoteHot -> RemoteWarm
+                                            _         -> remoteState
                                  else id
                                )
                              . updateMiniProtocol tConnId num completionAction

diff --git a/ouroboros-network-framework/src/Ouroboros/Network/InboundGovernor/State.hs b/ouroboros-network-framework/src/Ouroboros/Network/InboundGovernor/State.hs
@@ -22,6 +22,7 @@ module Ouroboros.Network.InboundGovernor.State
   , updateMiniProtocol
   , RemoteState (.., RemoteEstablished)
   , updateRemoteState
+  , mapRemoteState
   , MiniProtocolData (..)
   ) where
 
@@ -283,3 +284,18 @@ updateRemoteState connId csRemoteState state =
           connId
           (igsConnections state)
     }
+
+mapRemoteState :: Ord peerAddr
+               => ConnectionId peerAddr
+               -> (RemoteState m -> RemoteState m)
+               -> InboundGovernorState muxMode peerAddr m a b
+               -> InboundGovernorState muxMode peerAddr m a b
+mapRemoteState connId fn state =
+    state {
+      igsConnections =
+        Map.adjust
+          (\connState@ConnectionState { csRemoteState } ->
+            connState { csRemoteState = fn csRemoteState })
+          connId
+          (igsConnections state)
+    }