Ensure uncommitted forkers do not leak Ledger tables handles

Author: jasagredo

ouroboros-consensus/changelog.d/20250822_162104_jasataco_forkers.md

@@ -0,0 +1,24 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+For top level release notes, leave all the headers commented out.
+-->
+
+<!--
+### Patch
+
+- A bullet item for the Patch category.
+
+-->
+
+### Non-Breaking
+
+- Ensure uncommitted forkers do not leak Ledger tables handles.
+
+<!--
+### Breaking
+
+- A bullet item for the Breaking category.
+
+-->

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

@@ -17,13 +17,12 @@
 module Ouroboros.Consensus.Storage.LedgerDB.V2 (mkInitDb) where
 
 import Control.Arrow ((>>>))
-import qualified Control.Monad as Monad (void, (>=>))
+import qualified Control.Monad as Monad (join, void)
 import Control.Monad.Except
 import Control.RAWLock
 import qualified Control.RAWLock as RAWLock
 import Control.ResourceRegistry
 import Control.Tracer
-import Data.Foldable (traverse_)
 import qualified Data.Foldable as Foldable
 import Data.Functor.Contravariant ((>$<))
 import Data.Kind (Type)
@@ -573,13 +572,14 @@ getEnvSTM (LDBHandle varState) f =
 getStateRef ::
   (IOLike m, Traversable t) =>
   LedgerDBEnv m l blk ->
+  ResourceRegistry m ->
   (LedgerSeq m l -> t (StateRef m l)) ->
   m (t (StateRef m l))
-getStateRef ldbEnv project =
+getStateRef ldbEnv reg project =
   RAWLock.withReadAccess (ldbOpenHandlesLock ldbEnv) $ \() -> do
     tst <- project <$> readTVarIO (ldbSeq ldbEnv)
     for tst $ \st -> do
-      tables' <- duplicate $ tables st
+      (_, tables') <- allocate reg (\_ -> duplicate $ tables st) close
       pure st{tables = tables'}
 
 -- | Like 'StateRef', but takes care of closing the handle when the given action
@@ -590,8 +590,8 @@ withStateRef ::
   (LedgerSeq m l -> t (StateRef m l)) ->
   (t (StateRef m l) -> m a) ->
   m a
-withStateRef ldbEnv project =
-  bracket (getStateRef ldbEnv project) (traverse_ (close . tables))
+withStateRef ldbEnv project f =
+  withRegistry $ \reg -> getStateRef ldbEnv reg project >>= f
 
 acquireAtTarget ::
   ( HeaderHash l ~ HeaderHash blk
@@ -602,9 +602,10 @@ acquireAtTarget ::
   ) =>
   LedgerDBEnv m l blk ->
   Either Word64 (Target (Point blk)) ->
+  ResourceRegistry m ->
   m (Either GetForkerError (StateRef m l))
-acquireAtTarget ldbEnv target =
-  getStateRef ldbEnv $ \l -> case target of
+acquireAtTarget ldbEnv target reg =
+  getStateRef ldbEnv reg $ \l -> case target of
     Right VolatileTip -> pure $ currentHandle l
     Right ImmutableTip -> pure $ anchorHandle l
     Right (SpecificPoint pt) -> do
@@ -638,7 +639,7 @@ newForkerAtTarget ::
   Target (Point blk) ->
   m (Either GetForkerError (Forker m l blk))
 newForkerAtTarget h rr pt = getEnv h $ \ldbEnv ->
-  acquireAtTarget ldbEnv (Right pt) >>= traverse (newForker h ldbEnv rr)
+  acquireAtTarget ldbEnv (Right pt) rr >>= traverse (newForker h ldbEnv rr)
 
 newForkerByRollback ::
   ( HeaderHash l ~ HeaderHash blk
@@ -653,14 +654,14 @@ newForkerByRollback ::
   Word64 ->
   m (Either GetForkerError (Forker m l blk))
 newForkerByRollback h rr n = getEnv h $ \ldbEnv ->
-  acquireAtTarget ldbEnv (Left n) >>= traverse (newForker h ldbEnv rr)
+  acquireAtTarget ldbEnv (Left n) rr >>= traverse (newForker h ldbEnv rr)
 
 closeForkerEnv ::
   IOLike m => ForkerEnv m l blk -> m ()
 closeForkerEnv ForkerEnv{foeResourcesToRelease = (lock, key, toRelease)} =
   RAWLock.withWriteAccess lock $
     const $ do
-      id =<< atomically (swapTVar toRelease (pure ()))
+      Monad.join $ atomically (swapTVar toRelease (pure ()))
       _ <- release key
       pure ((), ())
 
@@ -757,7 +758,12 @@ newForker h ldbEnv rr st = do
   let tr = LedgerDBForkerEvent . TraceForkerEventWithKey forkerKey >$< ldbTracer ldbEnv
   traceWith tr ForkerOpen
   lseqVar <- newTVarIO . LedgerSeq . AS.Empty $ st
-  (k, toRelease) <- allocate rr (\_ -> newTVarIO (pure ())) (readTVarIO Monad.>=> id)
+  -- The closing action that we allocate in the TVar from the start is not
+  -- strictly necessary if the caller uses a short-lived registry like the ones
+  -- in Chain selection or the forging loop. Just in case the user passes a
+  -- long-lived registry, we store such closing action to make sure the handle
+  -- is closed even under @forkerClose@ if the registry outlives the forker.
+  (k, toRelease) <- allocate rr (\_ -> newTVarIO (close (tables st))) (Monad.join . readTVarIO)
   let forkerEnv =
         ForkerEnv
           { foeLedgerSeq = lseqVar

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

@@ -165,6 +165,17 @@ implForkerCommit env = do
                   _ AS.:< closeOld' -> closeLedgerSeq (LedgerSeq closeOld')
                 -- Finally, close the anchor of @lseq@ (which is a duplicate of
                 -- the head of @olddb'@).
+                --
+                -- Note if the resource registry used to create the Forker is
+                -- ephemeral as the one created on each Chain selection or each
+                -- Forging loop iteration, this first duplicated state will be
+                -- closed by the resource registry closing down, so this will be
+                -- a double release, which is fine. We prefer keeping this
+                -- action just in case some client passes a registry that
+                -- outlives the forker.
+                --
+                -- The rest of the states in the forker will be closed via
+                -- @foeResourcesToRelease@ instead of via the registry.
                 close $ tables $ AS.anchor lseq
           pure (closeDiscarded, s)
       )

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

@@ -96,8 +96,10 @@ newInMemoryLedgerTablesHandle tracer someFS@(SomeHasFS hasFS) l = do
   pure
     LedgerTablesHandle
       { close = do
-          atomically $ writeTVar tv LedgerTablesHandleClosed
-          traceWith tracer V2.TraceLedgerTablesHandleClose
+          p <- atomically $ swapTVar tv LedgerTablesHandleClosed
+          case p of
+            LedgerTablesHandleOpen{} -> traceWith tracer V2.TraceLedgerTablesHandleClose
+            _ -> pure ()
       , duplicate = do
           hs <- readTVarIO tv
           !x <- guardClosed hs $ newInMemoryLedgerTablesHandle tracer someFS