{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Test.Consensus.PeerSimulator.NodeLifecycle
  ( LiveInterval (..)
  , LiveIntervalResult (..)
  , LiveNode (..)
  , LiveResources (..)
  , NodeLifecycle (..)
  , lifecycleStart
  , lifecycleStop
  , restoreNode
  ) where

import Control.ResourceRegistry
import Control.Tracer (Tracer (..), traceWith)
import Data.Functor (void)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Typeable (Typeable)
import Ouroboros.Consensus.Block
import Ouroboros.Consensus.Config (TopLevelConfig (..))
import Ouroboros.Consensus.HeaderValidation (HeaderWithTime (..))
import Ouroboros.Consensus.MiniProtocol.ChainSync.Client
  ( ChainSyncClientHandleCollection (..)
  )
import Ouroboros.Consensus.Storage.ChainDB.API
import qualified Ouroboros.Consensus.Storage.ChainDB.API as ChainDB
import qualified Ouroboros.Consensus.Storage.ChainDB.Impl as ChainDB
import Ouroboros.Consensus.Storage.ChainDB.Impl.Args
  ( cdbsLoE
  , updateTracer
  )
import Ouroboros.Consensus.Util.IOLike
import Ouroboros.Network.AnchoredFragment (AnchoredFragment)
import qualified Ouroboros.Network.AnchoredFragment as AF
import System.FS.Sim.MockFS (MockFS)
import qualified System.FS.Sim.MockFS as MockFS
import Test.Consensus.PeerSimulator.Resources
import Test.Consensus.PeerSimulator.StateView
import Test.Consensus.PeerSimulator.Trace
import Test.Consensus.PointSchedule.Peers (PeerId)
import Test.Util.ChainDB
import Test.Util.Orphans.IOLike ()
import Test.Util.TestBlock (TestBlock, testInitExtLedger)

-- | Resources used for a single live interval of the node, constructed when the
-- node is started.
-- When the node is shut down, 'lnCopyToImmDb' is used to persist the current
-- chain.
data LiveNode blk m = LiveNode
  { forall blk (m :: * -> *). LiveNode blk m -> ChainDB m blk
lnChainDb :: ChainDB m blk
  , forall blk (m :: * -> *). LiveNode blk m -> StateViewTracers blk m
lnStateViewTracers :: StateViewTracers blk m
  , forall blk (m :: * -> *). LiveNode blk m -> Tracer m ()
lnStateTracer :: Tracer m ()
  , forall blk (m :: * -> *). LiveNode blk m -> m (WithOrigin SlotNo)
lnCopyToImmDb :: m (WithOrigin SlotNo)
  -- ^ Write persistent ChainDB state (the immutable and volatile DBs, but not
  -- the ledger and GSM state) to the VFS TVars to preserve it for the next
  -- interval.
  -- Returns the immutable tip's slot for tracing.
  , forall blk (m :: * -> *). LiveNode blk m -> Set PeerId
lnPeers :: Set PeerId
  -- ^ The set of peers that should be started.
  -- Based on the simulation results at node shutdown, disconnected peers are
  -- removed for the next live interval.
  }

-- | Result of a node shutdown at the end of a live interval.
data LiveIntervalResult blk = LiveIntervalResult
  { forall blk. LiveIntervalResult blk -> [PeerSimulatorResult blk]
lirPeerResults :: [PeerSimulatorResult blk]
  -- ^ Used to initialize the 'StateViewTracers' of the next run to preserve
  -- earlier disconnections for the final result.
  , forall blk. LiveIntervalResult blk -> Set PeerId
lirActive :: Set PeerId
  -- ^ The remaining peers, computed by removing all peers present in
  -- 'lrPeerResults' from the current state in 'lnPeers'.
  }

-- | Resources used by the handlers 'lifecycleStart' and 'lifecycleStop' to
-- shut down running components, construct tracers used for single intervals,
-- and reset and persist state.
data LiveResources blk m = LiveResources
  { forall blk (m :: * -> *). LiveResources blk m -> ResourceRegistry m
lrRegistry :: ResourceRegistry m
  , forall blk (m :: * -> *).
LiveResources blk m -> PeerSimulatorResources m blk
lrPeerSim :: PeerSimulatorResources m blk
  , forall blk (m :: * -> *).
LiveResources blk m -> Tracer m (TraceEvent blk)
lrTracer :: Tracer m (TraceEvent blk)
  , forall blk (m :: * -> *).
LiveResources blk m -> ChainDB m blk -> m (Tracer m ())
lrSTracer :: ChainDB m blk -> m (Tracer m ())
  , forall blk (m :: * -> *). LiveResources blk m -> TopLevelConfig blk
lrConfig :: TopLevelConfig blk
  , forall blk (m :: * -> *).
LiveResources blk m -> NodeDBs (StrictTMVar m MockFS)
lrCdb :: NodeDBs (StrictTMVar m MockFS)
  -- ^ The chain DB state consists of several transient parts and the
  -- immutable DB's virtual file system.
  -- After 'lnCopyToImmDb' was executed, the latter will contain the final
  -- state of an interval.
  -- The rest is reset when the chain DB is recreated.
  , forall blk (m :: * -> *).
LiveResources blk m
-> LoE (StrictTVar m (AnchoredFragment (HeaderWithTime blk)))
lrLoEVar :: LoE (StrictTVar m (AnchoredFragment (HeaderWithTime blk)))
  -- ^ The LoE fragment must be reset for each live interval.
  }

data LiveInterval blk m = LiveInterval
  { forall blk (m :: * -> *). LiveInterval blk m -> LiveResources blk m
liResources :: LiveResources blk m
  , forall blk (m :: * -> *).
LiveInterval blk m -> LiveIntervalResult blk
liResult :: LiveIntervalResult blk
  , forall blk (m :: * -> *). LiveInterval blk m -> LiveNode blk m
liNode :: LiveNode blk m
  }

-- | Handlers for starting the node and shutting it down for each live interval,
-- using the state of the previous run.
data NodeLifecycle blk m = NodeLifecycle
  { forall blk (m :: * -> *). NodeLifecycle blk m -> Maybe DiffTime
nlMinDuration :: Maybe DiffTime
  -- ^ The minimum tick duration that triggers a node downtime.
  -- If this is 'Nothing', downtimes are disabled.
  , forall blk (m :: * -> *).
NodeLifecycle blk m -> LiveIntervalResult blk -> m (LiveNode blk m)
nlStart :: LiveIntervalResult blk -> m (LiveNode blk m)
  -- ^ Start the node with prior state.
  -- For the first start, this must be called with an empty 'lirPeerResults'
  -- and the initial set of all peers in 'lirActive'.
  , forall blk (m :: * -> *).
NodeLifecycle blk m -> LiveNode blk m -> m (LiveIntervalResult blk)
nlShutdown :: LiveNode blk m -> m (LiveIntervalResult blk)
  }

-- | Create a ChainDB and start a BlockRunner that operate on the peers'
-- candidate fragments.
mkChainDb ::
  IOLike m =>
  LiveResources TestBlock m ->
  m (ChainDB m TestBlock, m (WithOrigin SlotNo))
mkChainDb :: forall (m :: * -> *).
IOLike m =>
LiveResources TestBlock m
-> m (ChainDB m TestBlock, m (WithOrigin SlotNo))
mkChainDb LiveResources TestBlock m
resources = do
  STM m () -> m ()
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m () -> m ()) -> STM m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    -- Reset only the non-persisted state of the ChainDB's file system mocks:
    -- - GSM state and Ledger DB are discarded
    -- - Immutable DB and Volatile DB are preserved for the next interval
    STM m MockFS -> STM m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM m MockFS -> STM m ()) -> STM m MockFS -> STM m ()
