Introduce bracketObjectDiffusionInbound

Author: agustinmista

ouroboros-consensus-diffusion/src/ouroboros-consensus-diffusion/Ouroboros/Consensus/Network/NodeToNode.hs

@@ -69,6 +69,10 @@ import Ouroboros.Consensus.MiniProtocol.ChainSync.Client
 import qualified Ouroboros.Consensus.MiniProtocol.ChainSync.Client as CsClient
 import Ouroboros.Consensus.MiniProtocol.ChainSync.Server
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Inbound (objectDiffusionInbound)
+import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Inbound.State
+  ( ObjectDiffusionInboundStateView
+  , bracketObjectDiffusionInbound
+  )
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.ObjectPool.PerasCert
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Outbound (objectDiffusionOutbound)
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.PerasCert
@@ -212,6 +216,7 @@ data Handlers m addr blk = Handlers
   , hPerasCertDiffusionClient ::
       NodeToNodeVersion ->
       ControlMessageSTM m ->
+      ObjectDiffusionInboundStateView m ->
       ConnectionId addr ->
       PerasCertDiffusionInboundPipelined blk m ()
   , hPerasCertDiffusionServer ::
@@ -314,7 +319,7 @@ mkHandlers
             (mapTxSubmissionMempoolReader txForgetValidated $ getMempoolReader getMempool)
             (getMempoolWriter getMempool)
             version
-      , hPerasCertDiffusionClient = \version controlMessageSTM peer ->
+      , hPerasCertDiffusionClient = \version controlMessageSTM state peer ->
           objectDiffusionInbound
             (contramap (TraceLabelPeer peer) (Node.perasCertDiffusionInboundTracer tracers))
             ( perasCertDiffusionMaxFifoLength miniProtocolParameters
@@ -324,6 +329,7 @@ mkHandlers
             (makePerasCertPoolWriterFromChainDB $ getChainDB)
             version
             controlMessageSTM
+            state
       , hPerasCertDiffusionServer = \version peer ->
           objectDiffusionOutbound
             (contramap (TraceLabelPeer peer) (Node.perasCertDiffusionOutboundTracer tracers))
@@ -862,17 +868,21 @@ mkApps kernel rng Tracers{..} mkCodecs ByteLimits{..} chainSyncTimeouts lopBucke
       }
     channel = do
       labelThisThread "PerasCertDiffusionClient"
-      ((), trailing) <-
-        runPipelinedPeerWithLimits
-          (TraceLabelPeer them `contramap` tPerasCertDiffusionTracer)
-          (cPerasCertDiffusionCodec (mkCodecs version))
-          blPerasCertDiffusion
-          timeLimitsObjectDiffusion
-          channel
-          ( objectDiffusionInboundPeerPipelined
-              (hPerasCertDiffusionClient version controlMessageSTM them)
-          )
-      return (NoInitiatorResult, trailing)
+      bracketObjectDiffusionInbound
+        (getPerasCertDiffusionHandles kernel)
+        them
+        $ \state -> do
+          ((), trailing) <-
+            runPipelinedPeerWithLimits
+              (TraceLabelPeer them `contramap` tPerasCertDiffusionTracer)
+              (cPerasCertDiffusionCodec (mkCodecs version))
+              blPerasCertDiffusion
+              timeLimitsObjectDiffusion
+              channel
+              ( objectDiffusionInboundPeerPipelined
+                  (hPerasCertDiffusionClient version controlMessageSTM state them)
+              )
+          return (NoInitiatorResult, trailing)
 
   aPerasCertDiffusionServer ::
     NodeToNodeVersion ->

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ObjectDiffusion/Inbound.hs

@@ -38,6 +38,9 @@ import Data.Word (Word64)
 import GHC.Generics (Generic)
 import Network.TypedProtocol.Core (N (Z), Nat (..), natToInt)
 import NoThunks.Class (NoThunks (..), unsafeNoThunks)
+import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Inbound.State
+  ( ObjectDiffusionInboundStateView
+  )
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.ObjectPool.API
 import Ouroboros.Network.ControlMessage
 import Ouroboros.Network.NodeToNode.Version (NodeToNodeVersion)
@@ -131,13 +134,15 @@ objectDiffusionInbound ::
   ObjectPoolWriter objectId object m ->
   NodeToNodeVersion ->
   ControlMessageSTM m ->
+  ObjectDiffusionInboundStateView m ->
   ObjectDiffusionInboundPipelined objectId object m ()
 objectDiffusionInbound
   tracer
   (maxFifoLength, maxNumIdsToReq, maxNumObjectsToReq)
   ObjectPoolWriter{..}
   _version
