{-# LANGUAGE NamedFieldPuns #-}

-- | This module contains code that is generic to any “scheduled server” (think
-- scheduled ChainSync or BlockFetch server). A scheduled server keeps track of
-- the current state of a point schedule and wakes up when new ticks arise. It
-- processes as many messages there are via its domain-specific handlers; once
-- there is nothing new to process, or what needs to process requires a
-- different state of the point schedule, the scheduled server goes back to
-- sleep, awaiting another tick.
module Test.Consensus.PeerSimulator.ScheduledServer (
    ScheduledServer (..)
  , awaitOnlineState
  , ensureCurrentState
  , runHandler
  , runHandlerWithTrace
  ) where

import           Control.Tracer (Tracer, traceWith)
import           Data.Foldable (traverse_)
import           Ouroboros.Consensus.Util.IOLike (IOLike,
                     MonadSTM (STM, atomically))
import           Test.Consensus.PeerSimulator.Trace
                     (TraceScheduledServerHandlerEvent (..))
import           Test.Consensus.PointSchedule.Peers (PeerId)

data ScheduledServer m state blk =
  ScheduledServer {
    forall (m :: * -> *) state blk.
ScheduledServer m state blk -> PeerId
ssPeerId       :: PeerId,
    forall (m :: * -> *) state blk.
ScheduledServer m state blk -> STM m (Maybe state)
ssCurrentState :: STM m (Maybe state),
    forall (m :: * -> *) state blk.
ScheduledServer m state blk -> STM m ()
ssTickStarted  :: STM m (),
    forall (m :: * -> *) state blk.
ScheduledServer m state blk
-> Tracer m (TraceScheduledServerHandlerEvent state blk)
ssCommonTracer :: Tracer m (TraceScheduledServerHandlerEvent state blk)
  }

nextTickState :: IOLike m => ScheduledServer m state blk -> m (Maybe state)
nextTickState :: forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> m (Maybe state)
nextTickState ScheduledServer {STM m (Maybe state)
ssCurrentState :: forall (m :: * -> *) state blk.
ScheduledServer m state blk -> STM m (Maybe state)
ssCurrentState :: STM m (Maybe state)
ssCurrentState, STM m ()
ssTickStarted :: forall (m :: * -> *) state blk.
ScheduledServer m state blk -> STM m ()
ssTickStarted :: STM m ()
ssTickStarted} =
  STM m (Maybe state) -> m (Maybe state)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m ()
ssTickStarted STM m () -> STM m (Maybe state) -> STM m (Maybe state)
forall a b. STM m a -> STM m b -> STM m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> STM m (Maybe state)
ssCurrentState)

retryOffline :: IOLike m => ScheduledServer m state blk -> Maybe state -> m state
retryOffline :: forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> Maybe state -> m state
retryOffline ScheduledServer m state blk
server = m state -> (state -> m state) -> Maybe state -> m state
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (ScheduledServer m state blk -> m state
forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> m state
awaitOnlineState ScheduledServer m state blk
server) state -> m state
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Block until the peer simulator has updated the concurrency primitive that
-- indicates that it's this peer's server's turn in the point schedule.
-- If the new state is 'Nothing', the point schedule has declared this peer as
-- offline for the current tick, so it will not resume operation and wait for
-- the next update.
awaitOnlineState :: IOLike m => ScheduledServer m state blk -> m state
awaitOnlineState :: forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> m state
awaitOnlineState ScheduledServer m state blk
server =
  ScheduledServer m state blk -> Maybe state -> m state
forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> Maybe state -> m state
retryOffline ScheduledServer m state blk
server (Maybe state -> m state) -> m (Maybe state) -> m state
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ScheduledServer m state blk -> m (Maybe state)
forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> m (Maybe state)
nextTickState ScheduledServer m state blk
server

-- | Fetch the current state from the STM action, and if it is 'Nothing',
-- wait for the next tick to be triggered in 'awaitOnlineState'.
--
-- Since processing of a tick always ends when the handler finishes
-- after serving the last point, this function is only relevant for the
-- initial state update.
ensureCurrentState :: IOLike m => ScheduledServer m state blk -> m state
ensureCurrentState :: forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> m state
ensureCurrentState ScheduledServer m state blk
server =
  ScheduledServer m state blk -> Maybe state -> m state
forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> Maybe state -> m state
retryOffline ScheduledServer m state blk
server (Maybe state -> m state) -> m (Maybe state) -> m state
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STM m (Maybe state) -> m (Maybe state)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (ScheduledServer m state blk -> STM m (Maybe state)
forall (m :: * -> *) state blk.
ScheduledServer m state blk -> STM m (Maybe state)
ssCurrentState ScheduledServer m state blk
server)

-- | Handler functions are STM actions for the usual race condition reasons,
-- which means that they cannot emit trace messages.
--
-- For that reason, we allow them to return their messages alongside the
-- protocol result and emit them here.
runHandlerWithTrace ::
  IOLike m =>
  Tracer m traceMsg ->
  STM m (state, [traceMsg]) ->
  m state
