LDB V2: Fix leaked handle when not committing a forker

Author: jasagredo

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)
@@ -151,7 +150,8 @@ mkInitDb args bss getBlock snapManager =
       InMemoryHandleEnv -> InMemory.newInMemoryLedgerTablesHandle v2Tracer lgrHasFS
       LSMHandleEnv lsmRes ->
         \values -> do
-          table <- LSM.tableFromValuesMK lgrRegistry (sessionResource lsmRes) (forgetLedgerTables st) values
+          table <-
+            LSM.tableFromValuesMK v2Tracer lgrRegistry (sessionResource lsmRes) (forgetLedgerTables st) values
           LSM.newLSMLedgerTablesHandle v2Tracer lgrRegistry table
 
   loadSnapshot ::
@@ -562,13 +562,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
@@ -579,8 +580,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
@@ -591,9 +592,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
@@ -627,7 +629,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
@@ -642,14 +644,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 ((), ())
 
@@ -746,7 +748,7 @@ 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)
+  (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

@@ -167,6 +167,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 it should (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 doesn't respect the
+                -- condition of registries being ephemeral.
+                --
+                -- 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/LSM.hs

@@ -163,19 +163,23 @@ instance LSM.SerialiseKey TxInBytes where
 tableFromValuesMK ::
   forall m l.
   (IOLike m, IndexedMemPack (l EmptyMK) (TxOut l), MemPack (TxIn l)) =>
+  Tracer m V2.FlavorImplSpecificTrace ->
   ResourceRegistry m ->
   Session m ->
   l EmptyMK ->
   LedgerTables l ValuesMK ->
   m (ResourceKey m, UTxOTable m)
-tableFromValuesMK rr session st (LedgerTables (ValuesMK values)) = do
+tableFromValuesMK tracer rr session st (LedgerTables (ValuesMK values)) = do
   res@(_, table) <-
     allocate
       rr
       ( \_ ->
           LSM.newTableWith (LSM.defaultTableConfig{LSM.confFencePointerIndex = LSM.OrdinaryIndex}) session
       )
-      LSM.closeTable
+      ( \tb -> do
+          traceWith tracer V2.TraceLedgerTablesHandleClose
+          LSM.closeTable tb
+      )
   mapM_ (go table) $ chunks 1000 $ Map.toList values
   pure res
  where
@@ -232,9 +236,15 @@ newLSMLedgerTablesHandle tracer rr (resKey, t) = do
     LedgerTablesHandle
       { close = do
           Monad.void $ release resKey
-          traceWith tracer V2.TraceLedgerTablesHandleClose
       , duplicate = do
-          table <- allocate rr (\_ -> LSM.duplicate t) LSM.closeTable
+          table <-
+            allocate
+              rr
+              (\_ -> LSM.duplicate t)
+              ( \t' -> do
+                  traceWith tracer V2.TraceLedgerTablesHandleClose
+                  LSM.closeTable t'
+              )
           newLSMLedgerTablesHandle tracer rr table
       , read = \st (LedgerTables (KeysMK keys)) -> do
           let vec' = V.create $ do
@@ -431,7 +441,10 @@ loadSnapshot tracer rr ccfg fs session ds = do
                   (fromString $ snapshotToDirName ds)
                   (LSM.SnapshotLabel $ Text.pack $ "UTxO table")
             )
-            LSM.closeTable
+            ( \t -> do
+                traceWith tracer V2.TraceLedgerTablesHandleClose
+                LSM.closeTable t
+            )
       Monad.when (checksumAsRead /= snapshotChecksum snapshotMeta) $
         throwE $
           InitFailureRead ReadSnapshotDataCorruption