forall a b. (a -> b) -> a -> b
$ StrictTMVar m MockFS -> MockFS -> STM m MockFS
forall (m :: * -> *) a.
MonadSTM m =>
StrictTMVar m a -> a -> STM m a
swapTMVar (NodeDBs (StrictTMVar m MockFS) -> StrictTMVar m MockFS
forall db. NodeDBs db -> db
nodeDBsGsm NodeDBs (StrictTMVar m MockFS)
lrCdb) MockFS
MockFS.empty
    STM m MockFS -> STM m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM m MockFS -> STM m ()) -> STM m MockFS -> STM m ()
forall a b. (a -> b) -> a -> b
$ StrictTMVar m MockFS -> MockFS -> STM m MockFS
forall (m :: * -> *) a.
MonadSTM m =>
StrictTMVar m a -> a -> STM m a
swapTMVar (NodeDBs (StrictTMVar m MockFS) -> StrictTMVar m MockFS
forall db. NodeDBs db -> db
nodeDBsLgr NodeDBs (StrictTMVar m MockFS)
lrCdb) MockFS
MockFS.empty
  chainDbArgs <- do
    let args :: ChainDbArgs Identity m TestBlock
args =
          Tracer m (TraceEvent TestBlock)
-> ChainDbArgs Identity m TestBlock
-> ChainDbArgs Identity m TestBlock
forall (m :: * -> *) blk (f :: * -> *).
Tracer m (TraceEvent blk)
-> ChainDbArgs f m blk -> ChainDbArgs f m blk
updateTracer
            ((TraceEvent TestBlock -> m ()) -> Tracer m (TraceEvent TestBlock)
forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer (Tracer m (TraceEvent TestBlock) -> TraceEvent TestBlock -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m (TraceEvent TestBlock)
lrTracer (TraceEvent TestBlock -> m ())
-> (TraceEvent TestBlock -> TraceEvent TestBlock)
-> TraceEvent TestBlock
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceEvent TestBlock -> TraceEvent TestBlock
forall blk. TraceEvent blk -> TraceEvent blk
TraceChainDBEvent))
            ( MinimalChainDbArgs m TestBlock -> ChainDbArgs Identity m TestBlock
forall (m :: * -> *) blk.
(MonadThrow m, MonadSTM m, ConsensusProtocol (BlockProtocol blk),
 PrimMonad m) =>
MinimalChainDbArgs m blk -> Complete ChainDbArgs m blk
fromMinimalChainDbArgs
                MinimalChainDbArgs
                  { mcdbTopLevelConfig :: TopLevelConfig TestBlock
mcdbTopLevelConfig = TopLevelConfig TestBlock
lrConfig
                  , mcdbChunkInfo :: ChunkInfo
mcdbChunkInfo = TopLevelConfig TestBlock -> ChunkInfo
mkTestChunkInfo TopLevelConfig TestBlock
lrConfig
                  , mcdbInitLedger :: ExtLedgerState TestBlock ValuesMK
mcdbInitLedger = ExtLedgerState TestBlock ValuesMK
testInitExtLedger
                  , mcdbRegistry :: ResourceRegistry m
mcdbRegistry = ResourceRegistry m
lrRegistry
                  , mcdbNodeDBs :: NodeDBs (StrictTMVar m MockFS)
mcdbNodeDBs = NodeDBs (StrictTMVar m MockFS)
lrCdb
                  }
            )
    ChainDbArgs Identity m TestBlock
