consensus: use bracketWithRegistry for ChainSync servers

We intend for this change to make it more difficult to repeat the mistake underlying Issue #2870.
IntersectMBO · Jan 26, 2021 · c8132f2 · c8132f2
1 parent 2f85f69
commit c8132f2
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 25 deletions.
diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/MiniProtocol/ChainSync/Server.hs b/ouroboros-consensus/src/Ouroboros/Consensus/MiniProtocol/ChainSync/Server.hs
@@ -33,20 +33,20 @@ import           Ouroboros.Consensus.Util.ResourceRegistry (ResourceRegistry)
 
 
 chainSyncHeaderServerFollower
-    :: Functor m
-    => ChainDB m blk
-    -> ResourceRegistry m
-    -> m (Follower m blk (WithPoint blk (SerialisedHeader blk)))
-chainSyncHeaderServerFollower chainDB registry =
-    ($ registry) <$> ChainDB.newFollower chainDB getSerialisedHeaderWithPoint
+    :: ChainDB m blk
+    -> m (   ResourceRegistry m
+          -> Follower m blk (WithPoint blk (SerialisedHeader blk))
+         )
+chainSyncHeaderServerFollower chainDB =
+    ChainDB.newFollower chainDB getSerialisedHeaderWithPoint
 
 chainSyncBlockServerFollower
-    :: Functor m
-    => ChainDB m blk
-    -> ResourceRegistry m
-    -> m (Follower m blk (WithPoint blk (Serialised blk)))
-chainSyncBlockServerFollower chainDB registry =
-    ($ registry) <$> ChainDB.newFollower chainDB getSerialisedBlockWithPoint
+    :: ChainDB m blk
+    -> m (   ResourceRegistry m
+          -> Follower m blk (WithPoint blk (Serialised blk))
+         )
+chainSyncBlockServerFollower chainDB =
+    ChainDB.newFollower chainDB getSerialisedBlockWithPoint
 
 -- | Chain Sync Server for block headers for a given a 'ChainDB'.
 --

diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/Network/NodeToClient.hs b/ouroboros-consensus/src/Ouroboros/Consensus/Network/NodeToClient.hs
@@ -359,17 +359,16 @@ mkApps kernel Tracers {..} Codecs {..} Handlers {..} =
       -> m ((), Maybe bCS)
     aChainSyncServer them channel = do
       labelThisThread "LocalChainSyncServer"
-      withRegistry $ \registry ->
-        bracket
-          (chainSyncBlockServerFollower (getChainDB kernel) registry)
-          ChainDB.followerClose
-          (\flr -> runPeer
+      bracketWithRegistry
+        (chainSyncBlockServerFollower (getChainDB kernel))
+        ChainDB.followerClose
+        $ \_registry flr ->
+          runPeer
             (contramap (TraceLabelPeer them) tChainSyncTracer)
             cChainSyncCodec
             channel
             $ chainSyncServerPeer
             $ hChainSyncServer flr
-          )
 
     aTxSubmissionServer
       :: localPeer

diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/Network/NodeToNode.hs b/ouroboros-consensus/src/Ouroboros/Consensus/Network/NodeToNode.hs
@@ -487,20 +487,19 @@ mkApps kernel Tracers {..} Codecs {..} genChainSyncTimeout Handlers {..} =
       -> m ((), Maybe bCS)
     aChainSyncServer version them channel = do
       labelThisThread "ChainSyncServer"
-      withRegistry $ \registry -> do
-        chainSyncTimeout <- genChainSyncTimeout
-        bracket
-          (chainSyncHeaderServerFollower (getChainDB kernel) registry)
-          ChainDB.followerClose
-          (\flr -> runPeerWithLimits
+      chainSyncTimeout <- genChainSyncTimeout
+      bracketWithRegistry
+        (chainSyncHeaderServerFollower (getChainDB kernel))
+        ChainDB.followerClose
+        $ \_registry flr ->
+          runPeerWithLimits
             (contramap (TraceLabelPeer them) tChainSyncSerialisedTracer)
             cChainSyncCodecSerialised
             (byteLimitsChainSync (const 0)) -- TODO: Real Bytelimits, see #1727
             (timeLimitsChainSync chainSyncTimeout)
             channel
             $ chainSyncServerPeer
             $ hChainSyncServer flr version
-          )
 
     aBlockFetchClient
       :: NodeToNodeVersion