{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE NamedFieldPuns      #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Ouroboros.Network.TxSubmission.Inbound.V2.Registry
  ( SharedTxStateVar
  , PeerTxRegistry
  , PeerTxAPI (..)
  , TxSubmissionCountersVar
  , newSharedTxStateVar
  , newPeerTxRegistry
  , newTxSubmissionCountersVar
  , txCountersThreadV2
  , withPeer
  ) where

import Control.Concurrent.Class.MonadSTM qualified as Lazy
import Control.Concurrent.Class.MonadSTM.Strict
import Control.Monad (when)
import Control.Monad.Class.MonadThrow
import Control.Monad.Class.MonadTime.SI
import Control.Monad.Class.MonadTimer.SI
import Control.Tracer (Tracer, traceWith)
import Data.IntMap.Strict qualified as IntMap
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Void (Void)
import Data.Word (Word64)
import GHC.Stack (HasCallStack)

import Ouroboros.Network.Protocol.TxSubmission2.Type
import Ouroboros.Network.Tx (HasRawTxId)
import Ouroboros.Network.TxSubmission.Inbound.V2.Policy (TxDecisionPolicy (..))
import Ouroboros.Network.TxSubmission.Inbound.V2.State qualified as State
import Ouroboros.Network.TxSubmission.Inbound.V2.Types
import Ouroboros.Network.TxSubmission.Mempool.Reader

-- | Shared STM handle for V2 coordination state.
type SharedTxStateVar m peeraddr txid = StrictTVar m (SharedTxState peeraddr txid)

-- | STM handle for V2 monotonic counters.  Used both for each peer's
-- private counters cell and for the shared retired-peers accumulator
-- (see 'withPeer' and 'txCountersThreadV2').
type TxSubmissionCountersVar m = StrictTVar m TxSubmissionCounters

-- | Per-peer in-flight TVar.
type PeerTxInFlightVar m = StrictTVar m PeerTxInFlight

-- | Registry of every live peer's coordination TVars: its in-flight
-- contributions and its private counters cell.
--
-- 'withPeer' adds the pair on bracket-enter and removes it on
-- bracket-exit (after scrubbing any contributions the peer still has
-- to shared state and flushing its counters into the retired
-- accumulator).  'snapshotLiveReferences' reads the in-flight vars to
-- compute the @liveAdvertised@ union; 'snapshotCounters' reads the
-- counter cells to aggregate for emission.
type PeerTxRegistry m peeraddr =
       StrictTVar m (Map peeraddr (PeerTxInFlightVar m, TxSubmissionCountersVar m))

newSharedTxStateVar
  :: MonadSTM m
  => SharedTxState peeraddr txid
  -> m (SharedTxStateVar m peeraddr txid)
newSharedTxStateVar :: forall (m :: * -> *) peeraddr txid.
MonadSTM m =>
SharedTxState peeraddr txid -> m (SharedTxStateVar m peeraddr txid)
newSharedTxStateVar = SharedTxState peeraddr txid
-> m (StrictTVar m (SharedTxState peeraddr txid))
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
newTVarIO

newTxSubmissionCountersVar
  :: MonadSTM m
  => TxSubmissionCounters
  -> m (TxSubmissionCountersVar m)
newTxSubmissionCountersVar :: forall (m :: * -> *).
MonadSTM m =>
TxSubmissionCounters -> m (TxSubmissionCountersVar m)
newTxSubmissionCountersVar = TxSubmissionCounters -> m (StrictTVar m TxSubmissionCounters)
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
newTVarIO

newPeerTxRegistry
  :: MonadSTM m
  => m (PeerTxRegistry m peeraddr)
newPeerTxRegistry :: forall (m :: * -> *) peeraddr.
MonadSTM m =>
m (PeerTxRegistry m peeraddr)
newPeerTxRegistry = Map peeraddr (PeerTxInFlightVar m, TxSubmissionCountersVar m)
-> m (StrictTVar
        m (Map peeraddr (PeerTxInFlightVar m, TxSubmissionCountersVar m)))
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
newTVarIO Map peeraddr (PeerTxInFlightVar m, TxSubmissionCountersVar m)
forall k a. Map k a
Map.empty

-- | Central bookkeeping thread for V2.
--
-- Wakes every @'bufferedTxsMinLifetime' policy \/ 4@ seconds (capped
-- between 100 ms and 1 s) to run 'State.sweepSharedState' on the
-- shared tx state.  The sweep needs the union of every live peer's
-- 'pifAdvertised' to decide which entries are still wanted; the
-- registry is read inside the same STM transaction as the sweep so
-- the snapshot is coherent.  On a slower cadence (every
-- 'countersInterval' seconds of elapsed time) it also emits the
-- aggregated counters -- the retired-peers accumulator plus every
-- live peer's cell -- when they differ from the last emission, and a
-- 'TraceSharedTxState' snapshot when @sharedRevision@ or
-- @sharedGeneration@ has moved since the last emission.
txCountersThreadV2
  :: forall m peeraddr txid tx.
     (MonadDelay m, MonadSTM m, HasRawTxId txid)
  => TxDecisionPolicy
  -> Tracer m TxSubmissionCounters
  -> Tracer m (TraceTxLogic peeraddr txid tx)
  -> TxSubmissionCountersVar m
  -> SharedTxStateVar m peeraddr txid
  -> PeerTxRegistry m peeraddr
  -> m Void
txCountersThreadV2 :: forall (m :: * -> *) peeraddr txid tx.
(MonadDelay m, MonadSTM m, HasRawTxId txid) =>
TxDecisionPolicy
-> Tracer m TxSubmissionCounters
-> Tracer m (TraceTxLogic peeraddr txid tx)
-> TxSubmissionCountersVar m
-> SharedTxStateVar m peeraddr txid
-> PeerTxRegistry m peeraddr
-> m Void
txCountersThreadV2 TxDecisionPolicy
policy Tracer m TxSubmissionCounters
countersTracer Tracer m (TraceTxLogic peeraddr txid tx)
sharedStateTracer TxSubmissionCountersVar m
retiredCountersVar
                   SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxRegistry m peeraddr
registry = do
    now <- m Time
forall (m :: * -> *). MonadMonotonicTime m => m Time
getMonotonicTime
    initialSt <- readTVarIO sharedStateVar
    go mempty
       (sharedRevision   initialSt)
       (sharedGeneration initialSt)
       (addTime countersInterval now)
  where
    sweepInterval :: DiffTime
    sweepInterval :: DiffTime
sweepInterval = DiffTime -> DiffTime -> DiffTime
forall a. Ord a => a -> a -> a
max DiffTime
0.1 (DiffTime -> DiffTime -> DiffTime
forall a. Ord a => a -> a -> a
min DiffTime
1 (TxDecisionPolicy -> DiffTime
bufferedTxsMinLifetime TxDecisionPolicy
policy DiffTime -> DiffTime -> DiffTime
forall a. Fractional a => a -> a -> a
/ DiffTime
4))

    countersInterval :: DiffTime
    countersInterval :: DiffTime
countersInterval = DiffTime
7

    go :: TxSubmissionCounters -> Word64 -> Word64 -> Time -> m b
go !TxSubmissionCounters
previous !Word64
lastRev !Word64
lastGen !Time
nextEmitAt = do
      DiffTime -> m ()
forall (m :: * -> *). MonadDelay m => DiffTime -> m ()
threadDelay DiffTime
sweepInterval
      now <- m Time
forall (m :: * -> *). MonadMonotonicTime m => m Time
getMonotonicTime
      atomically $ do
        liveReferences <- snapshotLiveReferences registry
        st <- readTVar sharedStateVar
        let st' = Time
-> IntSet
-> SharedTxState peeraddr txid
-> SharedTxState peeraddr txid
forall txid peeraddr.
HasRawTxId txid =>
Time
-> IntSet
-> SharedTxState peeraddr txid
-> SharedTxState peeraddr txid
State.sweepSharedState Time
now IntSet
liveReferences SharedTxState peeraddr txid
st
        when (sharedRevision st' /= sharedRevision st) $
          writeTVar sharedStateVar st'
      if now >= nextEmitAt
         then do
           current <- atomically (snapshotCounters retiredCountersVar registry)
           when (current /= previous) $ traceWith countersTracer current
           st <- readTVarIO sharedStateVar
           let curRev = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedRevision   SharedTxState peeraddr txid
st
               curGen = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedGeneration SharedTxState peeraddr txid
st
           (lastRev', lastGen') <-
             if curRev /= lastRev || curGen /= lastGen
                then do
                  traceWith sharedStateTracer (TraceSharedTxState st)
                  pure (curRev, curGen)
                else pure (lastRev, lastGen)
           go current lastRev' lastGen' (addTime countersInterval now)
         else go previous lastRev lastGen nextEmitAt

-- | Read every live peer's 'pifAcksPending' and union them.  This is the
-- set of keys still referenced by some peer (advertised for fetch, or
-- held in @peerUnacknowledgedTxIds@ awaiting ack) and is used by the
-- sweep to know which lookup table entries can be safely reclaimed.
--
-- 'pifAcksPending' is a superset of 'pifAdvertised', so the advertised
-- set need not be unioned in separately.
snapshotLiveReferences
  :: MonadSTM m
  => PeerTxRegistry m peeraddr
  -> STM m IntSet
snapshotLiveReferences :: forall (m :: * -> *) peeraddr.
MonadSTM m =>
PeerTxRegistry m peeraddr -> STM m IntSet
snapshotLiveReferences PeerTxRegistry m peeraddr
registry = do
    peers <- PeerTxRegistry m peeraddr
-> STM
     m
     (Map
        peeraddr (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m))
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar PeerTxRegistry m peeraddr
registry
    pifs  <- traverse (readTVar . fst) (Map.elems peers)
    pure $! IntSet.unions (map pifAcksPending pifs)

-- | Aggregate the V2 counters for emission: the retired-peers
-- accumulator plus every live peer's private counters cell.
--
-- Read in a single STM transaction so the snapshot is coherent with
-- 'withPeer' disconnect flushes (which move a peer's counters into the
-- retired accumulator and drop it from the registry atomically),
-- keeping the emitted aggregate monotonic.  The transaction is
-- read-only, so the hot-path per-peer writers never block on it; only
-- this reader (at the 'countersInterval' cadence) retries on conflict.
snapshotCounters
  :: MonadSTM m
  => TxSubmissionCountersVar m
  -> PeerTxRegistry m peeraddr
  -> STM m TxSubmissionCounters
snapshotCounters :: forall (m :: * -> *) peeraddr.
MonadSTM m =>
TxSubmissionCountersVar m
-> PeerTxRegistry m peeraddr -> STM m TxSubmissionCounters
snapshotCounters TxSubmissionCountersVar m
retiredCountersVar PeerTxRegistry m peeraddr
registry = do
    retired <- TxSubmissionCountersVar m -> STM m TxSubmissionCounters
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar TxSubmissionCountersVar m
retiredCountersVar
    peers   <- readTVar registry
    live    <- traverse (readTVar . snd) (Map.elems peers)
    pure $! mconcat (retired : live)

-- | Peer-facing coordination API.
--
-- The peer thread keeps its local protocol state in a local
-- variable. Registry helpers operate only on the shared STM state
-- and the per-peer 'PeerTxInFlight' TVar (closure-captured); any
-- helper that needs peer-local state should take it explicitly as an
-- argument.
data PeerTxAPI m txid tx = PeerTxAPI {
    -- | Wait until either 'sharedGeneration' moves past the given
    -- value or the optional timeout expires.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx -> Word64 -> Maybe DiffTime -> m ()
awaitSharedChange    :: Word64
                         -> Maybe DiffTime
                         -> m (),

    -- | Compute the next action for this peer in non-pipelined mode.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
runNextPeerAction    :: Time
                         -> PeerTxLocalState tx
                         -> m (PeerAction, PeerTxLocalState tx),

    -- | Compute the next action for this peer in pipelined mode.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionPipelined :: Time
                               -> PeerTxLocalState tx
                               -> m (PeerAction, PeerTxLocalState tx),

    -- | Process a batch of txids received from this peer.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx
-> Time
-> NumTxIdsToReq
-> [(txid, SizeInBytes)]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applyReceivedTxIds   :: Time
                         -> NumTxIdsToReq
                         -> [(txid, SizeInBytes)]
                         -> PeerTxLocalState tx
                         -> m (PeerTxLocalState tx),

    -- | Process a batch of tx bodies received from this peer.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx
-> Time
-> [(txid, tx)]
-> PeerTxLocalState tx
-> m (Int, PeerTxLocalState tx)
applyReceivedTxs     :: Time
                         -> [(txid, tx)]
                         -> PeerTxLocalState tx
                         -> m (Int, PeerTxLocalState tx),

    -- | Mark txs as submitted to the mempool and update shared state.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx
-> Time
-> [TxKey]
-> [TxKey]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applySubmittedTxs    :: Time
                         -> [TxKey]
                         -> [TxKey]
                         -> PeerTxLocalState tx
                         -> m (PeerTxLocalState tx),

    -- | Resolve txids and advertised sizes for a batch of tx keys to request.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx
-> PeerTxLocalState tx -> [TxKey] -> m (Map txid SizeInBytes)
resolveTxRequest     :: PeerTxLocalState tx
                         -> [TxKey]
                         -> m (Map txid SizeInBytes),
    -- | Resolve buffered tx bodies into full submission records.
    --
    -- pre-condition: the list of keys is in `peerDownloadedTxs` field of
    -- `PeerTxLocalState`.  It is ensured by `pickSubmitAction`.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx
-> PeerTxLocalState tx -> [TxKey] -> m [(TxKey, txid, tx)]
resolveBufferedTxs   :: PeerTxLocalState tx
                         -> [TxKey]
                         -> m [(TxKey, txid, tx)],

    -- | Add a delta to the V2 monotonic counters.
    forall (m :: * -> *) txid tx.
PeerTxAPI m txid tx -> TxSubmissionCounters -> m ()
addCounters          :: TxSubmissionCounters -> m ()
  }

--
-- | A bracket function which registers / de-registers a new peer in
-- `SharedTxStateVar`,  which exposes `PeerTxStateAPI`.
-- `PeerTxStateAPI` is only safe inside the  `withPeer` scope.
--
withPeer
  :: forall peeraddr txid tx idx m a.
     ( MonadMask m
     , MonadTimer m
     , Ord peeraddr
     , Show peeraddr
     , Ord txid
     , HasRawTxId txid
     )
  => TxDecisionPolicy
  -> TxSubmissionMempoolReader txid tx idx m
  -> SharedTxStateVar m peeraddr txid
  -> PeerTxRegistry m peeraddr
  -> TxSubmissionCountersVar m
  -> peeraddr
  -> (PeerTxAPI m txid tx -> m a)
  -> m a
withPeer :: forall peeraddr txid tx idx (m :: * -> *) a.
(MonadMask m, MonadTimer m, Ord peeraddr, Show peeraddr, Ord txid,
 HasRawTxId txid) =>
TxDecisionPolicy
-> TxSubmissionMempoolReader txid tx idx m
-> SharedTxStateVar m peeraddr txid
-> PeerTxRegistry m peeraddr
-> TxSubmissionCountersVar m
-> peeraddr
-> (PeerTxAPI m txid tx -> m a)
-> m a
withPeer TxDecisionPolicy
policy TxSubmissionMempoolReader { STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot :: STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot :: forall txid tx idx (m :: * -> *).
TxSubmissionMempoolReader txid tx idx m
-> STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot }
         SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxRegistry m peeraddr
registry TxSubmissionCountersVar m
retiredCountersVar peeraddr
peeraddr PeerTxAPI m txid tx -> m a
io =
    m (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m)
-> ((StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m)
    -> m ())
-> ((StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m)
    -> m a)
-> m a
forall a b c. m a -> (a -> m b) -> (a -> m c) -> m c
forall (m :: * -> *) a b c.
MonadThrow m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket m (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m)
acquire (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m) -> m ()
release (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m) -> m a
run
  where
    acquire :: m (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m)
acquire = do
      peerInFlightVar <- PeerTxInFlight -> m (StrictTVar m PeerTxInFlight)
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
newTVarIO PeerTxInFlight
emptyPeerTxInFlight
      peerCountersVar <- newTVarIO mempty
      atomically $ modifyTVar registry
                     (Map.insert peeraddr (peerInFlightVar, peerCountersVar))
      pure (peerInFlightVar, peerCountersVar)

    -- On exit: reverse this peer's in-flight contributions to shared
    -- state, flush its accumulated counters into the retired
    -- accumulator (so the emitted aggregate stays monotonic across
    -- disconnects), then drop it from the registry.  Flush and removal
    -- commit in one transaction so 'snapshotCounters' counts this peer
    -- exactly once: via its live cell before exit, via the retired
    -- accumulator after.
    release :: (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m) -> m ()
release (StrictTVar m PeerTxInFlight
peerInFlightVar, TxSubmissionCountersVar m
peerCountersVar) = do
      now <- m Time
forall (m :: * -> *). MonadMonotonicTime m => m Time
getMonotonicTime
      atomically $ do
        pif <- readTVar peerInFlightVar
        modifyTVar sharedStateVar (scrubFromPeerInFlight peeraddr now pif)
        peerCounters <- readTVar peerCountersVar
        modifyTVar retiredCountersVar (<> peerCounters)
        modifyTVar registry (Map.delete peeraddr)

    run :: (StrictTVar m PeerTxInFlight, TxSubmissionCountersVar m) -> m a
run (StrictTVar m PeerTxInFlight
peerInFlightVar, TxSubmissionCountersVar m
peerCountersVar) = PeerTxAPI m txid tx -> m a
io PeerTxAPI {
          awaitSharedChange :: Word64 -> Maybe DiffTime -> m ()
awaitSharedChange = SharedTxStateVar m peeraddr txid
-> Word64 -> Maybe DiffTime -> m ()
forall (m :: * -> *) peeraddr txid.
MonadTimer m =>
SharedTxStateVar m peeraddr txid
-> Word64 -> Maybe DiffTime -> m ()
awaitSharedChangeImp SharedTxStateVar m peeraddr txid
sharedStateVar
        , runNextPeerAction :: Time -> PeerTxLocalState tx -> m (PeerAction, PeerTxLocalState tx)
runNextPeerAction = TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> StrictTVar m PeerTxInFlight
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
forall (m :: * -> *) peeraddr txid tx.
(MonadSTM m, Ord peeraddr) =>
TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionImp TxDecisionPolicy
policy SharedTxStateVar m peeraddr txid
sharedStateVar
                                StrictTVar m PeerTxInFlight
peerInFlightVar TxSubmissionCountersVar m
peerCountersVar peeraddr
peeraddr
        , runNextPeerActionPipelined :: Time -> PeerTxLocalState tx -> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionPipelined = TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> StrictTVar m PeerTxInFlight
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
forall (m :: * -> *) peeraddr txid tx.
(MonadSTM m, Ord peeraddr) =>
TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionPipelinedImp TxDecisionPolicy
policy
                                         SharedTxStateVar m peeraddr txid
sharedStateVar StrictTVar m PeerTxInFlight
peerInFlightVar
                                         TxSubmissionCountersVar m
peerCountersVar peeraddr
peeraddr
        , applyReceivedTxIds :: Time
-> NumTxIdsToReq
-> [(txid, SizeInBytes)]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applyReceivedTxIds = TxDecisionPolicy
-> STM m (MempoolSnapshot txid tx idx)
-> SharedTxStateVar m peeraddr txid
-> StrictTVar m PeerTxInFlight
-> TxSubmissionCountersVar m
-> Time
-> NumTxIdsToReq
-> [(txid, SizeInBytes)]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
forall (m :: * -> *) txid tx idx peeraddr.
(MonadSTM m, HasRawTxId txid) =>
TxDecisionPolicy
-> STM m (MempoolSnapshot txid tx idx)
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> Time
-> NumTxIdsToReq
-> [(txid, SizeInBytes)]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applyReceivedTxIdsImp TxDecisionPolicy
policy STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot
                                 SharedTxStateVar m peeraddr txid
sharedStateVar StrictTVar m PeerTxInFlight
peerInFlightVar TxSubmissionCountersVar m
peerCountersVar
        , applyReceivedTxs :: Time
-> [(txid, tx)]
-> PeerTxLocalState tx
-> m (Int, PeerTxLocalState tx)
applyReceivedTxs = TxDecisionPolicy
-> STM m (MempoolSnapshot txid tx idx)
-> SharedTxStateVar m peeraddr txid
-> StrictTVar m PeerTxInFlight
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> [(txid, tx)]
-> PeerTxLocalState tx
-> m (Int, PeerTxLocalState tx)
forall (m :: * -> *) peeraddr txid tx idx.
(MonadSTM m, Eq peeraddr, Show peeraddr, HasRawTxId txid) =>
TxDecisionPolicy
-> STM m (MempoolSnapshot txid tx idx)
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> [(txid, tx)]
-> PeerTxLocalState tx
-> m (Int, PeerTxLocalState tx)
applyReceivedTxsImp TxDecisionPolicy
policy STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot
                               SharedTxStateVar m peeraddr txid
sharedStateVar StrictTVar m PeerTxInFlight
peerInFlightVar TxSubmissionCountersVar m
peerCountersVar peeraddr
peeraddr
        , applySubmittedTxs :: Time
-> [TxKey]
-> [TxKey]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applySubmittedTxs = TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> StrictTVar m PeerTxInFlight
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> [TxKey]
-> [TxKey]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
forall (m :: * -> *) peeraddr txid tx.
(MonadSTM m, Eq peeraddr) =>
TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> [TxKey]
-> [TxKey]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applySubmittedTxsImp TxDecisionPolicy
policy SharedTxStateVar m peeraddr txid
sharedStateVar
                                StrictTVar m PeerTxInFlight
peerInFlightVar TxSubmissionCountersVar m
peerCountersVar peeraddr
peeraddr
        , resolveTxRequest :: PeerTxLocalState tx -> [TxKey] -> m (Map txid SizeInBytes)
resolveTxRequest = SharedTxStateVar m peeraddr txid
-> PeerTxLocalState tx -> [TxKey] -> m (Map txid SizeInBytes)
forall (m :: * -> *) txid peeraddr tx.
(HasCallStack, MonadSTM m, Ord txid) =>
SharedTxStateVar m peeraddr txid
-> PeerTxLocalState tx -> [TxKey] -> m (Map txid SizeInBytes)
resolveTxRequestImp SharedTxStateVar m peeraddr txid
sharedStateVar
        , resolveBufferedTxs :: PeerTxLocalState tx -> [TxKey] -> m [(TxKey, txid, tx)]
resolveBufferedTxs = SharedTxStateVar m peeraddr txid
-> PeerTxLocalState tx -> [TxKey] -> m [(TxKey, txid, tx)]
forall (m :: * -> *) peeraddr txid tx.
(HasCallStack, MonadSTM m) =>
SharedTxStateVar m peeraddr txid
-> PeerTxLocalState tx -> [TxKey] -> m [(TxKey, txid, tx)]
resolveBufferedTxsImp SharedTxStateVar m peeraddr txid
sharedStateVar
        , addCounters :: TxSubmissionCounters -> m ()
addCounters = \TxSubmissionCounters
delta -> 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
$ TxSubmissionCountersVar m
-> (TxSubmissionCounters -> TxSubmissionCounters) -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> (a -> a) -> STM m ()
modifyTVar TxSubmissionCountersVar m
peerCountersVar (TxSubmissionCounters
-> TxSubmissionCounters -> TxSubmissionCounters
forall a. Semigroup a => a -> a -> a
<> TxSubmissionCounters
delta)
        }

-- | Reverse this peer's still-outstanding contributions to the shared
-- 'TxEntry' counters.  Run by the bracket finalizer; uses the per-peer
-- TVar snapshot taken at exit time.
--
-- 'pifLeased' is best-effort (another peer can steal the lease in the
-- meantime), so the lease release verifies the entry still names this
-- peer as owner before claiming.
scrubFromPeerInFlight
  :: Eq peeraddr
  => peeraddr
  -> Time
  -> PeerTxInFlight
  -> SharedTxState peeraddr txid
  -> SharedTxState peeraddr txid
scrubFromPeerInFlight :: forall peeraddr txid.
Eq peeraddr =>
peeraddr
-> Time
-> PeerTxInFlight
-> SharedTxState peeraddr txid
-> SharedTxState peeraddr txid
scrubFromPeerInFlight peeraddr
peeraddr Time
now PeerTxInFlight
pif SharedTxState peeraddr txid
st
  | Bool
nothingToDo = SharedTxState peeraddr txid
st
  | Bool
otherwise   = SharedTxState peeraddr txid
st {
        sharedTxTable    = sharedTxTable',
        sharedGeneration = sharedGeneration st + 1
      }
  where
    -- 'pifAdvertised' and 'pifAcksPending' are not checked here: the
    -- bracket releaser also removes the peer from the registry, so the
    -- next sweep no longer sees this peer's keys in @liveReferences@ and
    -- reclaims the lookup tables itself.  Only the per-tx counters on
    -- 'sharedTxTable' need active scrubbing.
    nothingToDo :: Bool
nothingToDo =
         IntSet -> Bool
IntSet.null (PeerTxInFlight -> IntSet
pifLeased PeerTxInFlight
pif)
      Bool -> Bool -> Bool
&& IntSet -> Bool
IntSet.null (PeerTxInFlight -> IntSet
pifAttempting PeerTxInFlight
pif)
      Bool -> Bool -> Bool
&& IntSet -> Bool
IntSet.null (PeerTxInFlight -> IntSet
pifSubmitting PeerTxInFlight
pif)

    -- Walk each unique key once and apply all three modifications in a
    -- single 'IntMap.adjust'.  Keys typically appear in two of the
    -- three sets ('pifLeased' is a subset of 'pifAttempting ∪
    -- pifSubmitting').
    allKeys :: IntSet
allKeys = PeerTxInFlight -> IntSet
pifLeased PeerTxInFlight
pif
        IntSet -> IntSet -> IntSet
`IntSet.union` PeerTxInFlight -> IntSet
pifAttempting PeerTxInFlight
pif
        IntSet -> IntSet -> IntSet
`IntSet.union` PeerTxInFlight -> IntSet
pifSubmitting PeerTxInFlight
pif

    sharedTxTable' :: IntMap (TxEntry peeraddr)
sharedTxTable' = (IntMap (TxEntry peeraddr) -> Int -> IntMap (TxEntry peeraddr))
-> IntMap (TxEntry peeraddr) -> IntSet -> IntMap (TxEntry peeraddr)
forall a. (a -> Int -> a) -> a -> IntSet -> a
IntSet.foldl' IntMap (TxEntry peeraddr) -> Int -> IntMap (TxEntry peeraddr)
scrubOne (SharedTxState peeraddr txid -> IntMap (TxEntry peeraddr)
forall peeraddr txid.
SharedTxState peeraddr txid -> IntMap (TxEntry peeraddr)
sharedTxTable SharedTxState peeraddr txid
st) IntSet
allKeys

    scrubOne :: IntMap (TxEntry peeraddr) -> Int -> IntMap (TxEntry peeraddr)
scrubOne IntMap (TxEntry peeraddr)
tbl Int
k = (TxEntry peeraddr -> TxEntry peeraddr)
-> Int -> IntMap (TxEntry peeraddr) -> IntMap (TxEntry peeraddr)
forall a. (a -> a) -> Int -> IntMap a -> IntMap a
IntMap.adjust (Int -> TxEntry peeraddr -> TxEntry peeraddr
modifyEntry Int
k) Int
k IntMap (TxEntry peeraddr)
tbl

    modifyEntry :: Int -> TxEntry peeraddr -> TxEntry peeraddr
modifyEntry Int
k =
        Bool
-> (TxEntry peeraddr -> TxEntry peeraddr)
-> TxEntry peeraddr
-> TxEntry peeraddr
forall {a}. Bool -> (a -> a) -> a -> a
applyIf (Int -> IntSet -> Bool
IntSet.member Int
k (PeerTxInFlight -> IntSet
pifSubmitting PeerTxInFlight
pif)) TxEntry peeraddr -> TxEntry peeraddr
forall {peeraddr}. TxEntry peeraddr -> TxEntry peeraddr
clearSubmission
      (TxEntry peeraddr -> TxEntry peeraddr)
-> (TxEntry peeraddr -> TxEntry peeraddr)
-> TxEntry peeraddr
-> TxEntry peeraddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool
-> (TxEntry peeraddr -> TxEntry peeraddr)
-> TxEntry peeraddr
-> TxEntry peeraddr
forall {a}. Bool -> (a -> a) -> a -> a
applyIf (Int -> IntSet -> Bool
IntSet.member Int
k (PeerTxInFlight -> IntSet
pifAttempting PeerTxInFlight
pif)) TxEntry peeraddr -> TxEntry peeraddr
forall {peeraddr}. TxEntry peeraddr -> TxEntry peeraddr
decAttempt
      (TxEntry peeraddr -> TxEntry peeraddr)
-> (TxEntry peeraddr -> TxEntry peeraddr)
-> TxEntry peeraddr
-> TxEntry peeraddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool
-> (TxEntry peeraddr -> TxEntry peeraddr)
-> TxEntry peeraddr
-> TxEntry peeraddr
forall {a}. Bool -> (a -> a) -> a -> a
applyIf (Int -> IntSet -> Bool
IntSet.member Int
k (PeerTxInFlight -> IntSet
pifLeased     PeerTxInFlight
pif)) TxEntry peeraddr -> TxEntry peeraddr
releaseLease

    applyIf :: Bool -> (a -> a) -> a -> a
applyIf Bool
True  a -> a
f = a -> a
f
    applyIf Bool
False a -> a
_ = a -> a
forall a. a -> a
id

    releaseLease :: TxEntry peeraddr -> TxEntry peeraddr
releaseLease TxEntry peeraddr
entry = case TxEntry peeraddr -> TxLease peeraddr
forall peeraddr. TxEntry peeraddr -> TxLease peeraddr
txLease TxEntry peeraddr
entry of
      TxLeased peeraddr
owner Time
_ | peeraddr
owner peeraddr -> peeraddr -> Bool
forall a. Eq a => a -> a -> Bool
== peeraddr
peeraddr -> TxEntry peeraddr
entry { txLease = TxClaimable now }
      TxLease peeraddr
_                                    -> TxEntry peeraddr
entry

    decAttempt :: TxEntry peeraddr -> TxEntry peeraddr
decAttempt TxEntry peeraddr
entry = TxEntry peeraddr
entry { txAttempt = max 0 (txAttempt entry - 1) }

    clearSubmission :: TxEntry peeraddr -> TxEntry peeraddr
clearSubmission TxEntry peeraddr
entry = TxEntry peeraddr
entry { txInSubmission = False }

-- | Wait until either 'sharedGeneration' moves past the given value or the
-- optional timeout expires.
--
-- Used by idle peers to avoid busy-waiting while still being woken when
-- shared state changes (lease expiries, new tx advertisements, mempool
-- resolutions).  A spurious wake on a change that doesn't grant this peer
-- new work is harmless: the peer immediately re-runs 'nextPeerAction',
-- selects 'PeerDoNothing' again, and goes back to sleep on the new
-- generation value.
awaitSharedChangeImp :: MonadTimer m
                     => SharedTxStateVar m peeraddr txid
                     -> Word64
                     -> Maybe DiffTime
                     -> m ()
awaitSharedChangeImp :: forall (m :: * -> *) peeraddr txid.
MonadTimer m =>
SharedTxStateVar m peeraddr txid
-> Word64 -> Maybe DiffTime -> m ()
awaitSharedChangeImp SharedTxStateVar m peeraddr txid
sharedStateVar Word64
generation Maybe DiffTime
mDelay =
  case Maybe DiffTime
mDelay of
       Maybe DiffTime
Nothing ->
         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
           sharedState <- SharedTxStateVar m peeraddr txid
-> STM m (SharedTxState peeraddr txid)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar SharedTxStateVar m peeraddr txid
sharedStateVar
           check (sharedGeneration sharedState /= generation)
       Just DiffTime
delay -> do
         delayVar <- DiffTime -> m (TVar m Bool)
forall (m :: * -> *). MonadTimer m => DiffTime -> m (TVar m Bool)
registerDelay DiffTime
delay
         atomically $ do
           sharedState <- readTVar sharedStateVar
           expired <- Lazy.readTVar delayVar
           check (sharedGeneration sharedState /= generation || expired)

-- | Avoid rewriting the shared TVar when the pure state step made no shared
-- change. Callers use 'sharedGeneration' as the dirty bit for shared state.
writeSharedStateIfChanged :: MonadSTM m
                          => SharedTxStateVar m peeraddr txid
                          -> Word64                 -- ^ pre-step 'sharedGeneration'
                          -> Word64                 -- ^ pre-step 'sharedRevision'
                          -> SharedTxState peeraddr txid
                          -> STM m ()
writeSharedStateIfChanged :: forall (m :: * -> *) peeraddr txid.
MonadSTM m =>
SharedTxStateVar m peeraddr txid
-> Word64 -> Word64 -> SharedTxState peeraddr txid -> STM m ()
writeSharedStateIfChanged SharedTxStateVar m peeraddr txid
sharedStateVar Word64
sharedGeneration0 Word64
sharedRevision0 SharedTxState peeraddr txid
sharedState'
  | SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedGeneration SharedTxState peeraddr txid
sharedState' Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
sharedGeneration0
  , SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedRevision   SharedTxState peeraddr txid
sharedState' Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
sharedRevision0
  = () -> STM m ()
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise = SharedTxStateVar m peeraddr txid
-> SharedTxState peeraddr txid -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> a -> STM m ()
writeTVar SharedTxStateVar m peeraddr txid
sharedStateVar SharedTxState peeraddr txid
sharedState'

-- | Avoid rewriting the per-peer TVar when nothing changed.
writePeerInFlightIfChanged :: MonadSTM m
                           => PeerTxInFlightVar m
                           -> PeerTxInFlight
                           -> PeerTxInFlight
                           -> STM m ()
writePeerInFlightIfChanged :: forall (m :: * -> *).
MonadSTM m =>
PeerTxInFlightVar m -> PeerTxInFlight -> PeerTxInFlight -> STM m ()
writePeerInFlightIfChanged PeerTxInFlightVar m
var PeerTxInFlight
before PeerTxInFlight
after
  | PeerTxInFlight
before PeerTxInFlight -> PeerTxInFlight -> Bool
forall a. Eq a => a -> a -> Bool
== PeerTxInFlight
after = () -> STM m ()
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise       = PeerTxInFlightVar m -> PeerTxInFlight -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> a -> STM m ()
writeTVar PeerTxInFlightVar m
var PeerTxInFlight
after

-- | Update the counters for the action chosen by the peer scheduler.
--
updateCountersForAction :: MonadSTM m
                        => TxSubmissionCountersVar m
                        -> PeerAction
                        -> STM m ()
updateCountersForAction :: forall (m :: * -> *).
MonadSTM m =>
TxSubmissionCountersVar m -> PeerAction -> STM m ()
updateCountersForAction TxSubmissionCountersVar m
countersVar PeerAction
peerAction =
  case PeerAction
peerAction of
    PeerRequestTxIds NumTxIdsToAck
txIdsToAck NumTxIdsToReq
txIdsToReq
      | NumTxIdsToAck
txIdsToAck NumTxIdsToAck -> NumTxIdsToAck -> Bool
forall a. Eq a => a -> a -> Bool
/= NumTxIdsToAck
0 Bool -> Bool -> Bool
|| NumTxIdsToReq
txIdsToReq NumTxIdsToReq -> NumTxIdsToReq -> Bool
forall a. Eq a => a -> a -> Bool
/= NumTxIdsToReq
0 ->
          TxSubmissionCountersVar m
-> (TxSubmissionCounters -> TxSubmissionCounters) -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> (a -> a) -> STM m ()
modifyTVar TxSubmissionCountersVar m
countersVar (TxSubmissionCounters
-> TxSubmissionCounters -> TxSubmissionCounters
forall a. Semigroup a => a -> a -> a
<> TxSubmissionCounters
forall a. Monoid a => a
mempty
            { txIdMessagesSent = 1
            , txIdsRequested   = fromIntegral txIdsToReq
            })
    PeerRequestTxs [TxKey]
txKeys ->
      TxSubmissionCountersVar m
-> (TxSubmissionCounters -> TxSubmissionCounters) -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
StrictTVar m a -> (a -> a) -> STM m ()
modifyTVar TxSubmissionCountersVar m
countersVar (TxSubmissionCounters
-> TxSubmissionCounters -> TxSubmissionCounters
forall a. Semigroup a => a -> a -> a
<> TxSubmissionCounters
forall a. Monoid a => a
mempty { txMessagesSent = 1
                                        , txsRequested   = fromIntegral (length txKeys) })
    PeerAction
_ -> () -> STM m ()
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Compute the next action for this peer in non-pipelined mode.
runNextPeerActionImp :: ( MonadSTM m
                        , Ord peeraddr )
                     => TxDecisionPolicy
                     -> SharedTxStateVar m peeraddr txid
                     -> PeerTxInFlightVar m
                     -> TxSubmissionCountersVar m
                     -> peeraddr
                     -> Time
                     -> PeerTxLocalState tx
                     -> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionImp :: forall (m :: * -> *) peeraddr txid tx.
(MonadSTM m, Ord peeraddr) =>
TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionImp TxDecisionPolicy
policy SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxInFlightVar m
peerInFlightVar TxSubmissionCountersVar m
countersVar peeraddr
peeraddr
                     Time
now PeerTxLocalState tx
peerState = STM m (PeerAction, PeerTxLocalState tx)
-> m (PeerAction, PeerTxLocalState tx)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (PeerAction, PeerTxLocalState tx)
 -> m (PeerAction, PeerTxLocalState tx))
-> STM m (PeerAction, PeerTxLocalState tx)
-> m (PeerAction, PeerTxLocalState tx)
forall a b. (a -> b) -> a -> b
$ do
  sharedState <- SharedTxStateVar m peeraddr txid
-> STM m (SharedTxState peeraddr txid)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar SharedTxStateVar m peeraddr txid
sharedStateVar
  peerInFlight <- readTVar peerInFlightVar
  let sharedGeneration0 = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedGeneration SharedTxState peeraddr txid
sharedState
      sharedRevision0   = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedRevision   SharedTxState peeraddr txid
sharedState
      (peerAction, peerState', peerInFlight', sharedState') =
        State.nextPeerAction now policy peeraddr peerState peerInFlight sharedState
  writeSharedStateIfChanged sharedStateVar sharedGeneration0 sharedRevision0 sharedState'
  writePeerInFlightIfChanged peerInFlightVar peerInFlight peerInFlight'
  updateCountersForAction countersVar peerAction
  return (peerAction, peerState')

-- | Compute the next action for this peer in pipelined mode.
runNextPeerActionPipelinedImp :: ( MonadSTM m
                                  , Ord peeraddr )
                              => TxDecisionPolicy
                              -> SharedTxStateVar m peeraddr txid
                              -> PeerTxInFlightVar m
                              -> TxSubmissionCountersVar m
                              -> peeraddr
                              -> Time
                              -> PeerTxLocalState tx
                              -> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionPipelinedImp :: forall (m :: * -> *) peeraddr txid tx.
(MonadSTM m, Ord peeraddr) =>
TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> PeerTxLocalState tx
-> m (PeerAction, PeerTxLocalState tx)
runNextPeerActionPipelinedImp TxDecisionPolicy
policy SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxInFlightVar m
peerInFlightVar TxSubmissionCountersVar m
countersVar
                              peeraddr
peeraddr Time
now PeerTxLocalState tx
peerState =
    STM m (PeerAction, PeerTxLocalState tx)
-> m (PeerAction, PeerTxLocalState tx)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (PeerAction, PeerTxLocalState tx)
 -> m (PeerAction, PeerTxLocalState tx))
-> STM m (PeerAction, PeerTxLocalState tx)
-> m (PeerAction, PeerTxLocalState tx)
forall a b. (a -> b) -> a -> b
$ do
      sharedState <- SharedTxStateVar m peeraddr txid
-> STM m (SharedTxState peeraddr txid)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar SharedTxStateVar m peeraddr txid
sharedStateVar
      peerInFlight <- readTVar peerInFlightVar
      let sharedGeneration0 = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedGeneration SharedTxState peeraddr txid
sharedState
          sharedRevision0   = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedRevision   SharedTxState peeraddr txid
sharedState
          (peerAction, peerState', peerInFlight', sharedState') =
            State.nextPeerActionPipelined now policy peeraddr peerState
                                          peerInFlight sharedState
      writeSharedStateIfChanged sharedStateVar sharedGeneration0 sharedRevision0 sharedState'
      writePeerInFlightIfChanged peerInFlightVar peerInFlight peerInFlight'
      updateCountersForAction countersVar peerAction
      return (peerAction, peerState')

-- | Process a batch of txids received from this peer.
applyReceivedTxIdsImp :: ( MonadSTM m
                         , HasRawTxId txid )
                      => TxDecisionPolicy
                      -> STM m (MempoolSnapshot txid tx idx)
                      -> SharedTxStateVar m peeraddr txid
                      -> PeerTxInFlightVar m
                      -> TxSubmissionCountersVar m
                      -> Time
                      -> NumTxIdsToReq
                      -- ^ number of requested txid's
                      -> [(txid, SizeInBytes)]
                      -- ^ actually received txid's with sizes
                      -> PeerTxLocalState tx
                      -> m (PeerTxLocalState tx)
applyReceivedTxIdsImp :: forall (m :: * -> *) txid tx idx peeraddr.
(MonadSTM m, HasRawTxId txid) =>
TxDecisionPolicy
-> STM m (MempoolSnapshot txid tx idx)
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> Time
-> NumTxIdsToReq
-> [(txid, SizeInBytes)]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applyReceivedTxIdsImp TxDecisionPolicy
policy STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxInFlightVar m
peerInFlightVar
                      TxSubmissionCountersVar m
countersVar Time
now NumTxIdsToReq
txIdsToReq [(txid, SizeInBytes)]
txidsAndSizes PeerTxLocalState tx
peerState = do
  -- Snapshot the mempool outside the per-peer STM transaction so mempool
  -- writers don't kick the hot path into retries.  Stale answers are
  -- benign: a false positive delays re-fetch by 'bufferedTxsMinLifetime'
  -- via the retained set; a false negative wastes one body fetch that
  -- 'handleReceivedTxs' will reclassify as late.
  MempoolSnapshot { mempoolHasTx } <- STM m (MempoolSnapshot txid tx idx)
-> m (MempoolSnapshot txid tx idx)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot
  atomically $ do
    sharedState <- readTVar sharedStateVar
    peerInFlight <- readTVar peerInFlightVar
    let sharedGeneration0 = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedGeneration SharedTxState peeraddr txid
sharedState
        sharedRevision0   = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedRevision   SharedTxState peeraddr txid
sharedState
        (peerState', peerInFlight', sharedState') =
          State.handleReceivedTxIds mempoolHasTx now policy txIdsToReq txidsAndSizes
                                    peerState peerInFlight sharedState
    writeSharedStateIfChanged sharedStateVar sharedGeneration0 sharedRevision0 sharedState'
    writePeerInFlightIfChanged peerInFlightVar peerInFlight peerInFlight'
    modifyTVar countersVar (<> mempty { txIdRepliesReceived = 1
                                      , txIdsReceived       = fromIntegral (length txidsAndSizes) })
    return peerState'

-- | Process a batch of tx bodies received from this peer.
applyReceivedTxsImp :: ( MonadSTM m
                       , Eq peeraddr
                       , Show peeraddr
                       , HasRawTxId txid )
                    => TxDecisionPolicy
                    -> STM m (MempoolSnapshot txid tx idx)
                    -> SharedTxStateVar m peeraddr txid
                    -> PeerTxInFlightVar m
                    -> TxSubmissionCountersVar m
                    -> peeraddr
                    -> Time
                    -> [(txid, tx)]
                    -> PeerTxLocalState tx
                    -> m (Int, PeerTxLocalState tx)
applyReceivedTxsImp :: forall (m :: * -> *) peeraddr txid tx idx.
(MonadSTM m, Eq peeraddr, Show peeraddr, HasRawTxId txid) =>
TxDecisionPolicy
-> STM m (MempoolSnapshot txid tx idx)
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> [(txid, tx)]
-> PeerTxLocalState tx
-> m (Int, PeerTxLocalState tx)
applyReceivedTxsImp TxDecisionPolicy
policy STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxInFlightVar m
peerInFlightVar
                    TxSubmissionCountersVar m
countersVar peeraddr
peeraddr Time
now [(txid, tx)]
txs PeerTxLocalState tx
peerState = do
  -- Mempool snapshot taken in its own STM tx; see 'applyReceivedTxIdsImp'.
  MempoolSnapshot { mempoolHasTx } <- STM m (MempoolSnapshot txid tx idx)
-> m (MempoolSnapshot txid tx idx)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically STM m (MempoolSnapshot txid tx idx)
mempoolGetSnapshot
  atomically $ do
    sharedState <- readTVar sharedStateVar
    peerInFlight <- readTVar peerInFlightVar
    let sharedGeneration0 = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedGeneration SharedTxState peeraddr txid
sharedState
        sharedRevision0   = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedRevision   SharedTxState peeraddr txid
sharedState
        (omittedCount, lateCount, peerState', peerInFlight', sharedState') =
          State.handleReceivedTxs mempoolHasTx now policy peeraddr txs
                                  peerState peerInFlight sharedState
    writeSharedStateIfChanged sharedStateVar sharedGeneration0 sharedRevision0 sharedState'
    writePeerInFlightIfChanged peerInFlightVar peerInFlight peerInFlight'
    modifyTVar countersVar (<> mempty {
                        txRepliesReceived = 1,
                        txsReceived       = fromIntegral (length txs),
                        txsOmitted        = fromIntegral omittedCount,
                        lateBodies        = fromIntegral lateCount
                      })
    return (omittedCount + lateCount, peerState')

-- | Mark txs as submitted to the mempool and update shared state.
applySubmittedTxsImp :: ( MonadSTM m
                       , Eq peeraddr )
                     => TxDecisionPolicy
                     -> SharedTxStateVar m peeraddr txid
                     -> PeerTxInFlightVar m
                     -> TxSubmissionCountersVar m
                     -> peeraddr
                     -> Time
                     -> [TxKey]
                     -> [TxKey]
                     -> PeerTxLocalState tx
                     -> m (PeerTxLocalState tx)
applySubmittedTxsImp :: forall (m :: * -> *) peeraddr txid tx.
(MonadSTM m, Eq peeraddr) =>
TxDecisionPolicy
-> SharedTxStateVar m peeraddr txid
-> PeerTxInFlightVar m
-> TxSubmissionCountersVar m
-> peeraddr
-> Time
-> [TxKey]
-> [TxKey]
-> PeerTxLocalState tx
-> m (PeerTxLocalState tx)
applySubmittedTxsImp TxDecisionPolicy
policy SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxInFlightVar m
peerInFlightVar TxSubmissionCountersVar m
countersVar peeraddr
peeraddr
                     Time
now [TxKey]
acceptedTxs [TxKey]
rejectedTxs PeerTxLocalState tx
peerState =
  STM m (PeerTxLocalState tx) -> m (PeerTxLocalState tx)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (PeerTxLocalState tx) -> m (PeerTxLocalState tx))
-> STM m (PeerTxLocalState tx) -> m (PeerTxLocalState tx)
forall a b. (a -> b) -> a -> b
$ do
    sharedState <- SharedTxStateVar m peeraddr txid
-> STM m (SharedTxState peeraddr txid)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar SharedTxStateVar m peeraddr txid
sharedStateVar
    peerInFlight <- readTVar peerInFlightVar
    let sharedGeneration0 = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedGeneration SharedTxState peeraddr txid
sharedState
        sharedRevision0   = SharedTxState peeraddr txid -> Word64
forall peeraddr txid. SharedTxState peeraddr txid -> Word64
sharedRevision   SharedTxState peeraddr txid
sharedState
        (peerState', peerInFlight', sharedState') =
          State.handleSubmittedTxs now policy peeraddr acceptedTxs
                                   rejectedTxs peerState peerInFlight sharedState
    writeSharedStateIfChanged sharedStateVar sharedGeneration0 sharedRevision0 sharedState'
    writePeerInFlightIfChanged peerInFlightVar peerInFlight peerInFlight'
    modifyTVar countersVar (<> mempty { txsAccepted = fromIntegral (length acceptedTxs)
                                      , txsRejected = fromIntegral (length rejectedTxs) })
    return peerState'

-- | Resolve txids and advertised sizes for a batch of tx keys to request.
resolveTxRequestImp :: ( HasCallStack
                       , MonadSTM m
                       , Ord txid )
                    => SharedTxStateVar m peeraddr txid
                    -> PeerTxLocalState tx
                    -> [TxKey]
                    -> m (Map txid SizeInBytes)
resolveTxRequestImp :: forall (m :: * -> *) txid peeraddr tx.
(HasCallStack, MonadSTM m, Ord txid) =>
SharedTxStateVar m peeraddr txid
-> PeerTxLocalState tx -> [TxKey] -> m (Map txid SizeInBytes)
resolveTxRequestImp SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxLocalState tx
peerState [TxKey]
txKeys = do
  sharedState <- SharedTxStateVar m peeraddr txid -> m (SharedTxState peeraddr txid)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> m a
readTVarIO SharedTxStateVar m peeraddr txid
sharedStateVar
  return $ Map.fromList (fmap (resolveOne sharedState) txKeys)
  where
    resolveOne :: SharedTxState peeraddr a -> TxKey -> (a, SizeInBytes)
resolveOne SharedTxState peeraddr a
sharedState key :: TxKey
key@(TxKey Int
k) =
      ( SharedTxState peeraddr a -> TxKey -> a
forall peeraddr txid.
HasCallStack =>
SharedTxState peeraddr txid -> TxKey -> txid
resolveTxKey SharedTxState peeraddr a
sharedState TxKey
key
      , case Int -> IntMap SizeInBytes -> Maybe SizeInBytes
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k (PeerTxLocalState tx -> IntMap SizeInBytes
forall tx. PeerTxLocalState tx -> IntMap SizeInBytes
peerAvailableTxIds PeerTxLocalState tx
peerState) of
             Just SizeInBytes
txSize -> SizeInBytes
txSize
             Maybe SizeInBytes
Nothing     -> [Char] -> SizeInBytes
forall a. HasCallStack => [Char] -> a
error ([Char] -> SizeInBytes) -> [Char] -> SizeInBytes
forall a b. (a -> b) -> a -> b
$
               [Char]
"TxSubmission.V2.resolveTxRequestImp: missing tx size for key "
               [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
k
      )

-- | Resolve buffered tx bodies into full submission records.
resolveBufferedTxsImp :: ( HasCallStack
                         , MonadSTM m
                         )
                       => SharedTxStateVar m peeraddr txid
                       -> PeerTxLocalState tx
                       -> [TxKey]
                       -> m [(TxKey, txid, tx)]
resolveBufferedTxsImp :: forall (m :: * -> *) peeraddr txid tx.
(HasCallStack, MonadSTM m) =>
SharedTxStateVar m peeraddr txid
-> PeerTxLocalState tx -> [TxKey] -> m [(TxKey, txid, tx)]
resolveBufferedTxsImp SharedTxStateVar m peeraddr txid
sharedStateVar PeerTxLocalState tx
peerState [TxKey]
txKeys = do
  sharedState <- SharedTxStateVar m peeraddr txid -> m (SharedTxState peeraddr txid)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> m a
readTVarIO SharedTxStateVar m peeraddr txid
sharedStateVar
  return $ fmap (resolveOne sharedState) txKeys
  where
    resolveOne :: SharedTxState peeraddr b -> TxKey -> (TxKey, b, tx)
resolveOne SharedTxState peeraddr b
sharedState key :: TxKey
key@(TxKey Int
k) =
      ( TxKey
key
      , SharedTxState peeraddr b -> TxKey -> b
forall peeraddr txid.
HasCallStack =>
SharedTxState peeraddr txid -> TxKey -> txid
resolveTxKey SharedTxState peeraddr b
sharedState TxKey
key
      , case Int -> IntMap tx -> Maybe tx
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k (PeerTxLocalState tx -> IntMap tx
forall tx. PeerTxLocalState tx -> IntMap tx
peerDownloadedTxs PeerTxLocalState tx
peerState) of
             Just tx
tx -> tx
tx
             Maybe tx
Nothing -> [Char] -> tx
forall a. HasCallStack => [Char] -> a
error ([Char] -> tx) -> [Char] -> tx
forall a b. (a -> b) -> a -> b
$
               [Char]
"TxSubmission.V2.resolveBufferedTxsImp: missing buffered tx for key "
               [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
k
      )