Expose DiskSnapshotChecksum in Ouroboros.Consensus.Node

IntersectMBO · Dec 5, 2024 · f27ca6d · f27ca6d
1 parent 7088477
commit f27ca6d
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 14 deletions.
diff --git a/ouroboros-consensus-diffusion/src/ouroboros-consensus-diffusion/Ouroboros/Consensus/Node.hs b/ouroboros-consensus-diffusion/src/ouroboros-consensus-diffusion/Ouroboros/Consensus/Node.hs
@@ -6,6 +6,7 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MonadComprehensions #-}
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE Rank2Types #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -51,6 +52,8 @@ module Ouroboros.Consensus.Node (
   , RunNodeArgs (..)
   , Tracers
   , Tracers' (..)
+  , pattern DiskSnapshotChecksum
+  , pattern NoDiskSnapshotChecksum
     -- * Internal helpers
   , mkNodeKernelArgs
   , nodeKernelArgsEnforceInvariants
@@ -107,7 +110,8 @@ import           Ouroboros.Consensus.Storage.ChainDB (ChainDB, ChainDbArgs,
 import qualified Ouroboros.Consensus.Storage.ChainDB as ChainDB
 import qualified Ouroboros.Consensus.Storage.ChainDB.Impl.Args as ChainDB
 import           Ouroboros.Consensus.Storage.LedgerDB.DiskPolicy
-                     (DiskPolicyArgs (..))
+                     (DiskPolicyArgs (..), pattern DiskSnapshotChecksum,
+                     pattern NoDiskSnapshotChecksum)
 import           Ouroboros.Consensus.Util.Args
 import           Ouroboros.Consensus.Util.IOLike
 import           Ouroboros.Consensus.Util.Orphans ()

diff --git a/...boros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/ChainDB/Impl/LgrDB.hs b/...boros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/ChainDB/Impl/LgrDB.hs
@@ -222,6 +222,7 @@ initFromDisk ::
 initFromDisk LgrDbArgs { lgrHasFS = hasFS, .. }
              replayTracer
              immutableDB = wrapFailure (Proxy @blk) $ do
+    let LedgerDB.DiskPolicyArgs _ _ doDiskSnapshotChecksum  = lgrDiskPolicyArgs
     (_initLog, db, replayed) <-
       LedgerDB.initLedgerDB
         replayTracer
@@ -232,6 +233,7 @@ initFromDisk LgrDbArgs { lgrHasFS = hasFS, .. }
         lgrConfig
         lgrGenesis
         (streamAPI immutableDB)
+        doDiskSnapshotChecksum
     return (db, replayed)
   where
     ccfg = configCodec $ getExtLedgerCfg $ LedgerDB.ledgerDbCfg lgrConfig
@@ -280,12 +282,12 @@ takeSnapshot ::
      , IsLedger (LedgerState blk)
      )
   => LgrDB m blk -> m (Maybe (LedgerDB.DiskSnapshot, RealPoint blk))
-takeSnapshot lgrDB@LgrDB{ cfg, tracer, hasFS } = wrapFailure (Proxy @blk) $ do
+takeSnapshot lgrDB@LgrDB{ cfg, tracer, hasFS, diskPolicy } = wrapFailure (Proxy @blk) $ do
     ledgerDB <- LedgerDB.ledgerDbAnchor <$> atomically (getCurrent lgrDB)
     LedgerDB.takeSnapshot
       tracer
       hasFS
-      LedgerDB.DiskSnapshotChecksum
+      (LedgerDB.onDiskShouldChecksumSnapshots diskPolicy)
       (encodeDiskExtLedgerState ccfg)
       ledgerDB
   where

diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/LedgerDB/Init.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Storage/LedgerDB/Init.hs
@@ -108,6 +108,7 @@ initLedgerDB ::
   -> LedgerDbCfg (ExtLedgerState blk)
   -> m (ExtLedgerState blk) -- ^ Genesis ledger state
   -> StreamAPI m blk blk