-> m (ChainDbArgs Identity m TestBlock)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ChainDbArgs Identity m TestBlock
 -> m (ChainDbArgs Identity m TestBlock))
-> ChainDbArgs Identity m TestBlock
-> m (ChainDbArgs Identity m TestBlock)
forall a b. (a -> b) -> a -> b
$
      ChainDbArgs Identity m TestBlock
args
        { ChainDB.cdbsArgs =
            (ChainDB.cdbsArgs args)
              { cdbsLoE = traverse readTVarIO lrLoEVar
              }
        }
  (_, (chainDB, internal)) <-
    allocate
      lrRegistry
      (\ResourceId
_ -> ChainDbArgs Identity m TestBlock
-> Bool -> m (ChainDB m TestBlock, Internal m TestBlock)
forall (m :: * -> *) blk.
(IOLike m, LedgerSupportsProtocol blk,
 BlockSupportsDiffusionPipelining blk, InspectLedger blk,
 HasHardForkHistory blk, ConvertRawHash blk,
 SerialiseDiskConstraints blk, HasCallStack,
 LedgerSupportsLedgerDB blk) =>
Complete ChainDbArgs m blk
-> Bool -> m (ChainDB m blk, Internal m blk)
ChainDB.openDBInternal ChainDbArgs Identity m TestBlock
chainDbArgs Bool
False)
      (ChainDB.closeDB . fst)
  let ChainDB.Internal{intCopyToImmutableDB, intAddBlockRunner} = internal
  void $ forkLinkedThread lrRegistry "AddBlockRunner" (void intAddBlockRunner)
  pure (chainDB, intCopyToImmutableDB)
 where
  LiveResources{ResourceRegistry m
lrRegistry :: forall blk (m :: * -> *). LiveResources blk m -> ResourceRegistry m
lrRegistry :: ResourceRegistry m
lrRegistry, Tracer m (TraceEvent TestBlock)
lrTracer :: forall blk (m :: * -> *).
LiveResources blk m -> Tracer m (TraceEvent blk)
lrTracer :: Tracer m (TraceEvent TestBlock)
lrTracer, TopLevelConfig TestBlock
lrConfig :: forall blk (m :: * -> *). LiveResources blk m -> TopLevelConfig blk
lrConfig :: TopLevelConfig TestBlock
lrConfig, NodeDBs (StrictTMVar m MockFS)
lrCdb :: forall blk (m :: * -> *).
LiveResources blk m -> NodeDBs (StrictTMVar m MockFS)
lrCdb :: NodeDBs (StrictTMVar m MockFS)
lrCdb, LoE (StrictTVar m (AnchoredFragment (HeaderWithTime TestBlock)))
lrLoEVar :: forall blk (m :: * -> *).
LiveResources blk m
-> LoE (StrictTVar m (AnchoredFragment (HeaderWithTime blk)))
lrLoEVar :: LoE (StrictTVar m (AnchoredFragment (HeaderWithTime TestBlock)))
lrLoEVar} = LiveResources TestBlock m
resources

