WIP: rework snapshots for ongoing merges

Author: jorisdral

lsm-tree.cabal

@@ -152,6 +152,7 @@ library
     Database.LSMTree.Internal.RunReaders
     Database.LSMTree.Internal.Serialise
     Database.LSMTree.Internal.Serialise.Class
+    Database.LSMTree.Internal.Snapshot
     Database.LSMTree.Internal.UniqCounter
     Database.LSMTree.Internal.Unsliced
     Database.LSMTree.Internal.Vector

src-extras/Database/LSMTree/Extras/NoThunks.hs

@@ -280,6 +280,9 @@ deriving anyclass instance (Typeable m, Typeable (PrimState m), Typeable h)
 deriving stock instance Generic MergePolicyForLevel
 deriving anyclass instance NoThunks MergePolicyForLevel
 
+deriving stock instance Generic NumRuns
+deriving anyclass instance NoThunks NumRuns
+
 {-------------------------------------------------------------------------------
   Entry
 -------------------------------------------------------------------------------}

src/Database/LSMTree/Internal.hs

@@ -68,7 +68,6 @@ import           Control.Monad.Primitive
 import           Control.TempRegistry
 import           Control.Tracer
 import           Data.Arena (ArenaManager, newArenaManager)
-import           Data.Bifunctor (Bifunctor (..))
 import qualified Data.ByteString.Char8 as BSC
 import           Data.Char (isNumber)
 import           Data.Foldable
@@ -90,19 +89,19 @@ import qualified Database.LSMTree.Internal.Entry as Entry
 import           Database.LSMTree.Internal.Lookup (ByteCountDiscrepancy,
                      ResolveSerialisedValue, lookupsIO)
 import           Database.LSMTree.Internal.MergeSchedule
-import           Database.LSMTree.Internal.Paths (RunFsPaths (..),
-                     SessionRoot (..), SnapshotName)
+import           Database.LSMTree.Internal.Paths (SessionRoot (..),
+                     SnapshotName)
 import qualified Database.LSMTree.Internal.Paths as Paths
 import           Database.LSMTree.Internal.Range (Range (..))
 import qualified Database.LSMTree.Internal.RawBytes as RB
 import           Database.LSMTree.Internal.Run (Run)
 import qualified Database.LSMTree.Internal.Run as Run
-import           Database.LSMTree.Internal.RunNumber
 import qualified Database.LSMTree.Internal.RunReader as Reader
 import           Database.LSMTree.Internal.RunReaders (OffsetKey (..))
 import qualified Database.LSMTree.Internal.RunReaders as Readers
 import           Database.LSMTree.Internal.Serialise (SerialisedBlob (..),
                      SerialisedKey, SerialisedValue)
+import           Database.LSMTree.Internal.Snapshot
 import           Database.LSMTree.Internal.UniqCounter
 import qualified Database.LSMTree.Internal.Vector as V
 import qualified Database.LSMTree.Internal.WriteBuffer as WB
@@ -1214,14 +1213,18 @@ snapshot resolve snap label th = do
     traceWith (tableTracer th) $ TraceSnapshot snap
     let conf = tableConfig th
     withOpenTable th $ \thEnv -> do
+      let hfs = tableHasFS thEnv
+      let snapPath = Paths.snapshot (tableSessionRoot thEnv) snap
+      FS.doesFileExist (tableHasFS thEnv) snapPath >>= \b ->
+              when b $ throwIO (ErrSnapshotExists snap)
+
       -- For the temporary implementation it is okay to just flush the buffer
       -- before taking the snapshot.
-      let hfs = tableHasFS thEnv
       content <- modifyWithTempRegistry
                     (RW.unsafeAcquireWriteAccess (tableContent thEnv))
                     (atomically . RW.unsafeReleaseWriteAccess (tableContent thEnv))
                     $ \reg content -> do
-        r <- flushWriteBuffer
+        content' <- flushWriteBuffer
               (TraceMerge `contramap` tableTracer th)
               conf
               resolve
