From 35e19ced5282ea144bf4dd41655b0600157b446e Mon Sep 17 00:00:00 2001
From: Javier Sagredo <javier.sagredo@iohk.io>
Date: Mon, 25 Oct 2021 13:48:34 +0200
Subject: [PATCH] Fix quickcheck state machine tests

---
 .../Ouroboros/Storage/ChainDB/StateMachine.hs  |  7 +++----
 .../Consensus/Util/ResourceRegistry.hs         | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs b/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs
index 4d90d74430b..d3061478058 100644
--- a/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs
+++ b/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs
@@ -298,7 +298,7 @@ type TestFollower m blk = WithEq (Follower m blk (AllComponents blk))
 data ChainDBState m blk = ChainDBState
     { chainDB       :: ChainDB m blk
     , internal      :: ChainDB.Internal m blk
-    , addBlockAsync :: Async m Void
+    , addBlockAsync :: Thread m Void
       -- ^ Background thread that adds blocks to the ChainDB
     }
   deriving NoThunks via AllowThunk (ChainDBState m blk)
@@ -319,8 +319,7 @@ open
   => ChainDbArgs Identity m blk -> m (ChainDBState m blk)
 open args = do
     (chainDB, internal) <- openDBInternal args False
-    addBlockAsync       <- async (intAddBlockRunner internal)
-    link addBlockAsync
+    addBlockAsync       <- forkThreadAndLink (cdbRegistry args) "Add block runner" (intAddBlockRunner internal)
     return ChainDBState { chainDB, internal, addBlockAsync }
 
 -- PRECONDITION: the ChainDB is closed
@@ -333,7 +332,7 @@ reopen ChainDBEnv { varDB, args } = do
 
 close :: IOLike m => ChainDBState m blk -> m ()
 close ChainDBState { chainDB, addBlockAsync } = do
-    cancel addBlockAsync
+    cancelThread addBlockAsync
     closeDB chainDB
 
 run :: forall m blk.
diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/Util/ResourceRegistry.hs b/ouroboros-consensus/src/Ouroboros/Consensus/Util/ResourceRegistry.hs
index be7f99b2522..f2d1f7aacad 100644
--- a/ouroboros-consensus/src/Ouroboros/Consensus/Util/ResourceRegistry.hs
+++ b/ouroboros-consensus/src/Ouroboros/Consensus/Util/ResourceRegistry.hs
@@ -49,6 +49,7 @@ module Ouroboros.Consensus.Util.ResourceRegistry (
     -- * Combinators primarily for testing
   , closeRegistry
   , countResources
+  , forkThreadAndLink
   , unsafeNewRegistry
     -- * opaque
   , ResourceRegistry
@@ -1132,6 +1133,23 @@ forkThread rr label body = snd <$>
           removeThread tid
           void $ removeResource rid
 
+-- | Convenience function to group actions done in ChainDB QuickCheck state
+-- machine tests.
+--
+-- It is just a combination of 'forkLinkedThread' together with linking the
+-- exceptions thrown in the async thread to the main one.
+forkThreadAndLink :: forall m a. (IOLike m, HasCallStack)
+           => ResourceRegistry m
+           -> String  -- ^ Label for the thread
+           -> m a
+           -> m (Thread m a)
+forkThreadAndLink rr label body = do
+  t <- forkThread rr label body
+  link $ threadAsync t
+  linkToRegistry t
+  return t
+
+
 -- | Bracketed version of 'forkThread'
 --
 -- The analogue of 'withAsync' for the registry.