Resurrect UTxO statistics in LSM-trees

Author: jasagredo

ouroboros-consensus-cardano/app/snapshot-converter.hs

@@ -458,7 +458,7 @@ main = withStdTerminalHandles $ do
         OutEnv
           fp
           ( \a b ->
-              SomeBackend <$> mkLSMSinkArgs lsmDbPath (last $ splitDirectories fp) stdMkBlockIOFS newStdGen a b
+              SomeBackend <$> mkLSMSinkArgs lsmDbPath fp (SomeHasFS . ioHasFS @IO) stdMkBlockIOFS newStdGen a b
           )
           Nothing
           (Just lsmDbPath)

ouroboros-consensus/changelog.d/20260109_160520_javier.sagredo_resurrect_utxo_count.md

@@ -0,0 +1,25 @@
+<!--
+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
+
+- A bullet item for the Non-Breaking category.
+
+-->
+
+### Breaking
+
+- LSM-trees backend is now able to track the size of the tables. This is exposed
+  via `tablesSize`, which changed from `Maybe Int` to `Int`.

ouroboros-consensus/ouroboros-consensus.cabal

@@ -382,6 +382,7 @@ library ouroboros-consensus-lsm
   build-depends:
     base >=4.14 && <4.22,
     blockio,
+    bytestring,
     containers >=0.5 && <0.8,
     contra-tracer,
     filepath,

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

@@ -42,11 +42,16 @@ module Ouroboros.Consensus.Storage.LedgerDB.V2.LSM
   ) where
 
 import Codec.Serialise (decode)
+import Control.Exception (assert)
 import qualified Control.Monad as Monad
 import Control.Monad.Trans (lift)
 import Control.Monad.Trans.Except
+import Control.Monad.Trans.Maybe (MaybeT (..), maybeToExceptT)
 import Control.ResourceRegistry
 import Control.Tracer
+import Data.ByteString (toStrict)
+import qualified Data.ByteString.Builder as BS
+import Data.ByteString.Char8 (readWord64)
 import qualified Data.Foldable as Foldable
 import Data.Functor.Contravariant ((>$<))
 import qualified Data.List as List
@@ -63,6 +68,7 @@ import qualified Data.Vector as V
 import qualified Data.Vector.Mutable as VM
 import qualified Data.Vector.Primitive as VP
 import Data.Void
+import Data.Word
 import Database.LSMTree (Salt, Session, Table)
 import qualified Database.LSMTree as LSM
 import GHC.Generics
@@ -87,6 +93,7 @@ import Ouroboros.Consensus.Util.IndexedMemPack
 import qualified Streaming as S
 import qualified Streaming.Prelude as S
 import System.FS.API
+import System.FS.API.Lazy (hGetAll, hPutAll)
 import qualified System.FS.BlockIO.API as BIO
 import System.FS.BlockIO.IO
 import System.FilePath (splitDirectories, splitFileName)
@@ -170,21 +177,24 @@ newLSMLedgerTablesHandle ::
   , IndexedMemPack (l EmptyMK) (TxOut l)
   ) =>
   Tracer m LedgerDBV2Trace ->
+  -- | The size of the tables
+  Word64 ->
   (ResourceKey m, UTxOTable m) ->
   m (LedgerTablesHandle m l)
-newLSMLedgerTablesHandle tracer (origResKey, t) =
+newLSMLedgerTablesHandle tracer origSize (origResKey, t) =
   encloseTimedWith (TraceLedgerTablesHandleCreate >$< tracer) $ do
     tv <- newTVarIO origResKey
+    tsize <- newTVarIO origSize
     pure
       LedgerTablesHandle
         { close = implClose tv
-        , duplicate = \rr -> implDuplicate rr t tracer
+        , duplicate = \rr -> implDuplicate rr tsize t tracer
         , read = implRead tracer t
         , readRange = implReadRange t
         , readAll = implReadAll t
-        , pushDiffs = implPushDiffs tracer t
+        , pushDiffs = implPushDiffs tracer t tsize
         , takeHandleSnapshot = implTakeHandleSnapshot tracer t
-        , tablesSize = pure Nothing
+        , tablesSize = fromIntegral <$> readTVarIO tsize
         , transfer = atomically . writeTVar tv
         }
 
@@ -206,16 +216,18 @@ implDuplicate ::
   , IndexedMemPack (l EmptyMK) (TxOut l)
   ) =>
   ResourceRegistry m ->
+  StrictTVar m Word64 ->
   UTxOTable m ->
   Tracer m LedgerDBV2Trace ->
   m (ResourceKey m, LedgerTablesHandle m l)