+  -> Flag "DiskSnapshotChecksum"
   -> m (InitLog blk, LedgerDB' blk, Word64)
 initLedgerDB replayTracer
              tracer
@@ -116,9 +117,10 @@ initLedgerDB replayTracer
              decHash
              cfg
              getGenesisLedger
-             stream = do
+             stream
+             doDiskSnapshotChecksum = do
     snapshots <- listSnapshots hasFS
-    tryNewestFirst DiskSnapshotChecksum id snapshots
+    tryNewestFirst doDiskSnapshotChecksum id snapshots
   where
     tryNewestFirst :: Flag "DiskSnapshotChecksum"
                    -> (InitLog blk -> InitLog blk)

diff --git a/ouroboros-consensus/test/storage-test/Test/Ouroboros/Storage/LedgerDB/OnDisk.hs b/ouroboros-consensus/test/storage-test/Test/Ouroboros/Storage/LedgerDB/OnDisk.hs
@@ -295,7 +295,7 @@ data Cmd ss =
   | Snap (Flag "DiskSnapshotChecksum")
 
     -- | Restore the DB from on-disk, then return it along with the init log
-  | Restore
+  | Restore (Flag "DiskSnapshotChecksum")
 
     -- | Corrupt a previously taken snapshot
   | Corrupt Corruption ss
@@ -500,7 +500,7 @@ runMock cmd initMock =
     go Current       mock = (Ledger (cur (mockLedger mock)), mock)
     go (Push b)      mock = first MaybeErr $ mockUpdateLedger (push b)      mock
     go (Switch n bs) mock = first MaybeErr $ mockUpdateLedger (switch n bs) mock
-    go Restore       mock = (Restored (initLog, cur (mockLedger mock')), mock')
+    go Restore{}     mock = (Restored (initLog, cur (mockLedger mock')), mock')
       where
         initLog = mockInitLog mock
         mock'   = applyMockLog initLog mock
@@ -568,7 +568,7 @@ runMock cmd initMock =
           Delete   -> Nothing
           Truncate -> Just (ref, SnapCorrupted)
     go (Drop n) mock =
-        go Restore $ mock {
+        go (Restore NoDiskSnapshotChecksum) $ mock {
             mockLedger = drop (fromIntegral n) (mockLedger mock)
           }
 
@@ -759,7 +759,7 @@ runDB standalone@DB{..} cmd =
             doChecksum
             S.encode
             (ledgerDbAnchor db)
-    go hasFS Restore = do
+    go hasFS (Restore doChecksum) = do
         (initLog, db, _replayed) <-
           initLedgerDB
             nullTracer
@@ -770,6 +770,7 @@ runDB standalone@DB{..} cmd =
             dbLedgerDbCfg
             (return (testInitExtLedgerWithState initialTestLedgerState))
             stream
+            doChecksum
         atomically $ modifyTVar dbState (\(rs, _) -> (rs, db))
         return $ Restored (fromInitLog initLog, ledgerDbCurrent db)
     go hasFS (Corrupt c ss) =
@@ -786,7 +787,7 @@ runDB standalone@DB{..} cmd =
         atomically $ do
             (rs, _db) <- readTVar dbState
             writeTVar dbState (drop (fromIntegral n) rs, error "ledger DB not initialized")
-        go hasFS Restore
+        go hasFS (Restore NoDiskSnapshotChecksum)
 
     push ::
          TestBlock
@@ -943,7 +944,7 @@ generator secParam (Model mock hs) = Just $ QC.oneof $ concat [
                                 (lastAppliedPoint . ledgerState . mockCurrent $ afterRollback)
             return $ Switch numRollback blocks
         , fmap At $ Snap <$> QC.arbitrary
-        , fmap At $ return Restore
+        , fmap At $ Restore <$> QC.arbitrary
         , fmap At $ Drop <$> QC.choose (0, mockChainLength mock)
         ]
 
@@ -970,7 +971,7 @@ shrinker _ (At cmd) =
       Current      -> []
       Push _b      -> []
       Snap{}       -> []
-      Restore      -> []
+      Restore{}    -> []
       Switch 0 [b] -> [At $ Push b]
       Switch n bs  -> if length bs > fromIntegral n
                         then [At $ Switch n (init bs)]
@@ -1166,8 +1167,8 @@ tagEvents k = C.classify [
         fmap (TagRestore mST . rangeK k) $ C.maximum $ \ev ->
           let mock = modelMock (eventBefore ev) in
           case eventCmd ev of
-            At Restore | mockRecentSnap mock == mST -> Just (mockChainLength mock)
-            _otherwise                              -> Nothing
+            At (Restore{}) | mockRecentSnap mock == mST -> Just (mockChainLength mock)
+            _otherwise                                  -> Nothing
 
 {-------------------------------------------------------------------------------
   Inspecting the labelling function