Allocate forker maps in the registries

Author: jasagredo

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

@@ -90,6 +90,8 @@ mkInitDb args bss getBlock =
         st <- lgrGenesis
         let genesis = forgetLedgerTables st
             chlog = DbCh.empty genesis
+        -- We ignore the ResourceKey here because we only want this to close
+        -- once the node is shutting down.
         (_, backingStore) <-
           allocate
             lgrRegistry
@@ -120,7 +122,9 @@ mkInitDb args bss getBlock =
         (varDB, prevApplied) <-
           (,) <$> newTVarIO db <*> newTVarIO Set.empty
         flushLock <- mkLedgerDBLock
-        forkers <- newTVarIO Map.empty
+        -- We ignore the ResourceKey here because we only want this to close
+        -- once the node is shutting down.
+        (_, forkers) <- allocate lgrRegistry (\_ -> newTVarIO Map.empty) closeAllForkers
         nextForkerKey <- newTVarIO (ForkerKey 0)
         let env =
               LedgerDBEnv
@@ -334,7 +338,7 @@ implCloseDB (LDBHandle varState) = do
 
   -- Only when the LedgerDB was open
   whenJust mbOpenEnv $ \env -> do
-    closeAllForkers env
+    closeAllForkers (ldbForkers env)
     bsClose (ldbBackingStore env)
 
 mkInternals ::
@@ -694,17 +698,14 @@ newForkerByRollback h rr n = getEnv h $ \ldbEnv -> do
 -- | Close all open block and header 'Forker's.
 closeAllForkers ::
   IOLike m =>
-  LedgerDBEnv m l blk ->
+  StrictTVar m (Map ForkerKey (ForkerEnv m l blk)) ->
   m ()
-closeAllForkers ldbEnv =
-  do
-    forkerEnvs <- atomically $ do
-      forkerEnvs <- Map.elems <$> readTVar forkersVar
-      writeTVar forkersVar Map.empty
-      return forkerEnvs
-    mapM_ closeForkerEnv forkerEnvs
- where
-  forkersVar = ldbForkers ldbEnv
+closeAllForkers forkersVar = do
+  forkerEnvs <- atomically $ do
+    forkerEnvs <- Map.elems <$> readTVar forkersVar
+    writeTVar forkersVar Map.empty
+    return forkerEnvs
+  mapM_ closeForkerEnv forkerEnvs
 
 -- | Acquire both a value handle and a db changelog at the tip. Holds a read lock
 -- while doing so.
@@ -775,9 +776,13 @@ newForker h ldbEnv rr dblog = readLocked $ do
           }
   atomically $ do
     -- We need to make sure to release this read access when we drop the value
-    -- handle, so in 'closeForkerEnv' if it wasn't promoted or when promoting it
-    -- in 'getValueHandle'.
+    -- handle, so in 'closeForkerEnv' (if it wasn't promoted) or in
+    -- 'getValueHandle' (if it was promoted).
     unsafeAcquireReadAccess (ldbLock ldbEnv)
+
+    -- Note that we add the forkerEnv to the 'ldbForkers' so that an exception
+    -- which will close all the forkers, also closes this one, releasing the
+    -- read access we acquired above.
     modifyTVar (ldbForkers ldbEnv) $ Map.insert forkerKey forkerEnv
   traceWith (foeTracer forkerEnv) ForkerOpen
   pure $ mkForker h (ldbQueryBatchSize ldbEnv) forkerKey

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

@@ -98,9 +98,11 @@ mkInitDb args flavArgs getBlock =
     , mkLedgerDb = \lseq -> do
         varDB <- newTVarIO lseq
         prevApplied <- newTVarIO Set.empty
-        forkers <- newTVarIO Map.empty
-        nextForkerKey <- newTVarIO (ForkerKey 0)
         lock <- RAWLock.new LDBLock
+        -- We ignore the ResourceKey here because we only want this to close
+        -- once the node is shutting down.
+        (_, forkers) <- allocate lgrRegistry (\_ -> newTVarIO Map.empty) (closeAllForkers lock)
+        nextForkerKey <- newTVarIO (ForkerKey 0)
         let env =
               LedgerDBEnv
                 { ldbSeq = varDB
@@ -421,7 +423,7 @@ implCloseDB (LDBHandle varState) = do
 
   -- Only when the LedgerDB was open
   whenJust mbOpenEnv $ \env -> do
-    closeAllForkers env
+    closeAllForkers (ldbOpenHandlesLock env) (ldbForkers env)
 
 {-------------------------------------------------------------------------------
   The LedgerDBEnv
@@ -655,17 +657,16 @@ newForkerByRollback h rr n = getEnv h $ \ldbEnv ->
 -- | Close all open 'Forker's.
 closeAllForkers ::
   IOLike m =>
-  LedgerDBEnv m l blk ->
+  RAWLock m LDBLock ->
+  StrictTVar m (Map ForkerKey (ForkerEnv m l blk)) ->
   m ()
-closeAllForkers ldbEnv = do
-  toClose <- fmap (ldbEnv,) <$> (atomically $ stateTVar forkersVar (,Map.empty))
-  mapM_ closeForkerEnv toClose
- where
-  forkersVar = ldbForkers ldbEnv
+closeAllForkers lock forkersVar = do
+  toClose <- atomically $ swapTVar forkersVar Map.empty
+  mapM_ (closeForkerEnv lock) toClose
 
-closeForkerEnv :: IOLike m => (LedgerDBEnv m l blk, ForkerEnv m l blk) -> m ()
-closeForkerEnv (LedgerDBEnv{ldbOpenHandlesLock}, frkEnv) =
-  RAWLock.withWriteAccess ldbOpenHandlesLock $
+closeForkerEnv :: IOLike m => RAWLock m LDBLock -> ForkerEnv m l blk -> m ()
+closeForkerEnv lock frkEnv =
+  RAWLock.withWriteAccess lock $
     const $ do
       id =<< (atomically $ swapTVar (snd $ foeResourcesToRelease frkEnv) (pure ()))
       _ <- release (fst (foeResourcesToRelease frkEnv))
@@ -729,11 +730,11 @@ implForkerClose (LDBHandle varState) forkerKey = do
       readTVar varState >>= \case
         LedgerDBClosed -> pure Nothing
         LedgerDBOpen ldbEnv ->
-          fmap (ldbEnv,)
+          fmap (ldbOpenHandlesLock ldbEnv,)
             <$> stateTVar
               (ldbForkers ldbEnv)
               (Map.updateLookupWithKey (\_ _ -> Nothing) forkerKey)
-  whenJust menv closeForkerEnv
+  whenJust menv (uncurry closeForkerEnv)
 
 newForker ::
   ( IOLike m