Implement opening of an LSM LedgerDB backend

Author: jasagredo

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

@@ -17,6 +17,7 @@ module Ouroboros.Consensus.Storage.LedgerDB
   , openDBInternal
   ) where
 
+import Control.ResourceRegistry
 import Data.Functor.Contravariant ((>$<))
 import Data.Word
 import Ouroboros.Consensus.Block
@@ -32,6 +33,9 @@ import Ouroboros.Consensus.Storage.LedgerDB.TraceEvent
 import qualified Ouroboros.Consensus.Storage.LedgerDB.V1 as V1
 import qualified Ouroboros.Consensus.Storage.LedgerDB.V1.Snapshots as V1
 import qualified Ouroboros.Consensus.Storage.LedgerDB.V2 as V2
+import qualified Ouroboros.Consensus.Storage.LedgerDB.V2.Args as V2
+import qualified Ouroboros.Consensus.Storage.LedgerDB.V2.InMemory as InMemory
+import qualified Ouroboros.Consensus.Storage.LedgerDB.V2.LSM as LSM
 import Ouroboros.Consensus.Util.Args
 import Ouroboros.Consensus.Util.CallStack
 import Ouroboros.Consensus.Util.IOLike
@@ -41,7 +45,6 @@ openDB ::
   forall m blk.
   ( IOLike m
   , LedgerSupportsProtocol blk
-  , LedgerDbSerialiseConstraints blk
   , InspectLedger blk
   , HasCallStack
   , HasHardForkHistory blk
@@ -76,31 +79,51 @@ openDB
               snapManager
        in doOpenDB args initDb snapManager stream replayGoal
     LedgerDbFlavorArgsV2 bss -> do
+      (snapManager, bss') <- case bss of
+        V2.V2Args V2.InMemoryHandleArgs -> pure (InMemory.snapshotManager args, V2.InMemoryHandleEnv)
+        V2.V2Args (V2.LSMHandleArgs (V2.LSMArgs path genSalt mkFS)) -> do
+          (rk1, V2.SomeHasFSAndBlockIO fs blockio) <- mkFS (lgrRegistry args)
+          session <-
+            allocate
+              (lgrRegistry args)
+              ( \_ -> do
+                  salt <- genSalt
+                  LSM.openSession
+                    (LedgerDBFlavorImplEvent . FlavorImplSpecificTraceV2 . V2.LSMTrace >$< lgrTracer args)
+                    fs
+                    blockio
+                    salt
+                    (mkFsPath [path])
+              )
+              LSM.closeSession
+          pure (LSM.snapshotManager (snd session) args, V2.LSMHandleEnv session rk1)
       let initDb =
             V2.mkInitDb
               args
-              bss
+              bss'
               getBlock
-       in doOpenDB args initDb stream replayGoal
+              snapManager
+      doOpenDB args initDb snapManager stream replayGoal
 
 {-------------------------------------------------------------------------------
   Opening a LedgerDB
 -------------------------------------------------------------------------------}
 
 doOpenDB ::
-  forall m blk db.
+  forall m n blk db st.
   ( IOLike m
   , LedgerSupportsProtocol blk
   , InspectLedger blk
   , HasCallStack
   ) =>
   Complete LedgerDbArgs m blk ->
   InitDB db m blk ->
+  SnapshotManager m n blk st ->
   StreamAPI m blk blk ->
   Point blk ->
   m (LedgerDB' m blk, Word64)
-doOpenDB args initDb stream replayGoal =
-  f <$> openDBInternal args initDb stream replayGoal
+doOpenDB args initDb snapManager stream replayGoal =
+  f <$> openDBInternal args initDb snapManager stream replayGoal
  where
   f (ldb, replayCounter, _) = (ldb, replayCounter)
 
@@ -113,29 +136,28 @@ openDBInternal ::
   ) =>
   Complete LedgerDbArgs m blk ->
   InitDB db m blk ->
-  SnapshotManager m n blk st ->
+  SnapshotManagement m n blk st ->
   StreamAPI m blk blk ->
   Point blk ->
   m (LedgerDB' m blk, Word64, TestInternals' m blk)
-openDBInternal args@(LedgerDbArgs{lgrHasFS = SomeHasFS fs}) initDb stream replayGoal = do
+openDBInternal args@(LedgerDbArgs{lgrHasFS = SomeHasFS fs}) initDb snapManager stream replayGoal = do
   createDirectoryIfMissing fs True (mkFsPath [])
   (_initLog, db, replayCounter) <-
     initialize
       replayTracer
       snapTracer
-      lgrHasFS
       lgrConfig
       stream
       replayGoal
       initDb
+      snapManager
       lgrStartSnapshot
   (ledgerDb, internal) <- mkLedgerDb initDb db
   return (ledgerDb, replayCounter, internal)
  where
   LedgerDbArgs
     { lgrConfig
     , lgrTracer
-    , lgrHasFS
     , lgrStartSnapshot
     } = args
 

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

@@ -184,17 +184,20 @@ instance FromJSON SnapshotMetadata where
 data SnapshotBackend
   = UTxOHDMemSnapshot
   | UTxOHDLMDBSnapshot
+  | UTxOHDLSMSnapshot
   deriving (Eq, Show)
 
 instance ToJSON SnapshotBackend where
   toJSON = \case
     UTxOHDMemSnapshot -> "utxohd-mem"
     UTxOHDLMDBSnapshot -> "utxohd-lmdb"
+    UTxOHDLSMSnapshot -> "utxohd-lsm"
 
 instance FromJSON SnapshotBackend where
   parseJSON = Aeson.withText "SnapshotBackend" $ \case
     "utxohd-mem" -> pure UTxOHDMemSnapshot
     "utxohd-lmdb" -> pure UTxOHDLMDBSnapshot
+    "utxohd-lsm" -> pure UTxOHDLSMSnapshot
     _ -> fail "unknown SnapshotBackend"
 
 data MetadataErr

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

@@ -2,6 +2,7 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE NamedFieldPuns #-}
@@ -33,7 +34,6 @@ import Data.Set (Set)
 import qualified Data.Set as Set
 import Data.Traversable (for)
 import Data.Tuple (Solo (..))
-import Data.Void
 import Data.Word
 import GHC.Generics
 import NoThunks.Class
@@ -56,6 +56,8 @@ import Ouroboros.Consensus.Storage.LedgerDB.TraceEvent
 import Ouroboros.Consensus.Storage.LedgerDB.V2.Args as V2
 import Ouroboros.Consensus.Storage.LedgerDB.V2.Forker
 import qualified Ouroboros.Consensus.Storage.LedgerDB.V2.InMemory as InMemory
+import Ouroboros.Consensus.Storage.LedgerDB.V2.LSM (snapshotToStatePath)
+import qualified Ouroboros.Consensus.Storage.LedgerDB.V2.LSM as LSM
 import Ouroboros.Consensus.Storage.LedgerDB.V2.LedgerSeq
 import Ouroboros.Consensus.Util.Args
 import Ouroboros.Consensus.Util.CallStack
@@ -66,6 +68,8 @@ import Ouroboros.Network.Protocol.LocalStateQuery.Type
 import System.FS.API
 import Prelude hiding (read)
 
+type SnapshotManagerV2 m blk = SnapshotManager m m blk (StateRef m (ExtLedgerState blk))
+
 mkInitDb ::
   forall m blk.
   ( LedgerSupportsProtocol blk
@@ -75,10 +79,11 @@ mkInitDb ::
   , LedgerSupportsV2LedgerDB (LedgerState blk)
   ) =>
   Complete LedgerDbArgs m blk ->
-  Complete V2.LedgerDbFlavorArgs m ->
+  HandleEnv m ->
   ResolveBlock m blk ->
+  SnapshotManagerV2 m blk ->
   InitDB (LedgerSeq' m blk) m blk
-mkInitDb args flavArgs getBlock =
+mkInitDb args bss getBlock snapManager =
   InitDB
     { initFromGenesis = emptyF =<< lgrGenesis
     , initFromSnapshot =
@@ -112,9 +117,12 @@ mkInitDb args flavArgs getBlock =
                 , ldbResolveBlock = getBlock
                 , ldbQueryBatchSize = lgrQueryBatchSize
                 , ldbOpenHandlesLock = lock
+                , testLdbResourceKeys = case bss of
+                    InMemoryHandleEnv -> Nothing
+                    LSMHandleEnv (s, _) k -> Just (s, k)
                 }
         h <- LDBHandle <$> newTVarIO (LedgerDBOpen env)
-        pure $ implMkLedgerDb h bss
+        pure $ implMkLedgerDb h snapManager
     }
  where
   LedgerDbArgs
@@ -127,8 +135,6 @@ mkInitDb args flavArgs getBlock =
     , lgrRegistry
     } = args
 
-  bss = case flavArgs of V2Args bss0 -> bss0
-
   v2Tracer :: Tracer m V2.FlavorImplSpecificTrace
   v2Tracer = LedgerDBFlavorImplEvent . FlavorImplSpecificTraceV2 >$< lgrTracer
 
@@ -137,17 +143,20 @@ mkInitDb args flavArgs getBlock =
     m (LedgerSeq' m blk)
   emptyF st =
     empty' st $ case bss of
-      InMemoryHandleArgs -> InMemory.newInMemoryLedgerTablesHandle v2Tracer lgrHasFS
-      LSMHandleArgs x -> absurd x
+      InMemoryHandleEnv -> InMemory.newInMemoryLedgerTablesHandle v2Tracer lgrHasFS
+      LSMHandleEnv (_, session) _ ->
+        \values -> do
+          table <- LSM.tableFromValuesMK lgrRegistry session values
+          LSM.newLSMLedgerTablesHandle v2Tracer lgrRegistry session table
 
   loadSnapshot ::
     CodecConfig blk ->
     SomeHasFS m ->
     DiskSnapshot ->
     m (Either (SnapshotFailure blk) (LedgerSeq' m blk, RealPoint blk))
   loadSnapshot ccfg fs ds = case bss of
-    InMemoryHandleArgs -> runExceptT $ InMemory.loadSnapshot v2Tracer lgrRegistry ccfg fs ds
-    LSMHandleArgs x -> absurd x
+    InMemoryHandleEnv -> runExceptT $ InMemory.loadSnapshot v2Tracer lgrRegistry ccfg fs ds
+    LSMHandleEnv (_, session) _ -> runExceptT $ LSM.loadSnapshot v2Tracer lgrRegistry ccfg fs session ds
 
 implMkLedgerDb ::
   forall m l blk.
@@ -433,6 +442,10 @@ data LedgerDBEnv m l blk = LedgerDBEnv
   --
   --  * Modify 'ldbSeq' while holding a write lock, and then close the removed
   --    handles without any locking.
+  , testLdbResourceKeys :: !(Maybe (ResourceKey m, ResourceKey m))
+  -- ^ Resource keys used in the LSM backend so that the closing function used
+  -- in tests can release such resources. These are the resource keys for the
+  -- LSM session and the resource key for the BlockIO interface.
   }
   deriving Generic