-- | Allocate all the resources that depend on the results of previous live
-- intervals, the ChainDB and its persisted state.
restoreNode ::
  IOLike m =>
  LiveResources TestBlock m ->
  LiveIntervalResult TestBlock ->
  m (LiveNode TestBlock m)
restoreNode :: forall (m :: * -> *).
IOLike m =>
LiveResources TestBlock m
-> LiveIntervalResult TestBlock -> m (LiveNode TestBlock m)
restoreNode LiveResources TestBlock m
resources LiveIntervalResult{[PeerSimulatorResult TestBlock]
lirPeerResults :: forall blk. LiveIntervalResult blk -> [PeerSimulatorResult blk]
lirPeerResults :: [PeerSimulatorResult TestBlock]
lirPeerResults, Set PeerId
lirActive :: forall blk. LiveIntervalResult blk -> Set PeerId
lirActive :: Set PeerId
lirActive} = do
  lnStateViewTracers <- [PeerSimulatorResult TestBlock] -> m (StateViewTracers TestBlock m)
forall (m :: * -> *) blk.
IOLike m =>
[PeerSimulatorResult blk] -> m (StateViewTracers blk m)
stateViewTracersWithInitial [PeerSimulatorResult TestBlock]
lirPeerResults
  (lnChainDb, lnCopyToImmDb) <- mkChainDb resources
  lnStateTracer <- lrSTracer resources lnChainDb
  pure
    LiveNode
      { lnChainDb
      , lnStateViewTracers
      , lnStateTracer
      , lnCopyToImmDb
      , lnPeers = lirActive
      }