-implDuplicate rr t tracer = do
+implDuplicate rr sizeTVar t tracer = do
   (rk, table) <-
     allocate
       rr
       (\_ -> encloseTimedWith (TraceLedgerTablesHandleDuplicate >$< tracer) $ LSM.duplicate t)
       (encloseTimedWith (TraceLedgerTablesHandleClose >$< tracer) . LSM.closeTable)
-  (rk,) <$> newLSMLedgerTablesHandle tracer (rk, table)
+  size <- readTVarIO sizeTVar
+  (rk,) <$> newLSMLedgerTablesHandle tracer size (rk, table)
 
 implRead ::
   forall m l.
@@ -303,8 +315,8 @@ implPushDiffs ::
   , HasLedgerTables l
   , IndexedMemPack (l EmptyMK) (TxOut l)
   ) =>
-  Tracer m LedgerDBV2Trace -> UTxOTable m -> l mk -> l DiffMK -> m ()
-implPushDiffs tracer t _ !st1 =
+  Tracer m LedgerDBV2Trace -> UTxOTable m -> StrictTVar m Word64 -> l mk -> l DiffMK -> m ()
+implPushDiffs tracer t sizeTVar _ !st1 =
   encloseTimedWith (TraceLedgerTablesHandleRead >$< tracer) $ do
     let LedgerTables (DiffMK (Diff.Diff diffs)) = projectLedgerTables st1
     let vec = V.create $ do
@@ -314,6 +326,22 @@ implPushDiffs tracer t _ !st1 =
             0
             $ Map.toList diffs
           pure vec'
+    let (ins, dels) =
+          Map.foldl'
+            ( \(i, d) delta -> case delta of
+                Diff.Insert{} -> (i + 1, d)
+                Diff.Delete -> (i, d + 1)
+            )
+            (0, 0)
+            diffs
+    atomically $
+      modifyTVar
+        sizeTVar
+        ( \x ->
+            assert (x + ins >= x) $
+              assert (x + ins - dels <= x + ins) $
+                x + ins - dels
+        )
     encloseTimedWith (BackendTrace . SomeBackendTrace . LSMUpdate >$< tracer) $ LSM.updates t vec
  where
   f (Diff.Insert v) = LSM.Insert (toTxOutBytes (forgetLedgerTables st1) v) Nothing
@@ -394,21 +422,41 @@ implTakeSnapshot ccfg tracer shfs@(SomeHasFS hasFs) suffix st =
         then
           return Nothing
         else do
+          sz <- tablesSize (tables st)
           encloseTimedWith (TookSnapshot snapshot t >$< tracer) $
-            writeSnapshot snapshot
+            writeSnapshot sz snapshot
           return $ Just (snapshot, t)
  where
-  writeSnapshot ds = do
+  writeSnapshot sz ds = do
     createDirectoryIfMissing hasFs True $ snapshotToDirPath ds
     crc1 <- writeExtLedgerState shfs (encodeDiskExtLedgerState ccfg) (snapshotToStatePath ds) $ state st
     crc2 <- takeHandleSnapshot (tables st) (state st) $ snapshotToDirName ds
+    writeUTxOSizeFile hasFs (snapshotToUTxOSizeFilePath ds) sz
     writeSnapshotMetadata shfs ds $
       SnapshotMetadata
         { snapshotBackend = UTxOHDLSMSnapshot
         , snapshotChecksum = maybe crc1 (crcOfConcat crc1) crc2
         , snapshotTablesCodecVersion = TablesCodecVersion1
         }
 
+snapshotToUTxOSizeFilePath :: DiskSnapshot -> FsPath
+snapshotToUTxOSizeFilePath ds = snapshotToDirPath ds </> mkFsPath ["utxoSize"]
+
+writeUTxOSizeFile :: MonadThrow f => HasFS f h -> FsPath -> Int -> f ()
+writeUTxOSizeFile hasFs p sz =
+  Monad.void $ withFile hasFs p (WriteMode MustBeNew) $ \h ->
+    hPutAll hasFs h $ BS.toLazyByteString $ BS.intDec sz
+
+readUTxOSizeFile :: MonadThrow m => HasFS m h -> FsPath -> ExceptT (SnapshotFailure blk) m Word64
+readUTxOSizeFile hfs p = do
+  exists <- lift $ doesFileExist hfs p
+  Monad.unless exists $ throwE (InitFailureRead ReadSnapshotDataCorruption)
+  fmap fst $
+    maybeToExceptT (InitFailureRead ReadSnapshotDataCorruption) $
+      MaybeT $
+        withFile hfs p ReadMode $ \h ->
+          readWord64 . toStrict <$> hGetAll hfs h
+
 -- | Delete snapshot from disk and also from the LSM tree database.
 implDeleteSnapshot ::
   IOLike m =>
@@ -471,6 +519,7 @@ loadSnapshot tracer rr ccfg fs@(SomeHasFS hfs) session ds =
       withExceptT
         (InitFailureRead . ReadSnapshotFailed)
         $ readExtLedgerState fs (decodeDiskExtLedgerState ccfg) decode (snapshotToStatePath ds)
