ChainSyncClient: respond to ControlMessageSTM

ouroboros-consensus/src/Ouroboros/Consensus/MiniProtocol/ChainSync/Client.hs

@@ -54,6 +54,7 @@ import Network.TypedProtocol.Pipelined
 import Ouroboros.Network.AnchoredFragment (AnchoredFragment (..))
 import qualified Ouroboros.Network.AnchoredFragment as AF
 import Ouroboros.Network.Block (Tip, getTipBlockNo, getTipPoint)
+import Ouroboros.Network.Mux (ControlMessage (..), ControlMessageSTM)
 import Ouroboros.Network.Protocol.ChainSync.ClientPipelined
 import Ouroboros.Network.Protocol.ChainSync.PipelineDecision
 
@@ -342,6 +343,7 @@ chainSyncClient
  -> TopLevelConfig blk
  -> ChainDbView m blk
  -> NodeToNodeVersion
+ -> ControlMessageSTM m
  -> StrictTVar m (AnchoredFragment (Header blk))
  -> Consensus ChainSyncClientPipelined blk m
 chainSyncClient mkPipelineDecision0 tracer cfg
@@ -351,8 +353,9 @@ chainSyncClient mkPipelineDecision0 tracer cfg
  , getOurTip
  , getIsInvalidBlock
  }
- _version
- varCandidate = ChainSyncClientPipelined $
+ _version
+ controlMessageSTM
+ varCandidate = ChainSyncClientPipelined $
  continueWithState () $ initialise
  where
  -- | Start ChainSync by looking for an intersection between our current
@@ -523,22 +526,31 @@ chainSyncClient mkPipelineDecision0 tracer cfg
  (KnownIntersectionState blk)
  (ClientPipelinedStIdle n)
  nextStep mkPipelineDecision n theirTip = Stateful $ \kis -> do
- mKis' <- atomically $ intersectsWithCurrentChain kis
- case mKis' of
- Just kis'@KnownIntersectionState { theirFrag } -> do
- -- Our chain (tip) didn't change or if it did, it still intersects
- -- with the candidate fragment, so we can continue requesting the
- -- next block.
- atomically $ writeTVar varCandidate theirFrag
- let candTipBlockNo = AF.headBlockNo theirFrag
- return $ requestNext kis' mkPipelineDecision n theirTip candTipBlockNo
- Nothing ->
- -- Our chain (tip) has changed and it no longer intersects with the
- -- candidate fragment, so we have to find a new intersection, but
- -- first drain the pipe.
+ atomically controlMessageSTM >>= \case
+ -- We have been asked to terminate the client
+ Terminate ->
  continueWithState ()
  $ drainThePipe n
- $ findIntersection NoMoreIntersection
+ $ Stateful $ const $ return $ terminate AskedToTerminate
+ -- Continue
+ _ -> do
+ mKis' <- atomically $ intersectsWithCurrentChain kis
+ case mKis' of
+ Just kis'@KnownIntersectionState { theirFrag } -> do
+ -- Our chain (tip) didn't change or if it did, it still intersects
+ -- with the candidate fragment, so we can continue requesting the
+ -- next block.
+ atomically $ writeTVar varCandidate theirFrag
+ let candTipBlockNo = AF.headBlockNo theirFrag
+ return $
+ requestNext kis' mkPipelineDecision n theirTip candTipBlockNo
+ Nothing ->
+ -- Our chain (tip) has changed and it no longer intersects with
+ -- the candidate fragment, so we have to find a new intersection,
+ -- but first drain the pipe.
+ continueWithState ()
+ $ drainThePipe n
+ $ findIntersection NoMoreIntersection
 
  -- | "Drain the pipe": collect and discard all in-flight responses and
  -- finally execute the given action.
@@ -1000,6 +1012,9 @@ data ChainSyncClientResult =
  (Our (Tip blk))
  (Their (Tip blk))
 
+ -- | We were asked to terminate via the 'ControlMessageSTM'
+ | AskedToTerminate
+
 deriving instance Show ChainSyncClientResult
 
 instance Eq ChainSyncClientResult where
@@ -1021,6 +1036,9 @@ instance Eq ChainSyncClientResult where
  Just Refl -> (a, b) == (a', b')
  NoMoreIntersection{} == _ = False
 
+ AskedToTerminate == AskedToTerminate = True
+ AskedToTerminate == _ = False
+
 {-------------------------------------------------------------------------------
  Exception
 -------------------------------------------------------------------------------}

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

@@ -99,6 +99,7 @@ import Ouroboros.Consensus.Storage.ChainDB.Serialisation
 data Handlers m peer blk = Handlers {
  hChainSyncClient
  :: NodeToNodeVersion
+ -> ControlMessageSTM m
  -> StrictTVar m (AnchoredFragment (Header blk))
  -> ChainSyncClientPipelined (Header blk) (Tip blk) m ChainSyncClientResult
  -- TODO: we should consider either bundling these context parameters