-- | Allocate resources with 'restoreNode' and pass them to the callback that
-- starts the node's threads.
lifecycleStart ::
  forall m.
  IOLike m =>
  (LiveInterval TestBlock m -> m ()) ->
  LiveResources TestBlock m ->
  LiveIntervalResult TestBlock ->
  m (LiveNode TestBlock m)
lifecycleStart :: forall (m :: * -> *).
IOLike m =>
(LiveInterval TestBlock m -> m ())
-> LiveResources TestBlock m
-> LiveIntervalResult TestBlock
-> m (LiveNode TestBlock m)
lifecycleStart LiveInterval TestBlock m -> m ()
start LiveResources TestBlock m
liResources LiveIntervalResult TestBlock
liResult = do
  TraceEvent TestBlock -> m ()
trace (TraceSchedulerEvent TestBlock -> TraceEvent TestBlock
forall blk. TraceSchedulerEvent blk -> TraceEvent blk
TraceSchedulerEvent TraceSchedulerEvent TestBlock
forall blk. TraceSchedulerEvent blk
TraceNodeStartupStart)
  liNode <- LiveResources TestBlock m
-> LiveIntervalResult TestBlock -> m (LiveNode TestBlock m)
forall (m :: * -> *).
IOLike m =>
LiveResources TestBlock m
-> LiveIntervalResult TestBlock -> m (LiveNode TestBlock m)
restoreNode LiveResources TestBlock m
liResources LiveIntervalResult TestBlock
liResult
  start LiveInterval{liResources, liResult, liNode}
  chain <- atomically (ChainDB.getCurrentChain (lnChainDb liNode))
  trace (TraceSchedulerEvent (TraceNodeStartupComplete chain))
  pure liNode
 where
  trace :: TraceEvent TestBlock -> m ()
trace = Tracer m (TraceEvent TestBlock) -> TraceEvent TestBlock -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith (LiveResources TestBlock m -> Tracer m (TraceEvent TestBlock)
forall blk (m :: * -> *).
LiveResources blk m -> Tracer m (TraceEvent blk)
lrTracer LiveResources TestBlock m
liResources)

-- | Shut down the node by killing all its threads after extracting the
-- persistent state used to restart the node later.
lifecycleStop ::
  (IOLike m, GetHeader blk, Typeable blk) =>
  LiveResources blk m ->
  LiveNode blk m ->
  m (LiveIntervalResult blk)
lifecycleStop :: forall (m :: * -> *) blk.
(IOLike m, GetHeader blk, Typeable blk) =>
LiveResources blk m -> LiveNode blk m -> m (LiveIntervalResult blk)
lifecycleStop LiveResources blk m
resources LiveNode{StateViewTracers blk m
lnStateViewTracers :: forall blk (m :: * -> *). LiveNode blk m -> StateViewTracers blk m
lnStateViewTracers :: StateViewTracers blk m
lnStateViewTracers, m (WithOrigin SlotNo)
lnCopyToImmDb :: forall blk (m :: * -> *). LiveNode blk m -> m (WithOrigin SlotNo)
lnCopyToImmDb :: m (WithOrigin SlotNo)
lnCopyToImmDb, Set PeerId
lnPeers :: forall blk (m :: * -> *). LiveNode blk m -> Set PeerId
lnPeers :: Set PeerId
lnPeers} = do
  -- Trigger writing the immutable tip to the MockFS in our TVar for restoring in 'startNode'
  immutableTip <- m (WithOrigin SlotNo)