@@ -1231,35 +1234,29 @@ snapshot resolve snap label th = do
               (tableSessionUniqCounter thEnv)
               reg
               content
-        pure (r, r)
+        pure (content', content')
       -- At this point, we've flushed the write buffer but we haven't created the
       -- snapshot file yet. If an asynchronous exception happens beyond this
       -- point, we'll take that loss, as the inner state of the table is still
       -- consistent.
-      runNumbers <- V.forM (tableLevels content) $ \(Level mr rs) -> do
-        (,V.map (runNumber . Run.runRunFsPaths) rs) <$>
-          case mr of
-            SingleRun r -> pure (True, runNumber (Run.runRunFsPaths r))
-            MergingRun _ _ var -> do
-              withMVar var $ \case
-                CompletedMerge r -> pure (False, runNumber (Run.runRunFsPaths r))
-                OngoingMerge{}   -> error "snapshot: OngoingMerge not yet supported" -- TODO: implement
-      let snapPath = Paths.snapshot (tableSessionRoot thEnv) snap
-      FS.doesFileExist (tableHasFS thEnv) snapPath >>= \b ->
-              when b $ throwIO (ErrSnapshotExists snap)
+
+      snappedLevels <- snapLevels (tableLevels content)
+      let snapContents = BSC.pack $ show (label, snappedLevels, tableConfig th)
+
       FS.withFile
         (tableHasFS thEnv)
         snapPath
         (FS.WriteMode FS.MustBeNew) $ \h ->
