ChainSync: let GSM disable and re-enable CSJ; also enable LoP in PreSyncing

nfrisby · nfrisby · commit a1e48213627d · 2025-04-30T12:06:08.000-07:00
diff --git a/ouroboros-consensus-diffusion/src/ouroboros-consensus-diffusion/Ouroboros/Consensus/Node/GSM.hs b/ouroboros-consensus-diffusion/src/ouroboros-consensus-diffusion/Ouroboros/Consensus/Node/GSM.hs
@@ -12,7 +12,6 @@ module Ouroboros.Consensus.Node.GSM (
   , DurationFromNow (..)
   , GsmEntryPoints (..)
   , GsmNodeKernelArgs (..)
-  , GsmState (..)
   , GsmView (..)
   , MarkerFileView (..)
   , WrapDurationUntilTooOld (..)
@@ -124,7 +123,7 @@ data GsmView m upstreamPeer selection chainSyncState = GsmView {
     setCaughtUpPersistentMark :: Bool -> m ()
     -- ^ EG touch/delete the marker file on disk
   ,
-    writeGsmState             :: GsmState -> m ()
+    writeGsmState             :: GsmTransition -> m ()
     -- ^ EG update the TVar that the Diffusion Layer monitors, or en-/disable
     -- certain components of Genesis
   ,
@@ -258,7 +257,7 @@ realGsmEntryPoints tracerArgs gsmView = GsmEntryPoints {
         (g', ev) <- blockWhileCaughtUp g
 
         setCaughtUpPersistentMark False
-        writeGsmState PreSyncing
+        writeGsmState CaughtUpPreSyncing
         traceWith tracer ev
 
         enterPreSyncing' g'
@@ -267,7 +266,7 @@ realGsmEntryPoints tracerArgs gsmView = GsmEntryPoints {
     enterPreSyncing' g = do
         blockUntilHonestAvailabilityAssumption
 
-        writeGsmState Syncing
+        writeGsmState PreSyncingSyncing
         traceWith tracer GsmEventPreSyncingToSyncing
 
         enterSyncing' g
@@ -283,12 +282,12 @@ realGsmEntryPoints tracerArgs gsmView = GsmEntryPoints {
 
         case mev of
           Nothing  -> do
-            writeGsmState PreSyncing
+            writeGsmState SyncingPreSyncing
             traceWith tracer GsmEventSyncingToPreSyncing
 
             enterPreSyncing' g
           Just ev -> do
-            writeGsmState CaughtUp
+            writeGsmState SyncingCaughtUp
             setCaughtUpPersistentMark True
             traceWith tracer ev
 
diff --git a/ouroboros-consensus-diffusion/src/ouroboros-consensus-diffusion/Ouroboros/Consensus/NodeKernel.hs b/ouroboros-consensus-diffusion/src/ouroboros-consensus-diffusion/Ouroboros/Consensus/NodeKernel.hs
@@ -275,11 +275,14 @@ initNodeKernel args@NodeKernelArgs { registry, cfg, tracers
               , GSM.setCaughtUpPersistentMark = \upd ->
                   (if upd then GSM.touchMarkerFile else GSM.removeMarkerFile)
                     gsmMarkerFileView
-              , GSM.writeGsmState = \gsmState ->
-                  atomicallyWithMonotonicTime $ \time -> do
-                    writeTVar varGsmState gsmState
-                    handles <- cschcMap varChainSyncHandles
-                    traverse_ (($ time) . ($ gsmState) . cschOnGsmStateChanged) handles
+              , GSM.writeGsmState             = \gsmTransition ->
+                  join $ atomicallyWithMonotonicTime $ \time -> do
+                    writeTVar varGsmState $ GSM.gsmNewState gsmTransition
+                    let nil           = pure ()
+                        snoc acc hndl = do
+                          io <- cschOnGsmStateChanged hndl gsmTransition time
+                          pure $ do acc; io
+                    cschcMap varChainSyncHandles >>= foldM snoc nil
               , GSM.isHaaSatisfied            = do
                   readTVar varOutboundConnectionsState <&> \case
                     -- See the upstream Haddocks for the exact conditions under
diff --git a/ouroboros-consensus-diffusion/test/consensus-test/Test/Consensus/GSM.hs b/ouroboros-consensus-diffusion/test/consensus-test/Test/Consensus/GSM.hs
@@ -221,8 +221,8 @@ prop_sequential1 ctx cmds = runSimQC $ do
                 atomically $ do
                     writeTVar varMarker $ if b then Present else Absent
           ,
-            GSM.writeGsmState = \x -> atomically $ do
-                writeTVar varGsmState x
+            GSM.writeGsmState =
+                atomically . writeTVar varGsmState . GSM.gsmNewState
           ,
             GSM.isHaaSatisfied =
                 isHaaSatisfied . Map.keysSet <$> readTVar varStates
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ChainSync/Client.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ChainSync/Client.hs
@@ -80,6 +80,7 @@ import           Control.Monad.Class.MonadTimer (MonadTimer)
 import           Control.Monad.Except (runExcept, throwError)
 import           Control.Tracer
 import           Data.Foldable (traverse_)
+import           Data.Function (fix)
 import           Data.Functor ((<&>))
 import           Data.Kind (Type)
 import           Data.Map.Strict (Map)
@@ -114,7 +115,7 @@ import qualified Ouroboros.Consensus.MiniProtocol.ChainSync.Client.HistoricityCh
 import qualified Ouroboros.Consensus.MiniProtocol.ChainSync.Client.InFutureCheck as InFutureCheck
 import qualified Ouroboros.Consensus.MiniProtocol.ChainSync.Client.Jumping as Jumping
 import           Ouroboros.Consensus.MiniProtocol.ChainSync.Client.State
-import           Ouroboros.Consensus.Node.GsmState (GsmState (..))
+import           Ouroboros.Consensus.Node.GsmState (GsmState (..), GsmTransition (..))
 import           Ouroboros.Consensus.Node.NetworkProtocolVersion
 import           Ouroboros.Consensus.Protocol.Abstract
 import           Ouroboros.Consensus.Storage.ChainDB (ChainDB)
@@ -404,14 +405,16 @@ bracketChainSyncClient
       cschJumping <- newTVarIO (Disengaged DisengagedDone)
       let handle = ChainSyncClientHandle {
               cschGDDKill = throwTo tid DensityTooLow
-            , cschOnGsmStateChanged = updateLopBucketConfig lopBucket
+            , cschOnGsmStateChanged = \gsmTransition time -> do
+                  updateLopBucketConfig lopBucket gsmTransition time
+                  pure $ pure ()
             , cschState
             , cschJumping
             , cschJumpInfo
             }
           insertHandle = atomicallyWithMonotonicTime $ \time -> do
             initialGsmState <- getGsmState
-            updateLopBucketConfig lopBucket initialGsmState time
+            initializeLopBucketConfig lopBucket initialGsmState time
             cschcAddHandle varHandles peer handle
           deleteHandle = atomically $ cschcRemoveHandle varHandles peer
       bracket_ insertHandle deleteHandle $ f Jumping.noJumping
@@ -425,16 +428,39 @@ bracketChainSyncClient
         tid <- myThreadId
         atomicallyWithMonotonicTime $ \time -> do
           initialGsmState <- getGsmState
-          updateLopBucketConfig lopBucket initialGsmState time
+          initializeLopBucketConfig lopBucket initialGsmState time
           cschJumpInfo <- newTVar Nothing
           context <- Jumping.makeContext varHandles jumpSize tracerCsj
-          Jumping.registerClient context peer cschState $ \cschJumping -> ChainSyncClientHandle
-            { cschGDDKill = throwTo tid DensityTooLow
-            , cschOnGsmStateChanged = updateLopBucketConfig lopBucket
-            , cschState
-            , cschJumping
-            , cschJumpInfo
-            }
+          Jumping.registerClient context peer cschState $ \cschJumping ->
+            fix $ \handle -> ChainSyncClientHandle   -- NB @handle@ only occurs under a lambda
+              { cschGDDKill = throwTo tid DensityTooLow
+              , cschOnGsmStateChanged = \gsmTransition time' -> do
+                  updateLopBucketConfig lopBucket gsmTransition time'
+                  let peerContext =
+                        context {Jumping.peer = peer, Jumping.handle = handle}
+                  -- 'cschJumpInfo' does not need to be reset. The ChainSync
+                  -- client constantly updates it, even when its not
+                  -- registered with CSJ. That's necessary: when this peer
+                  -- /re-registers/, it might immediately be the Dynamo, and
+                  -- the Dynamo must not have 'Nothing' in its
+                  -- 'cschJumpInfo'.
+                  case gsmTransition of
+                    PreSyncingSyncing  -> pure $ pure ()
+                    SyncingCaughtUp    -> do
+                      writeTVar cschJumping $ Disengaged DisengagedDone
+                         -- The GSM only transitions to CaughtUp if all peers
+                         -- have send MsgAwaitReply, so DisengagedDone is
+                         -- correct here.
+                      mbEv <- Jumping.unregisterClient peerContext
+                      pure $ traverse_ (traceWith (Jumping.tracer peerContext)) mbEv
+                    CaughtUpPreSyncing -> do
+                      mbEv <- Jumping.reregisterClient peerContext
+                      pure $ traverse_ (traceWith (Jumping.tracer peerContext)) mbEv
+                    SyncingPreSyncing  -> pure $ pure ()
+              , cschState
+              , cschJumping
+              , cschJumpInfo
+              }
 
     releaseContext (peerContext, _mbEv) = do
         mbEv <- atomically $ Jumping.unregisterClient peerContext
@@ -444,33 +470,44 @@ bracketChainSyncClient
         invalidBlockRejector
             tracer version pipelining getIsInvalidBlock (csCandidate <$> readTVar varState)
 
-    -- | Update the configuration of the bucket to match the given GSM state.
-    -- NOTE: The new level is currently the maximal capacity of the bucket;
-    -- maybe we want to change that later.
-    updateLopBucketConfig :: LeakyBucket.Handlers m -> GsmState -> Time -> STM m ()
-    updateLopBucketConfig lopBucket gsmState =
-      LeakyBucket.updateConfig lopBucket $ \_ ->
-        let config = lopBucketConfig gsmState in
-          (LeakyBucket.capacity config, config)
-
-    -- | Wrapper around 'LeakyBucket.execAgainstBucket' that handles the
-    -- disabled bucket by running the given action with dummy handlers.
-    lopBucketConfig :: GsmState -> LeakyBucket.Config m
-    lopBucketConfig gsmState =
-      case (gsmState, csBucketConfig) of
-        (Syncing, ChainSyncLoPBucketEnabled ChainSyncLoPBucketEnabledConfig {csbcCapacity, csbcRate}) ->
-            LeakyBucket.Config
-              { capacity = fromInteger $ csbcCapacity,
-                rate = csbcRate,
-                onEmpty = throwIO EmptyBucket,
-                fillOnOverflow = True
-              }
-        -- NOTE: If we decide to slow the bucket down when “almost caught-up”,
-        -- we should add a state to the GSM and corresponding configuration
-        -- fields and a bucket config here.
-        (_, ChainSyncLoPBucketDisabled) -> LeakyBucket.dummyConfig
-        (PreSyncing, ChainSyncLoPBucketEnabled _) -> LeakyBucket.dummyConfig
-        (CaughtUp, ChainSyncLoPBucketEnabled _) -> LeakyBucket.dummyConfig
+    mkLopConfig ChainSyncLoPBucketEnabledConfig {csbcCapacity, csbcRate} =
+      LeakyBucket.Config {
+          capacity       = fromInteger csbcCapacity
+        , rate           = csbcRate
+        , onEmpty        = throwIO EmptyBucket
+        , fillOnOverflow = True
+        }
+
+    -- | Update the configuration of the bucket as part of the given GSM transition.
+    updateLopBucketConfig :: LeakyBucket.Handlers m -> GsmTransition -> Time -> STM m ()
+    updateLopBucketConfig lopBucket gsmTransition =
+      LeakyBucket.updateConfig lopBucket $ \(oldLevel, _oldConfig) ->
+        case csBucketConfig of
+          ChainSyncLoPBucketDisabled     -> (oldLevel, LeakyBucket.dummyConfig)
+          ChainSyncLoPBucketEnabled csbc ->
+            let config = mkLopConfig csbc
+            in
+            case gsmTransition of
+              PreSyncingSyncing  -> (oldLevel, config)
+              SyncingCaughtUp    -> (oldLevel, LeakyBucket.dummyConfig)
+              CaughtUpPreSyncing -> (LeakyBucket.capacity config, config)
+              SyncingPreSyncing  -> (oldLevel, config)
+
+    initializeLopBucketConfig :: LeakyBucket.Handlers m -> GsmState -> Time -> STM m ()
+    initializeLopBucketConfig lopBucket gsmState =
+      LeakyBucket.updateConfig lopBucket $ \(uninitializedLevel, _oldConfig) ->
+        let disabled = (uninitializedLevel, LeakyBucket.dummyConfig)
+        in
+        case csBucketConfig of
+          ChainSyncLoPBucketDisabled     -> disabled
+          ChainSyncLoPBucketEnabled csbc ->
+            let config  = mkLopConfig csbc
+                enabled = (LeakyBucket.capacity config, config)
+            in
+            case gsmState of
+              PreSyncing -> enabled
+              Syncing    -> enabled
+              CaughtUp   -> disabled
 
 -- Our task: after connecting to an upstream node, try to maintain an
 -- up-to-date header-only fragment representing their chain. We maintain
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ChainSync/Client/Jumping.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ChainSync/Client/Jumping.hs
@@ -191,6 +191,7 @@ module Ouroboros.Consensus.MiniProtocol.ChainSync.Client.Jumping (
   , mkJumping
   , noJumping
   , registerClient
+  , reregisterClient
   , rotateDynamo
   , unregisterClient
   ) where
@@ -799,6 +800,30 @@ registerClient context peer csState mkHandle = do
   cschcAddHandle (handlesCol context) peer handle
   pure (context {peer, handle}, mbEv)
 
+-- | Same as 'registerClient', but for a peer that had been previously disengaged.
+--
+-- This is specifically used for when the GSM transitions from
+-- 'Ouroboros.Consensus.Node.GsmState.CaughtUp' back to
+-- 'Ouroboros.Consensus.Node.GsmState.PreSyncing'.
+reregisterClient ::
+  ( LedgerSupportsProtocol blk,
+    IOLike m
+  ) =>
+  PeerContext m peer blk ->
+  STM m (Maybe (TraceEventCsj peer blk))
+reregisterClient peerContext = do
+  let h = handle peerContext
+  (csjState, mbEv) <- getDynamo (handlesCol peerContext) >>= \case
+    Nothing -> do
+      fragment <- csCandidate <$> readTVar (cschState h)
+      pure (Dynamo DynamoStarted $ pointSlot $ AF.anchorPoint fragment, Just InitializedAsDynamo)
+    Just (_, dynHandle) -> do
+      mJustInfo <- readTVar (cschJumpInfo dynHandle)
+      (\x -> (x, Nothing)) <$> newJumper mJustInfo (Happy FreshJumper Nothing)
+  writeTVar (cschJumping h) csjState
+  cschcAddHandle (handlesCol peerContext) (peer peerContext) h
+  pure mbEv
+
 -- | Unregister a client from a 'PeerContext'; this might trigger the election
 -- of a new dynamo or objector if the peer was one of these two.
 unregisterClient ::
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ChainSync/Client/State.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/MiniProtocol/ChainSync/Client/State.hs
@@ -35,7 +35,7 @@ import           Ouroboros.Consensus.HeaderStateHistory (HeaderStateHistory)
 import           Ouroboros.Consensus.HeaderValidation (HeaderWithTime (..))
 import           Ouroboros.Consensus.Ledger.SupportsProtocol
                      (LedgerSupportsProtocol)
-import           Ouroboros.Consensus.Node.GsmState (GsmState)
+import           Ouroboros.Consensus.Node.GsmState (GsmTransition)
 import           Ouroboros.Consensus.Util.IOLike (IOLike, NoThunks (..), STM,
                      StrictTVar, Time, modifyTVar, newTVar, readTVar)
 import           Ouroboros.Network.AnchoredFragment (AnchoredFragment,
@@ -81,7 +81,7 @@ data ChainSyncClientHandle m blk = ChainSyncClientHandle {
 
     -- | Callback called by the GSM when the GSM state changes. They take the
     -- current time and should execute rapidly. Used to enable/disable the LoP.
-  , cschOnGsmStateChanged :: !(GsmState -> Time -> STM m ())
+  , cschOnGsmStateChanged :: !(GsmTransition -> Time -> STM m (m ()))
 
     -- | Data shared between the client and external components like GDD.
   , cschState             :: !(StrictTVar m (ChainSyncState blk))
diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Node/GsmState.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Node/GsmState.hs
@@ -1,10 +1,15 @@
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE LambdaCase #-}
 
 -- | This module contains the definition of a state in the Genesis State Machine
 -- (GSM). The GSM itself is defined in 'ouroboros-consensus-diffusion', but the
 -- ChainSync client relies on its state.
-module Ouroboros.Consensus.Node.GsmState (GsmState (..)) where
+module Ouroboros.Consensus.Node.GsmState (
+    GsmState (..)
+  , GsmTransition (..)
+  , gsmNewState
+  ) where
 
 import           GHC.Generics (Generic)
 import           NoThunks.Class (NoThunks)
@@ -21,3 +26,20 @@ data GsmState =
     CaughtUp
     -- ^ We are caught-up.
   deriving (Eq, Show, Read, Generic, NoThunks)
+
+data GsmTransition =
+    PreSyncingSyncing
+  |
+    SyncingCaughtUp
+  |
+    CaughtUpPreSyncing
+  |
+    SyncingPreSyncing
+  deriving (Eq, Show, Read)
+
+gsmNewState :: GsmTransition -> GsmState
+gsmNewState = \case
+    PreSyncingSyncing  -> Syncing
+    SyncingCaughtUp    -> CaughtUp
+    CaughtUpPreSyncing -> PreSyncing
+    SyncingPreSyncing  -> PreSyncing

Original file line number	Diff line number	Diff line change
`@@ -221,8 +221,8 @@ prop_sequential1 ctx cmds = runSimQC $ do`
`221`	`221`	`atomically $ do`
`222`	`222`	`writeTVar varMarker $ if b then Present else Absent`
`223`	`223`	`,`
`224`		`- GSM.writeGsmState = \x -> atomically $ do`
`225`		`- writeTVar varGsmState x`
	`224`	`+ GSM.writeGsmState =`
	`225`	`+ atomically . writeTVar varGsmState . GSM.gsmNewState`
`226`	`226`	`,`
`227`	`227`	`GSM.isHaaSatisfied =`
`228`	`228`	`isHaaSatisfied . Map.keysSet <$> readTVar varStates`