lnCopyToImmDb
  trace (TraceSchedulerEvent (TraceNodeShutdownStart immutableTip))
  -- Remember which peers were still running before shutdown
  lirPeerResults <- svtGetPeerSimulatorResults lnStateViewTracers
  let disconnectedPeers = [PeerId] -> Set PeerId
forall a. Ord a => [a] -> Set a
Set.fromList (PeerSimulatorResult blk -> PeerId
forall blk. PeerSimulatorResult blk -> PeerId
psePeerId (PeerSimulatorResult blk -> PeerId)
-> [PeerSimulatorResult blk] -> [PeerId]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [PeerSimulatorResult blk]
lirPeerResults)
      lirActive = Set PeerId
lnPeers Set PeerId -> Set PeerId -> Set PeerId
forall a. Ord a => Set a -> Set a -> Set a
Set.\\ Set PeerId
disconnectedPeers
  -- Killing the peer overview threads should hopefully clean up all connections promptly
  releaseAll lrRegistry
  -- Reset the resources in TVars that were allocated by the simulator
  atomically $ do
    cschcRemoveAllHandles psrHandles
    case lrLoEVar of
      LoEEnabled StrictTVar m (AnchoredFragment (HeaderWithTime blk))
var -> StrictTVar m (AnchoredFragment (HeaderWithTime blk))
-> (AnchoredFragment (HeaderWithTime blk)
    -> AnchoredFragment (HeaderWithTime blk))
-> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> (a -> a) -> STM m ()
modifyTVar StrictTVar m (AnchoredFragment (HeaderWithTime blk))
var (AnchoredFragment (HeaderWithTime blk)
-> AnchoredFragment (HeaderWithTime blk)
-> AnchoredFragment (HeaderWithTime blk)
forall a b. a -> b -> a
const (Anchor (HeaderWithTime blk)
-> AnchoredFragment (HeaderWithTime blk)
forall v a b. Anchorable v a b => a -> AnchoredSeq v a b
AF.Empty Anchor (HeaderWithTime blk)
forall block. Anchor block
AF.AnchorGenesis))
      LoE (StrictTVar m (AnchoredFragment (HeaderWithTime blk)))
LoEDisabled -> () -> STM m ()
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  trace (TraceSchedulerEvent TraceNodeShutdownComplete)
  pure LiveIntervalResult{lirActive, lirPeerResults}
 where
  trace :: TraceEvent blk -> m ()
trace = Tracer m (TraceEvent blk) -> TraceEvent blk -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m (TraceEvent blk)
lrTracer
  LiveResources
    { ResourceRegistry m
lrRegistry :: forall blk (m :: * -> *). LiveResources blk m -> ResourceRegistry m
lrRegistry :: ResourceRegistry m
lrRegistry
    , Tracer m (TraceEvent blk)
lrTracer :: forall blk (m :: * -> *).
LiveResources blk m -> Tracer m (TraceEvent blk)
lrTracer :: Tracer m (TraceEvent blk)
lrTracer
    , lrPeerSim :: forall blk (m :: * -> *).
LiveResources blk m -> PeerSimulatorResources m blk
lrPeerSim = PeerSimulatorResources{ChainSyncClientHandleCollection PeerId m TestBlock
psrHandles :: ChainSyncClientHandleCollection PeerId m TestBlock
psrHandles :: forall (m :: * -> *) blk.
PeerSimulatorResources m blk
-> ChainSyncClientHandleCollection PeerId m TestBlock
psrHandles}
    , LoE (StrictTVar m (AnchoredFragment (HeaderWithTime blk)))
lrLoEVar :: forall blk (m :: * -> *).
LiveResources blk m
-> LoE (StrictTVar m (AnchoredFragment (HeaderWithTime blk)))
lrLoEVar :: LoE (StrictTVar m (AnchoredFragment (HeaderWithTime blk)))
lrLoEVar
    } = LiveResources blk m
resources