-          void $ FS.hPutAllStrict (tableHasFS thEnv) h
-                      (BSC.pack $ show (label, runNumbers, tableConfig th))
-      pure $! V.sum (V.map (\((_ :: (Bool, RunNumber)), rs) -> 1 + V.length rs) runNumbers)
+          void $ FS.hPutAllStrict (tableHasFS thEnv) h snapContents
+
+      pure $! numSnapRuns snappedLevels
 
 {-# SPECIALISE open ::
      Session IO h
   -> SnapshotLabel
   -> TableConfigOverride
   -> SnapshotName
+  -> ResolveSerialisedValue
   -> IO (TableHandle IO h) #-}
 -- |  See 'Database.LSMTree.Normal.open'.
 open ::
@@ -1268,8 +1265,9 @@ open ::
   -> SnapshotLabel -- ^ Expected label
   -> TableConfigOverride -- ^ Optional config override
   -> SnapshotName
+  -> ResolveSerialisedValue
   -> m (TableHandle m h)
-open sesh label override snap = do
+open sesh label override snap resolve = do
     traceWith (sessionTracer sesh) $ TraceOpenSnapshot snap override
     withOpenSession sesh $ \seshEnv -> do
       withTempRegistry $ \reg -> do
@@ -1283,26 +1281,17 @@ open sesh label override snap = do
                 snapPath
                 FS.ReadMode $ \h ->
                   FS.hGetAll (sessionHasFS seshEnv) h
-        let (label', runNumbers, conf) =
-                -- why we are using read for this?
-                -- apparently this is a temporary solution, to be done properly in WP15
-                read @(SnapshotLabel, V.Vector ((Bool, RunNumber), V.Vector RunNumber), TableConfig) $
-                BSC.unpack $ BSC.toStrict $ bs
-
-        let conf' = applyOverride override conf
+        let (label', snappedLevels, conf) = read $ BSC.unpack $ BSC.toStrict $ bs
         unless (label == label') $ throwIO (ErrSnapshotWrongType snap)
-        let runPaths = V.map (bimap (second $ RunFsPaths (Paths.activeDir $ sessionRoot seshEnv))
-                                    (V.map (RunFsPaths (Paths.activeDir $ sessionRoot seshEnv))))
-                            runNumbers
-
+        let conf' = applyOverride override conf
         am <- newArenaManager
         blobpath <- Paths.tableBlobPath (sessionRoot seshEnv) <$>
                       incrUniqCounter (sessionUniqCounter seshEnv)
         tableWriteBufferBlobs
           <- allocateTemp reg
                (WBB.new hfs blobpath)
                WBB.removeReference
-        tableLevels <- openLevels reg hfs hbio (confDiskCachePolicy conf') runPaths
+        tableLevels <- openLevels reg hfs hbio conf (sessionRoot seshEnv) resolve snappedLevels
         tableCache <- mkLevelsCache reg tableLevels
         newWith reg sesh seshEnv conf' am $! TableContent {
             tableWriteBuffer = WB.empty
@@ -1311,37 +1300,6 @@ open sesh label override snap = do
           , tableCache
           }
 
-{-# SPECIALISE openLevels ::
-     TempRegistry IO
-  -> HasFS IO h
-  -> HasBlockIO IO h
-  -> DiskCachePolicy
-  -> V.Vector ((Bool, RunFsPaths), V.Vector RunFsPaths)
-  -> IO (Levels IO h) #-}
--- | Open multiple levels.
-openLevels ::
-     (MonadFix m, MonadMask m, MonadMVar m, MonadSTM m, PrimMonad m)
-  => TempRegistry m
-  -> HasFS m h
-  -> HasBlockIO m h
-  -> DiskCachePolicy
-  -> V.Vector ((Bool, RunFsPaths), V.Vector RunFsPaths)
-  -> m (Levels m h)
-openLevels reg hfs hbio diskCachePolicy levels =
-    flip V.imapMStrict levels $ \i (mrPath, rsPaths) -> do
-      let ln      = LevelNo (i+1) -- level 0 is the write buffer
-          caching = diskCachePolicyForLevel diskCachePolicy ln
-      !r <- allocateTemp reg
-              (Run.openFromDisk hfs hbio caching (snd mrPath))
-              Run.removeReference
-      !rs <- flip V.mapMStrict rsPaths $ \run ->
-        allocateTemp reg
-          (Run.openFromDisk hfs hbio caching run)
-          Run.removeReference
-      let !mr = if fst mrPath then SingleRun r
-                              else error "openLevels: OngoingMerge not yet supported"
-      pure $! Level mr rs
-
 {-# SPECIALISE deleteSnapshot ::
      Session IO h
   -> SnapshotName

src/Database/LSMTree/Internal/Config.hs

@@ -80,10 +80,6 @@ instance NFData TableConfig where
   rnf (TableConfig a b c d e f g) =
       rnf a `seq` rnf b `seq` rnf c `seq` rnf d `seq` rnf e `seq` rnf f `seq` rnf g
 
--- | TODO: this should be removed once we have proper snapshotting with proper
--- persistence of the config to disk.
-deriving stock instance Read TableConfig
-
 -- | A reasonable default 'TableConfig'.
 --
 -- This uses a write buffer with up to 20,000 elements and a generous amount of
@@ -169,10 +165,6 @@ data MergePolicy =
 instance NFData MergePolicy where
   rnf MergePolicyLazyLevelling = ()
 
--- | TODO: this should be removed once we have proper snapshotting with proper
--- persistence of the config to disk.
-deriving stock instance Read MergePolicy
-
 {-------------------------------------------------------------------------------
   Size ratio
 -------------------------------------------------------------------------------}
@@ -183,10 +175,6 @@ data SizeRatio = Four
 instance NFData SizeRatio where
   rnf Four = ()
 
--- | TODO: this should be removed once we have proper snapshotting with proper
--- persistence of the config to disk.
-deriving stock instance Read SizeRatio
-
 sizeRatioInt :: SizeRatio -> Int
 sizeRatioInt = \case Four -> 4
 
@@ -214,14 +202,6 @@ data WriteBufferAlloc =
 instance NFData WriteBufferAlloc where
   rnf (AllocNumEntries n) = rnf n
 
--- | TODO: this should be removed once we have proper snapshotting with proper
--- persistence of the config to disk.
-deriving stock instance Read WriteBufferAlloc
-
--- | TODO: this should be removed once we have proper snapshotting with proper
--- persistence of the config to disk.
-deriving stock instance Read NumEntries
-
 {-------------------------------------------------------------------------------
   Bloom filter allocation
 -------------------------------------------------------------------------------}
@@ -263,10 +243,6 @@ instance NFData BloomFilterAlloc where
   rnf (AllocRequestFPR fpr) = rnf fpr
   rnf (AllocMonkey a b)     = rnf a `seq` rnf b
 
--- | TODO: this should be removed once we have proper snapshotting with proper
--- persistence of the config to disk.
-deriving stock instance Read BloomFilterAlloc
-
 defaultBloomFilterAlloc :: BloomFilterAlloc
 defaultBloomFilterAlloc = AllocFixed 10
 
@@ -334,10 +310,6 @@ instance NFData FencePointerIndex where
   rnf CompactIndex  = ()
   rnf OrdinaryIndex = ()
 
--- | TODO: this should be removed once we have proper snapshotting with proper
--- persistence of the config to disk.
-deriving stock instance Read FencePointerIndex
-
 {-------------------------------------------------------------------------------
   Disk cache policy
 -------------------------------------------------------------------------------}

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -16,13 +16,17 @@ module Database.LSMTree.Internal.MergeSchedule (
   , Levels
   , Level (..)
   , MergingRun (..)
+  , NumRuns (..)
   , MergingRunState (..)
+  , NumStepsDone (..)
     -- * Flushes and scheduled merges
   , updatesWithInterleavedFlushes
   , flushWriteBuffer
     -- * Exported for cabal-docspec
   , MergePolicyForLevel (..)
   , maxRunSize
+    -- * Credits
+  , supplyCredits
   ) where
 
 import           Control.Concurrent.Class.MonadMVar.Strict
@@ -38,6 +42,8 @@ import           Control.TempRegistry
 import           Control.Tracer
 import           Data.BloomFilter (Bloom)
 import           Data.Foldable (traverse_)
+import           Data.Primitive (Prim)
+import           Data.Primitive.PrimVar
 import qualified Data.Vector as V
 import           Database.LSMTree.Internal.Assertions (assert,
                      fromIntegralChecked)
@@ -260,12 +266,19 @@ data Level m h = Level {
 
 -- | A merging run is either a single run, or some ongoing merge.
 data MergingRun m h =
-    MergingRun !MergePolicyForLevel !Int !(StrictMVar m (MergingRunState m h))
+    MergingRun !MergePolicyForLevel !NumRuns !(StrictMVar m (MergingRunState m h))
   | SingleRun !(Run m (Handle h))
 
+newtype NumRuns = NumRuns { unNumRuns :: Int }
+  deriving stock (Show, Eq)
+
 data MergingRunState m h =
     CompletedMerge !(Run m (Handle h))
-  | OngoingMerge !(V.Vector (Run m (Handle h))) !(Merge m h)
+  | OngoingMerge !(V.Vector (Run m (Handle h))) !(PrimVar (PrimState m) NumStepsDone) !(Merge m h)
+
+newtype NumStepsDone = NumStepsDone { unNumStepsDone :: Int }
+  deriving stock (Show, Eq)
+  deriving newtype Prim
 
 {-# SPECIALISE addReferenceLevels :: TempRegistry IO -> Levels IO h -> IO () #-}
 addReferenceLevels ::
@@ -305,7 +318,7 @@ forRunAndMergeM_ lvls k1 k2 = V.forM_ lvls $ \(Level mr rs) -> do
       SingleRun r    -> k1 r
       MergingRun _ _ var -> withMVar var $ \case
         CompletedMerge r -> k1 r
-        OngoingMerge irs m -> V.mapM_ k1 irs >> k2 m
+        OngoingMerge irs _ m -> V.mapM_ k1 irs >> k2 m
     V.mapM_ k1 rs
 
 {-# SPECIALISE foldRunM ::
@@ -324,7 +337,7 @@ foldRunM f x lvls = flip (flip V.foldM x) lvls $ \y (Level mr rs) -> do
       SingleRun r -> f y r
       MergingRun _ _ var -> withMVar var $ \case
         CompletedMerge r -> f y r
-        OngoingMerge irs _m -> V.foldM f y irs
+        OngoingMerge irs _ _m -> V.foldM f y irs
     V.foldM f z rs
 
 {-# SPECIALISE forRunM ::
@@ -695,7 +708,7 @@ addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root uc r0 reg levels =
         CompletedMerge r -> do
           traceWith tr $ AtLevel ln $ TraceExpectCompletedMerge (runNumber $ Run.runRunFsPaths r)
           pure r
-        OngoingMerge _rs _m -> error "expectCompletedMerge: OngoingMerge not yet supported" -- TODO: implement.
+        OngoingMerge _rs _ _m -> error "expectCompletedMerge: OngoingMerge not yet supported" -- TODO: implement.
 
     newMerge :: MergePolicyForLevel
              -> Merge.Level
@@ -722,19 +735,20 @@ addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root uc r0 reg levels =
             traceWith tr $ AtLevel ln $ TraceCompletedMerge (Run.runNumEntries r) (runNumber $ Run.runRunFsPaths r)
             V.mapM_ (freeTemp reg . Run.removeReference) rs
             var <- newMVar $! CompletedMerge r
-            pure $! MergingRun mergepolicy (V.length rs) var
+            pure $! MergingRun mergepolicy (NumRuns $ V.length rs) var
           Incremental -> do
             mergeMaybe <- allocateTemp reg
               (Merge.new hfs hbio caching alloc mergelast resolve runPaths rs)
               (traverse_ Merge.removeReference)
             case mergeMaybe of
               Nothing -> error "newMerge: merges can not be empty"
               Just m -> do
-                var <- newMVar $! OngoingMerge rs m
-                pure $! MergingRun mergepolicy (V.length rs) var
+                pvar <- newPrimVar $! NumStepsDone 0
+                var <- newMVar $! OngoingMerge rs pvar m
+                pure $! MergingRun mergepolicy (NumRuns $ V.length rs) var
 
 data MergePolicyForLevel = LevelTiering | LevelLevelling
-  deriving stock Show
+  deriving stock (Show, Eq)
 
 mergePolicyForLevel :: MergePolicy -> LevelNo -> Levels m h -> MergePolicyForLevel
 mergePolicyForLevel MergePolicyLazyLevelling (LevelNo n) nextLevels
@@ -835,9 +849,9 @@ supplyCredits c levels =
     supplyMergeCredits (ceiling (fromIntegral c * cr)) mr
 
 creditsForMerge :: MergingRun m h -> Rational
-creditsForMerge SingleRun{}                         = 0
-creditsForMerge (MergingRun LevelLevelling _ _)     = 1 + 4
-creditsForMerge (MergingRun LevelTiering numRuns _) = fromIntegral numRuns / 4
+creditsForMerge SingleRun{}                             = 0
+creditsForMerge (MergingRun LevelLevelling _ _)         = 1 + 4
+creditsForMerge (MergingRun LevelTiering (NumRuns n) _) = fromIntegral n / 4
 
 {-# SPECIALISE supplyMergeCredits ::
      Credit
@@ -856,13 +870,14 @@ supplyMergeCredits _ SingleRun{} = pure ()
 supplyMergeCredits c (MergingRun _ _ var) = do
     b <- withMVar var $ \case
       CompletedMerge{} -> pure False
-      (OngoingMerge _rs m) -> do
-        (_n, stepResult) <- Merge.steps m c
+      (OngoingMerge _rs pvar m) -> do
+        (n, stepResult) <- Merge.steps m c
+        writePrimVar pvar $! NumStepsDone n
         pure $ stepResult == MergeComplete
     when b $
       modifyMVarMasked_ var $ \case
         mr@CompletedMerge{} -> pure $! mr
-        (OngoingMerge rs m) -> do
+        (OngoingMerge rs _var m) -> do
           RefCount n <- Merge.readRefCount m
           let !n' = fromIntegralChecked n
           V.forM_ rs $ \r -> Run.removeReferenceN r n'