-  controlMessageSTM =
+  controlMessageSTM
+  _state =
     ObjectDiffusionInboundPipelined $ do
       continueWithStateM (go Zero) initialInboundSt
    where

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ObjectDiffusion/Inbound/State.hs

@@ -2,32 +2,40 @@
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE UndecidableInstances #-}
 
 module Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Inbound.State
   ( ObjectDiffusionInboundState (..)
   , ObjectDiffusionInboundHandle (..)
   , ObjectDiffusionInboundHandleCollection (..)
+  , ObjectDiffusionInboundStateView (..)
   , newObjectDiffusionInboundHandleCollection
+  , bracketObjectDiffusionInbound
   )
 where
 
+import Control.Monad.Class.MonadThrow (bracket)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import GHC.Generics (Generic)
 import NoThunks.Class (NoThunks)
 import Ouroboros.Consensus.Block (BlockSupportsProtocol, HasHeader, Header)
+import Ouroboros.Consensus.MiniProtocol.Util.Idling (Idling (..))
 import Ouroboros.Consensus.Util.IOLike
   ( IOLike (..)
   , MonadSTM (..)
   , StrictTVar
   , modifyTVar
   , newTVar
+  , newTVarIO
   , readTVar
   )
 
 -- | An ObjectDiffusion inbound client state that's used by other components.
+--
+-- NOTE: 'blk' is not needed for now, but we keep it for future use.
 data ObjectDiffusionInboundState blk = ObjectDiffusionInboundState
   { odisIdling :: !Bool
   -- ^ Whether we have received all objects from a peer
@@ -82,3 +90,47 @@ newObjectDiffusionInboundHandleCollection = do
       , odihcRemoveAllHandles =
           modifyTVar handlesMap (const mempty)
       }
+
+-- | Interface for the ObjectDiffusion client to its state allocated by
+-- 'bracketChainSyncClient'.
+data ObjectDiffusionInboundStateView m = ObjectDiffusionInboundStateView
+  { odisvIdling :: !(Idling m)
+  }
+  deriving stock Generic
+
+bracketObjectDiffusionInbound ::
+  forall m peer blk a.
+  (IOLike m, HasHeader blk, NoThunks (Header blk)) =>
+  ObjectDiffusionInboundHandleCollection peer m blk ->
+  peer ->
+  (ObjectDiffusionInboundStateView m -> m a) ->
+  m a
+bracketObjectDiffusionInbound handles peer body =
+  mkObjectDiffusionInboundState >>= \odiState ->
+    bracket
+      (acquireContext odiState)
+      releaseContext
+      body
+ where
+  acquireContext odiState = atomically $ do
+    odihcAddHandle handles peer $
+      ObjectDiffusionInboundHandle
+        { odihState = odiState
+        }
+    return
+      ObjectDiffusionInboundStateView
+        { odisvIdling =
+            Idling
+              { idlingStart = atomically $ modifyTVar odiState $ \s -> s{odisIdling = True}
+              , idlingStop = atomically $ modifyTVar odiState $ \s -> s{odisIdling = False}
+              }
+        }
+
+  releaseContext _ = atomically $ do
+    odihcRemoveHandle handles peer
+
+  mkObjectDiffusionInboundState =
+    newTVarIO
+      ObjectDiffusionInboundState
+        { odisIdling = False
+        }

ouroboros-consensus/test/consensus-test/Test/Consensus/MiniProtocol/ObjectDiffusion/Smoke.hs

@@ -31,11 +31,15 @@ import NoThunks.Class (NoThunks)
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Inbound
   ( objectDiffusionInbound
   )
+import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Inbound.State
+  ( ObjectDiffusionInboundStateView (..)
+  )
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.ObjectPool.API
   ( ObjectPoolReader (..)
   , ObjectPoolWriter (..)
   )
 import Ouroboros.Consensus.MiniProtocol.ObjectDiffusion.Outbound (objectDiffusionOutbound)
+import qualified Ouroboros.Consensus.MiniProtocol.Util.Idling as Idling
 import Ouroboros.Consensus.Util.IOLike
   ( IOLike
   , MonadDelay (..)
@@ -267,6 +271,7 @@ prop_smoke_object_diffusion
               inboundPoolWriter
               nodeToNodeVersion
               (readTVar controlMessage)
+              ObjectDiffusionInboundStateView{odisvIdling = Idling.noIdling}
 
           outbound =
             objectDiffusionOutbound