From bcdda1e768639bdbeeff2009803d09202b4e8d01 Mon Sep 17 00:00:00 2001
From: Nicolas Frisby <nick.frisby@iohk.io>
Date: Sun, 13 Oct 2019 12:58:42 -0700
Subject: [PATCH] ouroboros-consensus: add newTestGuardedBlockchainTime

---
 .../src/Ouroboros/Consensus/BlockchainTime.hs | 26 ++++++++++++++++---
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/BlockchainTime.hs b/ouroboros-consensus/src/Ouroboros/Consensus/BlockchainTime.hs
index 726f4445ecf..385abee641f 100644
--- a/ouroboros-consensus/src/Ouroboros/Consensus/BlockchainTime.hs
+++ b/ouroboros-consensus/src/Ouroboros/Consensus/BlockchainTime.hs
@@ -17,6 +17,7 @@ module Ouroboros.Consensus.BlockchainTime (
   , NumSlots(..)
   , TestBlockchainTime(..)
   , newTestBlockchainTime
+  , newTestGuardedBlockchainTime
     -- * Real blockchain time
   , realBlockchainTime
     -- * Time to slots and back again
@@ -178,7 +179,19 @@ newTestBlockchainTime
     -> NumSlots           -- ^ Number of slots
     -> DiffTime           -- ^ Slot duration
     -> m (TestBlockchainTime m)
-newTestBlockchainTime registry (NumSlots numSlots) slotLen = do
+newTestBlockchainTime registry numSlots slotLen =
+    newTestGuardedBlockchainTime registry numSlots slotLen (\_ -> pure ())
+
+-- | Like 'newTestBlockchainTime', but also allows other computations to delay
+-- the end of each slot
+newTestGuardedBlockchainTime
+    :: forall m. (IOLike m, HasCallStack)
+    => ResourceRegistry m
+    -> NumSlots           -- ^ Number of slots
+    -> DiffTime           -- ^ Slot duration
+    -> (SlotNo -> m ())   -- ^ Blocks until slot is finished
+    -> m (TestBlockchainTime m)
+newTestGuardedBlockchainTime registry (NumSlots numSlots) slotLen waitOn = do
     slotVar <- newTVarM initVal
     doneVar <- newEmptyMVar ()
 
@@ -207,10 +220,15 @@ newTestBlockchainTime registry (NumSlots numSlots) slotLen = do
     loop slotVar doneVar = do
         -- count off each requested slot
         replicateM_ numSlots $ do
-          atomically $ modifyTVar slotVar $ Running . \case
-            Initializing -> SlotNo 0
-            Running slot -> succ slot
+          s' <- atomically $ do
+            st <- readTVar slotVar
+            let s' = case st of
+                    Initializing -> SlotNo 0
+                    Running s    -> succ s
+            writeTVar slotVar (Running s')
+            pure s'
           threadDelay slotLen
+          waitOn s'
         -- signal the end of the final slot
         putMVar doneVar ()