+    msz <- readUTxOSizeFile hfs (snapshotToUTxOSizeFilePath ds)
     case pointToWithOriginRealPoint (castPoint (getTip extLedgerSt)) of
       Origin -> throwE InitFailureGenesis
       NotOrigin pt -> do
@@ -492,7 +541,7 @@ loadSnapshot tracer rr ccfg fs@(SomeHasFS hfs) session ds =
           $ InitFailureRead
             ReadSnapshotDataCorruption
         (,pt)
-          <$> lift (empty extLedgerSt (rk, values) (newLSMLedgerTablesHandle tracer))
+          <$> lift (empty extLedgerSt (rk, values) (newLSMLedgerTablesHandle tracer msz))
 
 -- | Create the initial LSM table from values, which should happen only at
 -- Genesis.
@@ -504,7 +553,7 @@ tableFromValuesMK ::
   Session m ->
   l EmptyMK ->
   LedgerTables l ValuesMK ->
-  m (ResourceKey m, UTxOTable m)
+  m (ResourceKey m, UTxOTable m, Word64)
 tableFromValuesMK tracer rr session st (LedgerTables (ValuesMK values)) = do
   (rk, table) <-
     allocate
@@ -515,7 +564,7 @@ tableFromValuesMK tracer rr session st (LedgerTables (ValuesMK values)) = do
       )
       (encloseTimedWith (TraceLedgerTablesHandleClose >$< tracer) . LSM.closeTable)
   mapM_ (go table) $ chunks 1000 $ Map.toList values
-  pure (rk, table)
+  pure (rk, table, fromIntegral $ Map.size values)
  where
   go table items =
     LSM.inserts table $
@@ -613,9 +662,9 @@ instance
     loadSnapshot trcr reg ccfg shfs (sessionResource res) ds
 
   newHandleFromValues trcr reg res st = do
-    table <-
+    (rk, table, sz) <-
       tableFromValuesMK trcr reg (sessionResource res) (forgetLedgerTables st) (ltprj st)
-    newLSMLedgerTablesHandle trcr table
+    newLSMLedgerTablesHandle trcr sz (rk, table)
 
   snapshotManager _ res = Ouroboros.Consensus.Storage.LedgerDB.V2.LSM.snapshotManager (sessionResource res)
 
@@ -636,13 +685,15 @@ instance
     = SinkLSM
         -- \| Chunk size
         Int
-        -- \| Snap name
-        String
+        -- \| LedgerDB snapshot fs
+        (SomeHasFS m)
+        -- \| DiskSnapshot
+        DiskSnapshot
         (Session m)
 
   yield _ (YieldLSM chunkSize hdl) = yieldLsmS chunkSize hdl
 
-  sink _ (SinkLSM chunkSize snapName session) = sinkLsmS chunkSize snapName session
+  sink _ (SinkLSM chunkSize shfs ds session) = sinkLsmS chunkSize shfs ds session
 
 data SomeHasFSAndBlockIO m where
   SomeHasFSAndBlockIO ::
@@ -684,35 +735,37 @@ sinkLsmS ::
   , IndexedMemPack (l EmptyMK) (TxOut l)
   ) =>
   Int ->
-  String ->
+  SomeHasFS m ->
+  DiskSnapshot ->
   Session m ->
   Sink m l
-sinkLsmS writeChunkSize snapName session st s = do
-  tb :: UTxOTable m <- lift $ LSM.newTable session
-  r <- go tb writeChunkSize mempty s
+sinkLsmS writeChunkSize (SomeHasFS hfs) ds session st stream = do
+  lsmTable :: UTxOTable m <- lift $ LSM.newTable session
+  (r, utxosSize) <- go (0 :: Int) lsmTable writeChunkSize mempty stream
   lift $
     LSM.saveSnapshot
-      (LSM.toSnapshotName snapName)
+      (LSM.toSnapshotName (snapshotToDirName ds))
       (LSM.SnapshotLabel $ T.pack "UTxO table")
-      tb
-  lift $ LSM.closeTable tb
+      lsmTable
+  lift $ LSM.closeTable lsmTable
+  lift $ writeUTxOSizeFile hfs (snapshotToUTxOSizeFilePath ds) utxosSize
   pure (fmap (,Nothing) r)
  where
