Do unsafeRelease when closing all forkers

Author: jasagredo

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/LedgerDB/V1.hs

@@ -705,7 +705,7 @@ closeAllForkers forkersVar = do
     forkerEnvs <- Map.elems <$> readTVar forkersVar
     writeTVar forkersVar Map.empty
     return forkerEnvs
-  mapM_ closeForkerEnv forkerEnvs
+  mapM_ (closeForkerEnv (Flag @"ClosingAllForkers" True)) forkerEnvs
 
 -- | Acquire both a value handle and a db changelog at the tip. Holds a read lock
 -- while doing so.
@@ -822,4 +822,4 @@ implForkerClose (LDBHandle varState) forkerKey = do
           stateTVar
             (ldbForkers ldbEnv)
             (Map.updateLookupWithKey (\_ _ -> Nothing) forkerKey)
-  whenJust envMay closeForkerEnv
+  whenJust envMay (closeForkerEnv (Flag @"ClosingAllForkers" False))

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/LedgerDB/V1/Forker.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE FlexibleContexts #-}
@@ -44,6 +45,7 @@ import Ouroboros.Consensus.Storage.LedgerDB.V1.DiffSeq
   )
 import qualified Ouroboros.Consensus.Storage.LedgerDB.V1.DiffSeq as DS
 import Ouroboros.Consensus.Storage.LedgerDB.V1.Lock
+import Ouroboros.Consensus.Util
 import Ouroboros.Consensus.Util.IOLike
 
 {-------------------------------------------------------------------------------
@@ -90,9 +92,17 @@ deriving instance
   Close
 -------------------------------------------------------------------------------}
 
-closeForkerEnv :: IOLike m => ForkerEnv m l blk -> m ()
-closeForkerEnv ForkerEnv{foeBackingStoreValueHandle} = do
-  either (\(l, _, _) -> atomically . unsafeReleaseReadAccess $ l) (Monad.void . release . fst)
+closeForkerEnv :: IOLike m => Flag "ClosingAllForkers" -> ForkerEnv m l blk -> m ()
+closeForkerEnv closingAll ForkerEnv{foeBackingStoreValueHandle} = do
+  either
+    (\(l, _, _) -> atomically . unsafeReleaseReadAccess $ l)
+    ( Monad.void
+        -- During shutdown, we will close all forkers from the main thread
+        -- instead of from the thread that opened the forker, hence we need
+        -- unsafeRelease.
+        . (if getFlag closingAll then unsafeRelease else release)
+        . fst
+    )
     =<< takeMVar foeBackingStoreValueHandle
 
 {-------------------------------------------------------------------------------

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/LedgerDB/V2.hs

@@ -662,14 +662,20 @@ closeAllForkers ::
   m ()
 closeAllForkers lock forkersVar = do
   toClose <- atomically $ swapTVar forkersVar Map.empty
-  mapM_ (closeForkerEnv lock) toClose
+  mapM_ (closeForkerEnv (Flag @"ClosingAllForkers" True) lock) toClose
 
-closeForkerEnv :: IOLike m => RAWLock m LDBLock -> ForkerEnv m l blk -> m ()
-closeForkerEnv lock frkEnv =
+closeForkerEnv ::
+  IOLike m => Flag "ClosingAllForkers" -> RAWLock m LDBLock -> ForkerEnv m l blk -> m ()
+closeForkerEnv closingAll lock frkEnv =
   RAWLock.withWriteAccess lock $
     const $ do
       id =<< (atomically $ swapTVar (snd $ foeResourcesToRelease frkEnv) (pure ()))
-      _ <- release (fst (foeResourcesToRelease frkEnv))
+      _ <-
+        -- During shutdown, we will close all forkers from the main thread
+        -- instead of from the thread that opened the forker, hence we need
+        -- unsafeRelease.
+        (if getFlag closingAll then unsafeRelease else release)
+          (fst (foeResourcesToRelease frkEnv))
       pure ((), LDBLock)
 
 getForkerEnv ::
@@ -734,7 +740,7 @@ implForkerClose (LDBHandle varState) forkerKey = do
             <$> stateTVar
               (ldbForkers ldbEnv)
               (Map.updateLookupWithKey (\_ _ -> Nothing) forkerKey)
-  whenJust menv (uncurry closeForkerEnv)
+  whenJust menv (uncurry (closeForkerEnv $ Flag @"ClosingAllForkers" False))
 
 newForker ::
   ( IOLike m