runHandlerWithTrace :: forall (m :: * -> *) traceMsg state.
IOLike m =>
Tracer m traceMsg -> STM m (state, [traceMsg]) -> m state
runHandlerWithTrace Tracer m traceMsg
tracer STM m (state, [traceMsg])
handler = do
  (state
result, [traceMsg]
handlerMessages) <- STM m (state, [traceMsg]) -> m (state, [traceMsg])
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically STM m (state, [traceMsg])
handler
  (traceMsg -> m ()) -> [traceMsg] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Tracer m traceMsg -> traceMsg -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m traceMsg
tracer) [traceMsg]
handlerMessages
  state -> m state
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure state
result

-- | Run a peer server's message handler by fetching state from the scheduler's STM interface.
--
-- The handler is an STM action that returns a protocol result and log messages.
--
-- If the result is 'Nothing', the server's activity for the current tick is complete
-- and we listen for the scheduler's signal to start the next tick, which we continue without
-- updating the protocol handler (in @restart@).
--
-- Otherwise, the result is passed to @dispatchMessage@, which produces a native protocol handler
-- message with the server's continuation in it.
runHandler ::
  IOLike m =>
  ScheduledServer m state blk ->
  String ->
  (state -> STM m (Maybe msg, [traceMsg])) ->
  Tracer m traceMsg ->
  (msg -> m h) ->
  m h
runHandler :: forall (m :: * -> *) state blk msg traceMsg h.
IOLike m =>
ScheduledServer m state blk
-> String
-> (state -> STM m (Maybe msg, [traceMsg]))
-> Tracer m traceMsg
-> (msg -> m h)
-> m h
runHandler server :: ScheduledServer m state blk
server@ScheduledServer{Tracer m (TraceScheduledServerHandlerEvent state blk)
ssCommonTracer :: forall (m :: * -> *) state blk.
ScheduledServer m state blk
-> Tracer m (TraceScheduledServerHandlerEvent state blk)
ssCommonTracer :: Tracer m (TraceScheduledServerHandlerEvent state blk)
ssCommonTracer} String
handlerName state -> STM m (Maybe msg, [traceMsg])
handler Tracer m traceMsg
handlerTracer msg -> m h
dispatchMessage =
  m h
run
  where
    run :: m h
run = do
      state
currentState <- ScheduledServer m state blk -> m state
forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> m state
ensureCurrentState ScheduledServer m state blk
server
      Tracer m (TraceScheduledServerHandlerEvent state blk)
-> TraceScheduledServerHandlerEvent state blk -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m (TraceScheduledServerHandlerEvent state blk)
ssCommonTracer (TraceScheduledServerHandlerEvent state blk -> m ())
-> TraceScheduledServerHandlerEvent state blk -> m ()
forall a b. (a -> b) -> a -> b
$ String -> state -> TraceScheduledServerHandlerEvent state blk
forall state blk.
String -> state -> TraceScheduledServerHandlerEvent state blk
TraceHandling String
handlerName state
currentState
      m h -> (msg -> m h) -> Maybe msg -> m h
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m h
restart msg -> m h
done (Maybe msg -> m h) -> m (Maybe msg) -> m h
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Tracer m traceMsg -> STM m (Maybe msg, [traceMsg]) -> m (Maybe msg)
forall (m :: * -> *) traceMsg state.
IOLike m =>
Tracer m traceMsg -> STM m (state, [traceMsg]) -> m state
runHandlerWithTrace Tracer m traceMsg
handlerTracer (state -> STM m (Maybe msg, [traceMsg])
handler state
currentState)

    restart :: m h
restart = do
      Tracer m (TraceScheduledServerHandlerEvent state blk)
-> TraceScheduledServerHandlerEvent state blk -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m (TraceScheduledServerHandlerEvent state blk)
ssCommonTracer (TraceScheduledServerHandlerEvent state blk -> m ())
-> TraceScheduledServerHandlerEvent state blk -> m ()
forall a b. (a -> b) -> a -> b
$ String -> TraceScheduledServerHandlerEvent state blk
forall state blk.
String -> TraceScheduledServerHandlerEvent state blk
TraceRestarting String
handlerName
      ScheduledServer m state blk -> m state
forall (m :: * -> *) state blk.
IOLike m =>
ScheduledServer m state blk -> m state
awaitOnlineState ScheduledServer m state blk
server m state -> m h -> m h
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m h
run

    done :: msg -> m h
done msg
msg = do
      Tracer m (TraceScheduledServerHandlerEvent state blk)
-> TraceScheduledServerHandlerEvent state blk -> m ()
forall (m :: * -> *) a. Tracer m a -> a -> m ()
traceWith Tracer m (TraceScheduledServerHandlerEvent state blk)
ssCommonTracer (TraceScheduledServerHandlerEvent state blk -> m ())
-> TraceScheduledServerHandlerEvent state blk -> m ()
forall a b. (a -> b) -> a -> b
$ String -> TraceScheduledServerHandlerEvent state blk
forall state blk.
String -> TraceScheduledServerHandlerEvent state blk
TraceDoneHandling String
handlerName
      msg -> m h
dispatchMessage msg
msg