-  go tb 0 m s' = do
-    lift $
-      LSM.inserts tb $
-        V.fromList [(toTxInBytes (Proxy @l) k, toTxOutBytes st v, Nothing) | (k, v) <- m]
-    go tb writeChunkSize mempty s'
-  go tb n m s' = do
-    mbs <- S.uncons s'
-    case mbs of
+  writeToTable lsmTable accUTxOs =
+    LSM.inserts lsmTable $
+      V.fromList
+        [(toTxInBytes (Proxy @l) txin, toTxOutBytes st txout, Nothing) | (txin, txout) <- accUTxOs]
+
+  go utxosSize lsmTable 0 accUTxOs stream' = do
+    lift $ writeToTable lsmTable accUTxOs
+    go utxosSize lsmTable writeChunkSize mempty stream'
+  go utxosSize lsmTable numToRead accUTxOs stream' = do
+    mItem <- S.uncons stream'
+    case mItem of
       Nothing -> do
-        lift $
-          LSM.inserts tb $
-            V.fromList
-              [(toTxInBytes (Proxy @l) k, toTxOutBytes st v, Nothing) | (k, v) <- m]
-        S.effects s'
-      Just (item, s'') -> go tb (n - 1) (item : m) s''
+        lift $ writeToTable lsmTable accUTxOs
+        (,utxosSize) <$> S.effects stream'
+      Just (item, stream'') -> go (utxosSize + 1) lsmTable (numToRead - 1) (item : accUTxOs) stream''
 
 -- | Create Yield arguments for LSM
 mkLSMYieldArgs ::
@@ -746,15 +799,17 @@ mkLSMYieldArgs fp snapName mkFS mkGen _ reg = do
             (LSM.SnapshotLabel $ T.pack "UTxO table")
       )
       LSM.closeTable
-  YieldLSM 1000 <$> newLSMLedgerTablesHandle nullTracer tb
+  YieldLSM 1000 <$> newLSMLedgerTablesHandle nullTracer 0 tb
 
 -- | Create Sink arguments for LSM
 mkLSMSinkArgs ::
   IOLike m =>
   -- | The filepath in which the LSM database should be opened. Must not have a trailing slash!
   FilePath ->
-  -- | The complete name of the snapshot to be created, so @<slotno>[_<suffix>]@.
-  String ->
+  -- | The filepath to the snapshot to be created, so @.../.../ledger/<slotno>[_<suffix>]@.
+  FilePath ->
+  -- | Usually 'ioHasFS'
+  (MountPoint -> SomeHasFS m) ->
   -- | Usually 'stdMkBlockIOFS'
   (FilePath -> ResourceRegistry m -> m (a, SomeHasFSAndBlockIO m)) ->
   -- | Usually 'newStdGen'
@@ -764,18 +819,20 @@ mkLSMSinkArgs ::
   m (SinkArgs m LSM l)
 mkLSMSinkArgs
   (splitFileName -> (fp, lsmDir))
-  snapName
-  mkFS
+  snapFP
+  mkFs
+  mkBlockIOFS
   mkGen
   _
   reg =
     do
-      (_, SomeHasFSAndBlockIO hasFS blockIO) <- mkFS fp reg
+      (_, SomeHasFSAndBlockIO hasFS blockIO) <- mkBlockIOFS fp reg
       removeDirectoryRecursive hasFS lsmFsPath
       createDirectory hasFS lsmFsPath
       salt <- fst . genWord64 <$> mkGen
       (_, session) <-
         allocate reg (\_ -> LSM.newSession nullTracer hasFS blockIO salt lsmFsPath) LSM.closeSession
-      pure (SinkLSM 1000 snapName session)
+      let snapFS = mkFs (MountPoint snapFP)
+      pure (SinkLSM 1000 snapFS (fromJust $ snapshotFromPath $ last $ splitDirectories snapFP) session)
    where
     lsmFsPath = mkFsPath [lsmDir]

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/ChainDB/API.hs

@@ -397,7 +397,7 @@ data ChainDB m blk = ChainDB
   , getChainSelStarvation :: STM m ChainSelStarvation
   -- ^ Whether ChainSel is currently starved, or when was last time it
   -- stopped being starved.
-  , getStatistics :: m (Maybe Statistics)
+  , getStatistics :: m Statistics
   -- ^ Get statistics from the LedgerDB, in particular the number of entries
   -- in the tables.
   , addPerasCertAsync :: ValidatedPerasCert blk -> m (AddPerasCertPromise m)

ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/ChainDB/Impl/Query.hs

@@ -289,7 +289,7 @@ getReadOnlyForkerAtPoint ::
   m (Either LedgerDB.GetForkerError (LedgerDB.ReadOnlyForker' m blk))
 getReadOnlyForkerAtPoint CDB{..} = LedgerDB.getReadOnlyForker cdbLedgerDB
 
-getStatistics :: IOLike m => ChainDbEnv m blk -> m (Maybe LedgerDB.Statistics)
+getStatistics :: IOLike m => ChainDbEnv m blk -> m LedgerDB.Statistics
 getStatistics CDB{..} = LedgerDB.getTipStatistics cdbLedgerDB
 
 getPerasWeightSnapshot ::