Skip to content

Commit 3f16a57

Browse files
jasagredojasagredo
committed
Define withTMVar and use it in Mempool
1 parent a27dc3f commit 3f16a57

File tree

3 files changed

+116
-73
lines changed

3 files changed

+116
-73
lines changed

‎ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Mempool/Init.hs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ mkMempool ::
104104
mkMempool mpEnv = Mempool
105105
{ addTx = implAddTx mpEnv
106106
, removeTxs = implRemoveTxs mpEnv
107-
, syncWithLedger = fst <$> implSyncWithLedger mpEnv
107+
, syncWithLedger = implSyncWithLedger mpEnv
108108
, getSnapshot = snapshotFromIS <$> readTMVar istate
109109
, getSnapshotFor = implGetSnapshotFor mpEnv
110110
, getCapacity = isCapacity <$> readTMVar istate

‎ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Mempool/Update.hs

Lines changed: 75 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,13 @@ module Ouroboros.Consensus.Mempool.Update (
1010

1111
import Cardano.Slotting.Slot
1212
import Control.Concurrent.Class.MonadMVar (withMVar)
13+
import Control.Monad (void)
1314
import Control.Monad.Except (runExcept)
1415
import Control.Tracer
1516
import qualified Data.List.NonEmpty as NE
1617
import Data.Maybe (fromMaybe)
1718
import qualified Data.Measure as Measure
1819
import qualified Data.Set as Set
19-
import Ouroboros.Consensus.Block.Abstract (castHash, castPoint,
20-
pointHash)
2120
import Ouroboros.Consensus.HeaderValidation
2221
import Ouroboros.Consensus.Ledger.Abstract
2322
import Ouroboros.Consensus.Ledger.SupportsMempool
@@ -27,8 +26,9 @@ import Ouroboros.Consensus.Mempool.Capacity
2726
import Ouroboros.Consensus.Mempool.Impl.Common
2827
import Ouroboros.Consensus.Mempool.TxSeq (TxTicket (..))
2928
import qualified Ouroboros.Consensus.Mempool.TxSeq as TxSeq
30-
import Ouroboros.Consensus.Util (whenJust)
29+
import Ouroboros.Consensus.Util (whenJust, withTMVarAnd)
3130
import Ouroboros.Consensus.Util.IOLike hiding (withMVar)
31+
import Ouroboros.Network.Block
3232

3333
{-------------------------------------------------------------------------------
3434
Add transactions
@@ -158,32 +158,37 @@ doAddTx mpEnv wti tx =
158158

159159
doAddTx' s = do
160160
traceWith trcr $ TraceMempoolAttemptingAdd tx
161-
is <- atomically $ do
162-
i <- takeTMVar istate
161+
res <- withTMVarAnd istate (\is ->
163162
case s of
164163
Nothing -> pure ()
165-
Just s' -> check $ isMempoolSize i /= s'
166-
pure i
167-
mTbs <- getLedgerTablesAtFor ldgrInterface (isTip is) [tx]
168-
case mTbs of
169-
Just tbs -> do
170-
traceWith trcr $ TraceMempoolLedgerFound (isTip is)
171-
case pureTryAddTx cfg wti tx is tbs of
172-
NotEnoughSpaceLeft -> do
173-
atomically $ putTMVar istate is
174-
doAddTx' (Just $ isMempoolSize is)
175-
Processed outcome@(TransactionProcessingResult is' _ _) -> do
176-
atomically $ putTMVar istate $ fromMaybe is is'
177-
pure outcome
178-
Nothing -> do
179-
traceWith trcr $ TraceMempoolLedgerNotFound (isTip is)
180-
-- We couldn't retrieve the values because the state is no longer on
181-
-- the db. We need to resync.
182-
atomically $ putTMVar istate is
183-
(_, mTrace) <- implSyncWithLedger mpEnv
184-
whenJust mTrace (traceWith trcr)
164+
Just s' -> check $ isMempoolSize is /= s')
165+
$ \is () -> do
166+
mTbs <- getLedgerTablesAtFor ldgrInterface (isTip is) [tx]
167+
case mTbs of
168+
Just tbs -> do
169+
traceWith trcr $ TraceMempoolLedgerFound (isTip is)
170+
case pureTryAddTx cfg wti tx is tbs of
171+
NotEnoughSpaceLeft -> do
172+
pure (Retry (isMempoolSize is), is)
173+
Processed outcome@(TransactionProcessingResult is' _ _) -> do
174+
pure (OK outcome, fromMaybe is is')
175+
Nothing -> do
176+
traceWith trcr $ TraceMempoolLedgerNotFound (isTip is)
177+
-- We couldn't retrieve the values because the state is no longer on
178+
-- the db. We need to resync.
179+
pure (Resync, is)
180+
case res of
181+
Retry s' -> doAddTx' (Just s')
182+
OK outcome -> pure outcome
183+
Resync -> do
184+
void $ implSyncWithLedger mpEnv
185185
doAddTx' s
186186

187+
data WithTMVarOutcome retry ok =
188+
Retry retry
189+
| OK ok
190+
| Resync
191+
187192
-- | Craft a 'TriedToAddTx' value containing the resulting state if
188193
-- applicable, the tracing event and the result of adding this transaction. See
189194
-- the documentation of 'implAddTx' for some more context.
@@ -318,10 +323,8 @@ implRemoveTxs ::
318323
-> NE.NonEmpty (GenTxId blk)
319324
-> m ()
320325
implRemoveTxs mpEnv toRemove = do
321-
(is, ls) <- atomically $ do
322-
is <- takeTMVar istate
323-
ls <- getCurrentLedgerState ldgrInterface
324-
pure (is, ls)
326+
out <- withTMVarAnd istate (const $ getCurrentLedgerState ldgrInterface)
327+
$ \is ls -> do
325328
let toKeep = filter
326329
( (`notElem` Set.fromList (NE.toList toRemove))
327330
. txId
@@ -333,9 +336,7 @@ implRemoveTxs mpEnv toRemove = do
333336
toKeep' = [ txForgetValidated . TxSeq.txTicketTx $ tx | tx <- toKeep ]
334337
mTbs <- getLedgerTablesAtFor ldgrInterface (castPoint (getTip ls)) toKeep'
335338
case mTbs of
336-
Nothing -> do
337-
atomically $ putTMVar istate is
338-
implRemoveTxs mpEnv toRemove
339+
Nothing -> pure (Resync, is)
339340
Just tbs -> do
340341
let (is', t) = pureRemoveTxs
341342
capacityOverride
@@ -346,8 +347,14 @@ implRemoveTxs mpEnv toRemove = do
346347
(isLastTicketNo is)
347348
toKeep
348349
toRemove
349-
atomically $ putTMVar istate is'
350350
traceWith trcr t
351+
pure (OK (), is')
352+
case out of
353+
Resync -> do
354+
void $ implSyncWithLedger mpEnv
355+
implRemoveTxs mpEnv toRemove
356+
OK () -> pure ()
357+
Retry _ -> error "Impossible!"
351358
where
352359
MempoolEnv { mpEnvStateVar = istate
353360
, mpEnvLedger = ldgrInterface
@@ -399,47 +406,43 @@ implSyncWithLedger ::
399406
, HasTxId (GenTx blk)
400407
)
401408
=> MempoolEnv m blk
402-
-> m (MempoolSnapshot blk, Maybe (TraceEventMempool blk))
409+
-> m (MempoolSnapshot blk)
403410
implSyncWithLedger mpEnv = do
404411
traceWith trcr TraceMempoolAttemptingSync
405-
(is, ls) <- atomically $ do
406-
is <- takeTMVar istate
407-
ls <- getCurrentLedgerState ldgrInterface
408-
pure (is, ls)
409-
410-
let (slot, ls') = tickLedgerState cfg $ ForgeInUnknownSlot ls
411-
412-
if pointHash (isTip is) == castHash (getTipHash ls) &&
413-
isSlotNo is == slot
414-
then do
415-
-- The tip didn't change, put the same state.
416-
atomically $ putTMVar istate is
417-
traceWith trcr $ TraceMempoolSyncNotNeeded (isTip is) (castPoint $ getTip ls)
418-
pure (snapshotFromIS is, Nothing)
419-
else do
420-
-- We need to revalidate
421-
let pt = castPoint (getTip ls)
422-
txs = [ txForgetValidated . TxSeq.txTicketTx $ tx
423-
| tx <- TxSeq.toList $ isTxs is
424-
]
425-
mTbs <- getLedgerTablesAtFor ldgrInterface pt txs
426-
case mTbs of
427-
Just tbs -> do
428-
let (is', mTrace) = pureSyncWithLedger
429-
capacityOverride
430-
cfg
431-
slot
432-
ls'
433-
tbs
434-
is
435-
atomically $ putTMVar istate is'
436-
whenJust mTrace (traceWith trcr)
437-
traceWith trcr TraceMempoolSyncDone
438-
return (snapshotFromIS is', mTrace)
439-
Nothing -> do
440-
-- If the point is gone, resync
441-
atomically $ putTMVar istate is
442-
implSyncWithLedger mpEnv
412+
res <- withTMVarAnd istate (const $ getCurrentLedgerState ldgrInterface) $
413+
\is ls -> do
414+
let (slot, ls') = tickLedgerState cfg $ ForgeInUnknownSlot ls
415+
if pointHash (isTip is) == castHash (getTipHash ls) && isSlotNo is == slot
416+
then do
417+
-- The tip didn't change, put the same state.
418+
traceWith trcr $ TraceMempoolSyncNotNeeded (isTip is) (castPoint $ getTip ls)
419+
pure (OK (snapshotFromIS is), is)
420+
else do
421+
-- We need to revalidate
422+
let pt = castPoint (getTip ls)
423+
txs = [ txForgetValidated . TxSeq.txTicketTx $ tx
424+
| tx <- TxSeq.toList $ isTxs is
425+
]
426+
mTbs <- getLedgerTablesAtFor ldgrInterface pt txs
427+
case mTbs of
428+
Just tbs -> do
429+
let (is', mTrace) = pureSyncWithLedger
430+
capacityOverride
431+
cfg
432+
slot
433+
ls'
434+
tbs
435+
is
436+
whenJust mTrace (traceWith trcr)
437+
traceWith trcr TraceMempoolSyncDone
438+
pure (OK (snapshotFromIS is'), is')
439+
Nothing -> do
440+
-- If the point is gone, resync
441+
pure (Resync, is)
442+
case res of
443+
OK v -> pure v
444+
Resync -> implSyncWithLedger mpEnv
445+
Retry _ -> error "Impossible!"
443446
where
444447
MempoolEnv { mpEnvStateVar = istate
445448
, mpEnvLedger = ldgrInterface

‎ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Util.hs

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,9 @@ module Ouroboros.Consensus.Util (
7979
, electric
8080
, newFuse
8181
, withFuse
82+
-- * withTMVar
83+
, withTMVar
84+
, withTMVarAnd
8285
) where
8386

8487
import Cardano.Crypto.Hash (Hash, HashAlgorithm, hashFromBytes,
@@ -463,3 +466,40 @@ withFuse (Fuse name m) (Electric io) = do
463466
newtype FuseBlownException = FuseBlownException Text
464467
deriving (Show)
465468
deriving anyclass (Exception)
469+
470+
{-------------------------------------------------------------------------------
471+
withTMVar
472+
-------------------------------------------------------------------------------}
473+
474+
-- | Apply @f@ with the content of @tv@ as state, restoring the original value when an
475+
-- exception occurs
476+
withTMVar ::
477+
IOLike m
478+
=> StrictTMVar m a
479+
-> (a -> m (c, a))
480+
-> m c
481+
withTMVar tv f = withTMVarAnd tv (const $ pure ()) (\a -> const $ f a)
482+
483+
-- | Apply @f@ with the content of @tv@ as state, restoring the original value
484+
-- when an exception occurs. Additionally run a @STM@ action when acquiring the
485+
-- value.
486+
withTMVarAnd ::
487+
IOLike m
488+
=> StrictTMVar m a
489+
-> (a -> STM m b) -- ^ Additional STM action to run in the same atomically
490+
-- block as the TMVar is acquired
491+
-> (a -> b -> m (c, a)) -- ^ Action
492+
-> m c
493+
withTMVarAnd tv guard f =
494+
bracketOnError
495+
(atomically $ do
496+
i <- takeTMVar tv
497+
g <- guard i
498+
pure (i, g)
499+
)
500+
(atomically . putTMVar tv . fst)
501+
(\(s, g) -> do
502+
(x, s') <- f s g
503+
atomically $ putTMVar tv s'
504+
return x
505+
)

0 commit comments

Comments
 (0)