{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Ouroboros.Consensus.Util.STM
  ( -- * 'Watcher'
    Watcher (..)
  , forkLinkedWatcher
  , forkLinkedWatcherAllocate
  , forkLinkedWatcherFinalize
  , withWatcher

    -- * Misc
  , Fingerprint (..)
  , WithFingerprint (..)
  , blockUntilAllJust
  , blockUntilChanged
  , blockUntilJust
  , runWhenJust

    -- * Simulate various monad stacks in STM
  , Sim (..)
  , simId
  , simStateT

    -- * withTMVar
  , withTMVar
  , withTMVarAnd
  ) where

import Control.Monad (void)
import Control.Monad.State (StateT (..))
import Control.ResourceRegistry
import Data.Void
import GHC.Stack
import Ouroboros.Consensus.Util.IOLike
import Ouroboros.Network.BlockFetch.ConsensusInterface
  ( Fingerprint (..)
  , WithFingerprint (..)
  )

{-------------------------------------------------------------------------------
  Misc
-------------------------------------------------------------------------------}

-- | Wait until the TVar changed
blockUntilChanged ::
  forall m a b.
  (MonadSTM m, Eq b) =>
  (a -> b) -> b -> STM m a -> STM m (a, b)
blockUntilChanged :: forall (m :: * -> *) a b.
(MonadSTM m, Eq b) =>
(a -> b) -> b -> STM m a -> STM m (a, b)
blockUntilChanged a -> b
f b
b STM m a
getA = do
  a <- STM m a
getA
  let b' = a -> b
f a
a
  if b' == b
    then retry
    else return (a, b')

-- | Spawn a new thread that waits for an STM value to become 'Just'
--
-- The thread will be linked to the registry.
runWhenJust ::
  IOLike m =>
  ResourceRegistry m ->
  -- | Label for the thread
  String ->
  STM m (Maybe a) ->
  (a -> m ()) ->
  m ()
runWhenJust :: forall (m :: * -> *) a.
IOLike m =>
ResourceRegistry m
-> String -> STM m (Maybe a) -> (a -> m ()) -> m ()
runWhenJust ResourceRegistry m
registry String
label STM m (Maybe a)
getMaybeA a -> m ()
action =
  m (Thread m ()) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Thread m ()) -> m ()) -> m (Thread m ()) -> m ()
forall a b. (a -> b) -> a -> b
$
    ResourceRegistry m -> String -> m () -> m (Thread m ())
forall (m :: * -> *) a.
(MonadAsync m, MonadFork m, MonadMask m, HasCallStack) =>
ResourceRegistry m -> String -> m a -> m (Thread m a)
forkLinkedThread ResourceRegistry m
registry String
label (m () -> m (Thread m ())) -> m () -> m (Thread m ())
forall a b. (a -> b) -> a -> b
$
      a -> m ()
action (a -> m ()) -> m a -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STM m a -> m a
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (Maybe a) -> STM m a
forall (m :: * -> *) a. MonadSTM m => STM m (Maybe a) -> STM m a
blockUntilJust STM m (Maybe a)
getMaybeA)

blockUntilJust :: MonadSTM m => STM m (Maybe a) -> STM m a
blockUntilJust :: forall (m :: * -> *) a. MonadSTM m => STM m (Maybe a) -> STM m a
blockUntilJust STM m (Maybe a)
getMaybeA = do
  ma <- STM m (Maybe a)
getMaybeA
  case ma of
    Maybe a
Nothing -> STM m a
forall a. STM m a
forall (m :: * -> *) a. MonadSTM m => STM m a
retry
    Just a
a -> a -> STM m a
forall a. a -> STM m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

