{-# 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           Ouroboros.Consensus.Block
import           Ouroboros.Consensus.Config (TopLevelConfig (..))
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 qualified System.FS.Sim.MockFS as MockFS
import           System.FS.Sim.MockFS (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 ()

    -- | 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 -> m (WithOrigin SlotNo)
lnCopyToImmDb      :: m (WithOrigin SlotNo)

    -- | 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.
  , forall blk (m :: * -> *). LiveNode blk m -> Set PeerId
lnPeers            :: Set PeerId
  }

-- | Result of a node shutdown at the end of a live interval.
data LiveIntervalResult blk = LiveIntervalResult {
    -- | Used to initialize the 'StateViewTracers' of the next run to preserve
    -- earlier disconnections for the final result.
    forall blk. LiveIntervalResult blk -> [PeerSimulatorResult blk]
lirPeerResults :: [PeerSimulatorResult blk]

    -- | The remaining peers, computed by removing all peers present in
    -- 'lrPeerResults' from the current state in 'lnPeers'.
  , forall blk. LiveIntervalResult blk -> Set PeerId
lirActive      :: Set PeerId
  }

-- | 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

    -- | 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 -> NodeDBs (StrictTMVar m MockFS)
lrCdb      :: NodeDBs (StrictTMVar m MockFS)

    -- | The LoE fragment must be reset for each live interval.
  , forall blk (m :: * -> *).
LiveResources blk m
-> LoE (StrictTVar m (AnchoredFragment (Header blk)))
lrLoEVar   :: LoE (StrictTVar m (AnchoredFragment (Header blk)))
  }

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 {
    -- | The minimum tick duration that triggers a node downtime.
    -- If this is 'Nothing', downtimes are disabled.
    forall blk (m :: * -> *). NodeLifecycle blk m -> Maybe DiffTime
nlMinDuration :: Maybe DiffTime

    -- | 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 -> LiveIntervalResult blk -> m (LiveNode blk m)
nlStart       :: LiveIntervalResult blk -> m (LiveNode blk m)
  , 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 (Header TestBlock)))
lrLoEVar :: forall blk (m :: * -> *).
LiveResources blk m
-> LoE (StrictTVar m (AnchoredFragment (Header blk)))
lrLoEVar :: LoE (StrictTVar m (AnchoredFragment (Header 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) =>
  LiveResources blk m ->
  LiveNode blk m ->
  m (LiveIntervalResult blk)
lifecycleStop :: forall (m :: * -> *) blk.
(IOLike m, GetHeader 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 (Header blk))
var -> StrictTVar m (AnchoredFragment (Header blk))
-> (AnchoredFragment (Header blk) -> AnchoredFragment (Header blk))
-> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> (a -> a) -> STM m ()
modifyTVar StrictTVar m (AnchoredFragment (Header blk))
var (AnchoredFragment (Header blk)
-> AnchoredFragment (Header blk) -> AnchoredFragment (Header blk)
forall a b. a -> b -> a
const (Anchor (Header blk) -> AnchoredFragment (Header blk)
forall v a b. Anchorable v a b => a -> AnchoredSeq v a b
AF.Empty Anchor (Header blk)
forall block. Anchor block
AF.AnchorGenesis))
      LoE (StrictTVar m (AnchoredFragment (Header 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 (Header blk)))
lrLoEVar :: forall blk (m :: * -> *).
LiveResources blk m
-> LoE (StrictTVar m (AnchoredFragment (Header blk)))
lrLoEVar :: LoE (StrictTVar m (AnchoredFragment (Header blk)))
lrLoEVar
      } = LiveResources blk m
resources