blockUntilAllJust :: MonadSTM m => [STM m (Maybe a)] -> STM m [a]
blockUntilAllJust :: forall (m :: * -> *) a.
MonadSTM m =>
[STM m (Maybe a)] -> STM m [a]
blockUntilAllJust = (STM m (Maybe a) -> STM m a) -> [STM m (Maybe a)] -> STM m [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM STM m (Maybe a) -> STM m a
forall (m :: * -> *) a. MonadSTM m => STM m (Maybe a) -> STM m a
blockUntilJust

{-------------------------------------------------------------------------------
  Simulate monad stacks
-------------------------------------------------------------------------------}

newtype Sim n m = Sim {forall (n :: * -> *) (m :: * -> *).
Sim n m -> forall a. n a -> STM m a
runSim :: forall a. n a -> STM m a}

simId :: Sim (STM m) m
simId :: forall (m :: * -> *). Sim (STM m) m
simId = (forall a. STM m a -> STM m a) -> Sim (STM m) m
forall (n :: * -> *) (m :: * -> *).
(forall a. n a -> STM m a) -> Sim n m
Sim STM m a -> STM m a
forall a. a -> a
forall a. STM m a -> STM m a
id

simStateT :: IOLike m => StrictTVar m st -> Sim n m -> Sim (StateT st n) m
simStateT :: forall (m :: * -> *) st (n :: * -> *).
IOLike m =>
StrictTVar m st -> Sim n m -> Sim (StateT st n) m
simStateT StrictTVar m st
stVar (Sim forall a. n a -> STM m a
k) = (forall a. StateT st n a -> STM m a) -> Sim (StateT st n) m
forall (n :: * -> *) (m :: * -> *).
(forall a. n a -> STM m a) -> Sim n m
Sim ((forall a. StateT st n a -> STM m a) -> Sim (StateT st n) m)
-> (forall a. StateT st n a -> STM m a) -> Sim (StateT st n) m
forall a b. (a -> b) -> a -> b
$ \(StateT st -> n (a, st)
f) -> do
  st <- StrictTVar m st -> STM m st
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar StrictTVar m st
stVar
  (a, st') <- k (f st)
  writeTVar stVar st'
  return a

{-------------------------------------------------------------------------------
  Watchers
-------------------------------------------------------------------------------}

-- | Specification for a thread that watches a variable, and reports interesting
-- changes.
--
-- NOTE: STM does not guarantee that 'wNotify' will /literally/ be called on
-- /every/ change: when the system is under heavy load, some updates may be
-- missed.
data Watcher m a fp = Watcher
  { forall (m :: * -> *) a fp. Watcher m a fp -> a -> fp
wFingerprint :: a -> fp
  -- ^ Obtain a fingerprint from a value of the monitored variable.
  , forall (m :: * -> *) a fp. Watcher m a fp -> Maybe fp
wInitial :: Maybe fp
  -- ^ The initial fingerprint
  --
  -- If 'Nothing', the action is executed once immediately to obtain the
  -- initial fingerprint.
  , forall (m :: * -> *) a fp. Watcher m a fp -> a -> m ()
wNotify :: a -> m ()
  -- ^ An action executed each time the fingerprint changes.
  , forall (m :: * -> *) a fp. Watcher m a fp -> STM m a
wReader :: STM m a
  -- ^ The variable to monitor.
  }

-- | Execute a 'Watcher'
--
-- NOT EXPORTED
runWatcher ::
  forall m a fp.
  (IOLike m, Eq fp, HasCallStack) =>
  Watcher m a fp ->
  m Void
runWatcher :: forall (m :: * -> *) a fp.
(IOLike m, Eq fp, HasCallStack) =>
Watcher m a fp -> m Void
runWatcher Watcher m a fp
watcher = do
  initB <- case Maybe fp
mbInitFP of
    Just fp
fp -> fp -> m fp
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return fp
fp
    Maybe fp
Nothing -> do
      a <- STM m a -> m a
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically STM m a
getA
      notify a
      return $ f a
  loop initB
 where
  Watcher
    { wFingerprint :: forall (m :: * -> *) a fp. Watcher m a fp -> a -> fp
wFingerprint = a -> fp
f
    , wInitial :: forall (m :: * -> *) a fp. Watcher m a fp -> Maybe fp
wInitial = Maybe fp
mbInitFP
    , wNotify :: forall (m :: * -> *) a fp. Watcher m a fp -> a -> m ()
wNotify = a -> m ()
notify
    , wReader :: forall (m :: * -> *) a fp. Watcher m a fp -> STM m a
wReader = STM m a
getA
    } = Watcher m a fp
watcher

  loop :: fp -> m Void
  loop :: fp -> m Void
loop fp
fp = do
    (a, fp') <- STM m (a, fp) -> m (a, fp)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (a, fp) -> m (a, fp)) -> STM m (a, fp) -> m (a, fp)
forall a b. (a -> b) -> a -> b
$ (a -> fp) -> fp -> STM m a -> STM m (a, fp)
forall (m :: * -> *) a b.
(MonadSTM m, Eq b) =>
(a -> b) -> b -> STM m a -> STM m (a, b)
blockUntilChanged a -> fp
f fp
fp STM m a
getA
    notify a
    loop fp'

-- | Spawn a new thread that runs a 'Watcher'
--
-- The thread will be linked to the registry.
forkLinkedWatcher ::
  forall m a fp.
  (IOLike m, Eq fp, HasCallStack) =>
  ResourceRegistry m ->
  -- | Label for the thread
  String ->
  Watcher m a fp ->
  m (Thread m Void)
forkLinkedWatcher :: forall (m :: * -> *) a fp.
(IOLike m, Eq fp, HasCallStack) =>
ResourceRegistry m -> String -> Watcher m a fp -> m (Thread m Void)
forkLinkedWatcher ResourceRegistry m
registry String
label Watcher m a fp
watcher =
  ResourceRegistry m -> String -> m Void -> m (Thread m Void)
forall (m :: * -> *) a.
(MonadAsync m, MonadFork m, MonadMask m, HasCallStack) =>
ResourceRegistry m -> String -> m a -> m (Thread m a)
forkLinkedThread ResourceRegistry m
registry String
label (m Void -> m (Thread m Void)) -> m Void -> m (Thread m Void)
forall a b. (a -> b) -> a -> b
$ Watcher m a fp -> m Void
forall (m :: * -> *) a fp.
(IOLike m, Eq fp, HasCallStack) =>
Watcher m a fp -> m Void
runWatcher Watcher m a fp
watcher

-- | Spawn a new thread that safely allocates a resource, then runs a 'Watcher'
-- using the resource, finalizing the resource safely when the watcher
-- terminates in any way.
--
-- The allocated resource and thread are both linked to the registry.
forkLinkedWatcherAllocate ::
  forall m r a fp.
  (IOLike m, Eq fp, HasCallStack) =>
  ResourceRegistry m ->
  -- | Label for the thread
  String ->
  m r ->
  (r -> m ()) ->
  (r -> Watcher m a fp) ->
  m (Thread m Void)
forkLinkedWatcherAllocate :: forall (m :: * -> *) r a fp.
(IOLike m, Eq fp, HasCallStack) =>
ResourceRegistry m
-> String
-> m r
-> (r -> m ())
-> (r -> Watcher m a fp)
-> m (Thread m Void)
forkLinkedWatcherAllocate ResourceRegistry m
registry String
label m r
allocater r -> m ()
finalizer r -> Watcher m a fp
f =
  ResourceRegistry m -> String -> m Void -> m (Thread m Void)
forall (m :: * -> *) a.
(MonadAsync m, MonadFork m, MonadMask m, HasCallStack) =>
ResourceRegistry m -> String -> m a -> m (Thread m a)
forkLinkedThread ResourceRegistry m
registry String
label (m Void -> m (Thread m Void)) -> m Void -> m (Thread m Void)
forall a b. (a -> b) -> a -> b
$ do
    (rk, r) <- ResourceRegistry m
-> (ResourceId -> m r) -> (r -> m ()) -> m (ResourceKey m, r)
forall (m :: * -> *) a.
(MonadSTM m, MonadMask m, MonadThread m, HasCallStack) =>
ResourceRegistry m
-> (ResourceId -> m a) -> (a -> m ()) -> m (ResourceKey m, a)
allocate ResourceRegistry m
registry (m r -> ResourceId -> m r
forall a b. a -> b -> a
const m r
allocater) r -> m ()
finalizer
    runWatcher (f r) `finally` release rk

-- | Spawn a new thread that runs a 'Watcher', executing a finalizer when the
-- thread terminates.
--
-- The thread will be linked to the registry.
forkLinkedWatcherFinalize ::
  forall m a fp.
  (IOLike m, Eq fp, HasCallStack) =>
  ResourceRegistry m ->
  -- | Label for the thread
  String ->
  Watcher m a fp ->
  m () ->
  m (Thread m Void)
forkLinkedWatcherFinalize :: forall (m :: * -> *) a fp.
(IOLike m, Eq fp, HasCallStack) =>
ResourceRegistry m
-> String -> Watcher m a fp -> m () -> m (Thread m Void)
forkLinkedWatcherFinalize ResourceRegistry m
registry String
label Watcher m a fp
watcher m ()
finalizer =
  ResourceRegistry m -> String -> m Void -> m (Thread m Void)
forall (m :: * -> *) a.
(MonadAsync m, MonadFork m, MonadMask m, HasCallStack) =>
ResourceRegistry m -> String -> m a -> m (Thread m a)
forkLinkedThread ResourceRegistry m
registry String
label (m Void -> m (Thread m Void)) -> m Void -> m (Thread m Void)
forall a b. (a -> b) -> a -> b
$ Watcher m a fp -> m Void
forall (m :: * -> *) a fp.
(IOLike m, Eq fp, HasCallStack) =>
Watcher m a fp -> m Void
runWatcher Watcher m a fp
watcher m Void -> m () -> m Void
forall a b. m a -> m b -> m a
forall (m :: * -> *) a b. MonadThrow m => m a -> m b -> m a
`finally` m ()
finalizer

-- | Spawn a new thread that runs a 'Watcher'
--
-- The thread is bracketed via 'withAsync' and 'link'ed.
--
-- We do not provide the 'Async' handle only because our anticipated use cases
-- don't need it.
withWatcher ::
  forall m a fp r.
  (IOLike m, Eq fp, HasCallStack) =>
  -- | Label for the thread
  String ->
  Watcher m a fp ->
  m r ->
  m r
withWatcher :: forall (m :: * -> *) a fp r.
(IOLike m, Eq fp, HasCallStack) =>
String -> Watcher m a fp -> m r -> m r
withWatcher String
label Watcher m a fp
watcher m r
k =
  m Void -> (Async m Void -> m r) -> m r
forall a b. m a -> (Async m a -> m b) -> m b
forall (m :: * -> *) a b.
MonadAsync m =>
m a -> (Async m a -> m b) -> m b
withAsync
    (do String -> m ()
forall (m :: * -> *). MonadThread m => String -> m ()
labelThisThread String
label; Watcher m a fp -> m Void
forall (m :: * -> *) a fp.
(IOLike m, Eq fp, HasCallStack) =>
Watcher m a fp -> m Void
runWatcher Watcher m a fp
watcher)
    (\Async m Void
h -> do Async m Void -> m ()
forall (m :: * -> *) a.
(MonadAsync m, MonadFork m, MonadMask m) =>
Async m a -> m ()
link Async m Void
h; m r
k)

{-------------------------------------------------------------------------------
  withTMVar
-------------------------------------------------------------------------------}

-- | Apply @f@ with the content of @tv@ as state, restoring the original value when an
-- exception occurs
withTMVar ::
  IOLike m =>
  StrictTMVar m a ->
  (a -> m (c, a)) ->
  m c
withTMVar :: forall (m :: * -> *) a c.
IOLike m =>
StrictTMVar m a -> (a -> m (c, a)) -> m c
withTMVar StrictTMVar m a
tv a -> m (c, a)
f = StrictTMVar m a -> (a -> STM m ()) -> (a -> () -> m (c, a)) -> m c
forall (m :: * -> *) a b c.
IOLike m =>
StrictTMVar m a -> (a -> STM m b) -> (a -> b -> m (c, a)) -> m c
withTMVarAnd StrictTMVar m a
tv (STM m () -> a -> STM m ()
forall a b. a -> b -> a
const (STM m () -> a -> STM m ()) -> STM m () -> a -> STM m ()
forall a b. (a -> b) -> a -> b
$ () -> STM m ()
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (\a
a -> m (c, a) -> () -> m (c, a)
forall a b. a -> b -> a
const (m (c, a) -> () -> m (c, a)) -> m (c, a) -> () -> m (c, a)
forall a b. (a -> b) -> a -> b
$ a -> m (c, a)
f a
a)

-- | Apply @f@ with the content of @tv@ as state, restoring the original value
-- when an exception occurs. Additionally run a @STM@ action when acquiring the
-- value.
withTMVarAnd ::
  IOLike m =>
  StrictTMVar m a ->
  -- | Additional STM action to run in the same atomically
  -- block as the TMVar is acquired
  (a -> STM m b) ->
  -- | Action
  (a -> b -> m (c, a)) ->
  m c
withTMVarAnd :: forall (m :: * -> *) a b c.
IOLike m =>
StrictTMVar m a -> (a -> STM m b) -> (a -> b -> m (c, a)) -> m c
withTMVarAnd StrictTMVar m a
tv a -> STM m b
guard a -> b -> m (c, a)
f =
  (c, a) -> c
forall a b. (a, b) -> a
fst ((c, a) -> c) -> (((c, a), ()) -> (c, a)) -> ((c, a), ()) -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((c, a), ()) -> (c, a)
forall a b. (a, b) -> a
fst
    (((c, a), ()) -> c) -> m ((c, a), ()) -> m c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (a, b)
-> ((a, b) -> ExitCase (c, a) -> m ())
-> ((a, b) -> m (c, a))
-> m ((c, a), ())
forall a b c.
m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c)
forall (m :: * -> *) a b c.
MonadCatch m =>
m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c)
generalBracket
      ( STM m (a, b) -> m (a, b)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (a, b) -> m (a, b)) -> STM m (a, b) -> m (a, b)
forall a b. (a -> b) -> a -> b
$ do
          istate <- StrictTMVar m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => StrictTMVar m a -> STM m a
takeTMVar StrictTMVar m a
tv
          guarded <- guard istate
          pure (istate, guarded)
      )
      ( \(a
origState, b
_) -> \case
          ExitCaseSuccess (c
_, a
newState) ->
            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
$ StrictTMVar m a -> a -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTMVar m a -> a -> STM m ()
putTMVar StrictTMVar m a
tv a
newState
          ExitCaseException SomeException
_ ->
            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
$ StrictTMVar m a -> a -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTMVar m a -> a -> STM m ()
putTMVar StrictTMVar m a
tv a
origState
          ExitCase (c, a)
ExitCaseAbort ->
            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
$ StrictTMVar m a -> a -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTMVar m a -> a -> STM m ()
putTMVar StrictTMVar m a
tv a
origState
      )
      ((a -> b -> m (c, a)) -> (a, b) -> m (c, a)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> m (c, a)
f)