create MergingRun through safe interface

Initialising a regular MergingRun even for OneShot merging avoid the
the need for directly creating completed MergingRuns and also re-uses
the existing code path to step and complete a Merge, which means there
is one less place to think about resource safety.

To be able to assert that the OneShot merge really completed, as well as
access the resulting run, we make expectCompleted not release the
MergingRun.

Similarly, we restore a completed MergingRun snapshot through a safe
interface. We also remove MergeKnownCompleted from the snapshot format,
which previously allowed to represent inconsistent states. It can
instead be re-constructed from the MergingRunState.

We also don't pass the TempRegistry across the module boundary any more.
The module can use its own where necessary.

Author: mheinzel

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -28,7 +28,6 @@ module Database.LSMTree.Internal.MergeSchedule (
   ) where
 
 import           Control.Concurrent.Class.MonadMVar.Strict
-import           Control.Monad ((<$!>))
 import           Control.Monad.Class.MonadST (MonadST)
 import           Control.Monad.Class.MonadSTM (MonadSTM (..))
 import           Control.Monad.Class.MonadThrow (MonadMask, MonadThrow (..))
@@ -38,7 +37,6 @@ import           Control.TempRegistry
 import           Control.Tracer
 import           Data.BloomFilter (Bloom)
 import           Data.Foldable (fold)
-import           Data.Primitive.PrimVar
 import qualified Data.Vector as V
 import           Database.LSMTree.Internal.Assertions (assert)
 import           Database.LSMTree.Internal.Config
@@ -93,9 +91,10 @@ data MergeTrace =
       RunBloomFilterAlloc
       MergePolicyForLevel
       Merge.Level
-  | TraceCompletedMerge
+  | TraceCompletedMerge  -- TODO: currently not traced for Incremental merges
       NumEntries -- ^ Size of output run
       RunNumber
+    -- | This is traced at the latest point the merge could complete.
   | TraceExpectCompletedMerge
       RunNumber
   | TraceNewMergeSingleRun
@@ -181,7 +180,7 @@ mkLevelsCache ::
 mkLevelsCache reg lvls = do
     rs <- foldRunAndMergeM
       (fmap V.singleton . dupRun)
-      (MR.duplicateRuns reg)
+      (\mr -> allocateTemp reg (MR.duplicateRuns mr) (V.mapM_ releaseRef))
       lvls
     pure $! LevelsCache_ {
         cachedRuns      = rs
@@ -698,18 +697,21 @@ addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root uc r0 reg levels =
     expectCompletedMerge ln ir = do
       r <- case ir of
         Single r   -> pure r
-        Merging mr -> MR.expectCompleted reg mr
+        Merging mr -> do
+          r <- allocateTemp reg (MR.expectCompleted mr) releaseRef
+          freeTemp reg (releaseRef mr)
+          pure r
       traceWith tr $ AtLevel ln $
         TraceExpectCompletedMerge (Run.runFsPathsNumber r)
       pure r
 
-    -- TODO: refactor, pull to top level?
+    -- Takes ownership of the runs passed.
     newMerge :: MergePolicyForLevel
              -> Merge.Level
              -> LevelNo
              -> V.Vector (Ref (Run m h))
              -> m (IncomingRun m h)
-    newMerge mergePolicy mergelast ln rs
+    newMerge mergePolicy mergeLevel ln rs
       | Just (r, rest) <- V.uncons rs
       , V.null rest = do
           traceWith tr $ AtLevel ln $
@@ -723,29 +725,33 @@ addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root uc r0 reg levels =
             !alloc = bloomFilterAllocForLevel conf ln
             !runPaths = Paths.runPath root (uniqueToRunNumber n)
         traceWith tr $ AtLevel ln $
-          TraceNewMerge (V.map Run.size rs) (runNumber runPaths) caching alloc mergePolicy mergelast
-        let numInputRuns = NumRuns $ V.length rs
-        let numInputEntries = V.foldMap' Run.size rs
+          TraceNewMerge (V.map Run.size rs) (runNumber runPaths) caching alloc mergePolicy mergeLevel
+        -- TODO: There currently is a resource management bug that happens if an
+        -- exception occurs after calling MR.new. In this case, all changes roll
+        -- back, so some of the runs in rs will live in the Levels structure at
+        -- their original places again. However, we passed their references to
+        -- the MergingRun, which gets aborted, releasing the run references.
+        -- Instead of passing the original references into newMerge, we have to
+        -- duplicate the ones that previously existed in the level and then
+        -- freeTemp the original ones. This way, on the happy path the result is
+        -- the same, but if an exception occurs, the original references do not
+        -- get released.
+        mr <- allocateTemp reg
+          (MR.new hfs hbio resolve caching alloc mergeLevel mergePolicy runPaths rs)
+          releaseRef
         case confMergeSchedule of
+          Incremental -> pure ()
           OneShot -> do
-            r <- allocateTemp reg
-                  (mergeRuns resolve hfs hbio caching alloc runPaths mergelast rs)
-                  releaseRef
-            traceWith tr $ AtLevel ln $
-              TraceCompletedMerge (Run.size r)
-                                  (Run.runFsPathsNumber r)
-            V.mapM_ (freeTemp reg . releaseRef) rs
-            Merging <$!> MR.unsafeNew mergePolicy numInputRuns numInputEntries MR.MergeKnownCompleted (MR.CompletedMerge r)
-
-          Incremental -> do
-            mergeMaybe <- allocateMaybeTemp reg
-              (Merge.new hfs hbio caching alloc mergelast resolve runPaths rs)
-              Merge.abort
-            case mergeMaybe of
-              Nothing -> error "newMerge: merges can not be empty"
-              Just m -> do
-                spentCreditsVar <- MR.SpentCreditsVar <$> newPrimVar 0
-                Merging <$!> MR.unsafeNew mergePolicy numInputRuns numInputEntries MR.MergeMaybeCompleted (MR.OngoingMerge rs spentCreditsVar m)
+            let !required = MR.Credits (unNumEntries (V.foldMap' Run.size rs))
+            let !thresh = creditThresholdForLevel conf ln
+            MR.supplyCredits required thresh mr
+            -- This ensures the merge is really completed. However, we don't
+            -- release the merge yet and only briefly inspect the resulting run.
+            bracket (MR.expectCompleted mr) releaseRef $ \r ->
+              traceWith tr $ AtLevel ln $
+                TraceCompletedMerge (Run.size r) (Run.runFsPathsNumber r)
+
+        return (Merging mr)
 
 -- $setup
 -- >>> import Database.LSMTree.Internal.Entry
@@ -795,23 +801,6 @@ mergeLastForLevel levels
 levelIsFull :: SizeRatio -> V.Vector run -> Bool
 levelIsFull sr rs = V.length rs + 1 >= (sizeRatioInt sr)
 
-{-# SPECIALISE mergeRuns :: ResolveSerialisedValue -> HasFS IO h -> HasBlockIO IO h -> RunDataCaching -> RunBloomFilterAlloc -> RunFsPaths -> Merge.Level -> V.Vector (Ref (Run IO h)) -> IO (Ref (Run IO h)) #-}
-mergeRuns ::
-     (MonadMask m, MonadST m, MonadSTM m)
-  => ResolveSerialisedValue
-  -> HasFS m h
-  -> HasBlockIO m h
-  -> RunDataCaching
-  -> RunBloomFilterAlloc
-  -> RunFsPaths
-  -> Merge.Level
-  -> V.Vector (Ref (Run m h))
-  -> m (Ref (Run m h))
-mergeRuns resolve hfs hbio caching alloc runPaths mergeLevel runs = do
-    Merge.new hfs hbio caching alloc mergeLevel resolve runPaths runs >>= \case
-      Nothing -> error "mergeRuns: no inputs"
-      Just m -> Merge.stepsToCompletion m 1024
-
 {-------------------------------------------------------------------------------
   Credits
 -------------------------------------------------------------------------------}

src/Database/LSMTree/Internal/MergingRun.hs

@@ -5,7 +5,8 @@
 -- | An incremental merge of multiple runs.
 module Database.LSMTree.Internal.MergingRun (
     MergingRun (..)
-  , unsafeNew
+  , new
+  , newCompleted
   , duplicateRuns
   , supplyCredits
   , expectCompleted
@@ -32,14 +33,21 @@ import           Control.Monad.Class.MonadThrow (MonadCatch (bracketOnError),
 import           Control.Monad.Primitive
 import           Control.RefCount
 import           Control.TempRegistry
+import           Data.Maybe (fromMaybe)
 import           Data.Primitive.MutVar
 import           Data.Primitive.PrimVar
 import qualified Data.Vector as V
 import           Database.LSMTree.Internal.Assertions (assert)
 import           Database.LSMTree.Internal.Entry (NumEntries (..), unNumEntries)
+import           Database.LSMTree.Internal.Lookup (ResolveSerialisedValue)
 import           Database.LSMTree.Internal.Merge (Merge, StepResult (..))
 import qualified Database.LSMTree.Internal.Merge as Merge
+import           Database.LSMTree.Internal.Paths (RunFsPaths (..))
 import           Database.LSMTree.Internal.Run (Run)
+import qualified Database.LSMTree.Internal.Run as Run
+import           Database.LSMTree.Internal.RunAcc (RunBloomFilterAlloc)
+import           System.FS.API (HasFS)
+import           System.FS.BlockIO.API (HasBlockIO)
 
 data MergingRun m h = MergingRun {
       mergePolicy         :: !MergePolicyForLevel
@@ -103,16 +111,65 @@ instance NFData MergeKnownCompleted where
   rnf MergeKnownCompleted = ()
   rnf MergeMaybeCompleted = ()
 
-{-# SPECIALISE unsafeNew ::
+{-# SPECIALISE new ::
+     HasFS IO h
+  -> HasBlockIO IO h
+  -> ResolveSerialisedValue
+  -> Run.RunDataCaching
+  -> RunBloomFilterAlloc
+  -> Merge.Level
+  -> MergePolicyForLevel
+  -> RunFsPaths
+  -> V.Vector (Ref (Run IO h))
+  -> IO (Ref (MergingRun IO h)) #-}
+-- | Create a new merging run, returning a reference to it that must ultimately
+-- be released via 'releaseRef'.
+--
+-- Takes over ownership of the references to the runs passed.
+--
+-- This function should be run with asynchronous exceptions masked to prevent
+-- failing after internal resources have already been created.
+new ::
+     (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)
+  => HasFS m h
+  -> HasBlockIO m h
+  -> ResolveSerialisedValue
+  -> Run.RunDataCaching
+  -> RunBloomFilterAlloc
+  -> Merge.Level
+  -> MergePolicyForLevel
+  -> RunFsPaths
+  -> V.Vector (Ref (Run m h))
+  -> m (Ref (MergingRun m h))
+new hfs hbio resolve caching alloc mergeLevel mergePolicy runPaths runs = do
+    merge <- fromMaybe (error "newMerge: merges can not be empty")
+      <$> Merge.new hfs hbio caching alloc mergeLevel resolve runPaths runs
+    let numInputRuns = NumRuns $ V.length runs
+    let numInputEntries = V.foldMap' Run.size runs
+    spentCreditsVar <- SpentCreditsVar <$> newPrimVar 0
+    unsafeNew mergePolicy numInputRuns numInputEntries MergeMaybeCompleted $
+      OngoingMerge runs spentCreditsVar merge
+
+{-# SPECIALISE newCompleted ::
      MergePolicyForLevel
   -> NumRuns
   -> NumEntries
-  -> MergeKnownCompleted
-  -> MergingRunState IO h
+  -> Ref (Run IO h)
   -> IO (Ref (MergingRun IO h)) #-}
--- | This allows constructing ill-formed MergingRuns, but the flexibility is
--- needed for creating a merging run that is already Completed, as well as
--- opening a merging run from a snapshot.
+-- | Create a merging run that is already in the completed state, returning a
+-- reference that must ultimately be released via 'releaseRef'.
+newCompleted ::
+     (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)
+  => MergePolicyForLevel
+  -> NumRuns
+  -> NumEntries
+  -> Ref (Run m h)
+  -> m (Ref (MergingRun m h))
+newCompleted mergePolicy numInputRuns numInputEntries run = do
+    unsafeNew mergePolicy numInputRuns numInputEntries MergeKnownCompleted $
+      CompletedMerge run
+
+{-# INLINE unsafeNew #-}
 unsafeNew ::
      (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)
   => MergePolicyForLevel
@@ -150,20 +207,18 @@ unsafeNew mergePolicy mergeNumRuns mergeNumEntries knownCompleted state = do
 
 -- | Create references to the runs that should be queried for lookups.
 -- In particular, if the merge is not complete, these are the input runs.
-{-# SPECIALISE duplicateRuns :: TempRegistry IO -> Ref (MergingRun IO h) -> IO (V.Vector (Ref (Run IO h))) #-}
+{-# SPECIALISE duplicateRuns :: Ref (MergingRun IO h) -> IO (V.Vector (Ref (Run IO h))) #-}
 duplicateRuns ::
      (PrimMonad m, MonadMVar m, MonadMask m)
-  => TempRegistry m
-  -> Ref (MergingRun m h)
+  => Ref (MergingRun m h)
   -> m (V.Vector (Ref (Run m h)))
-duplicateRuns reg (DeRef mr) =
+duplicateRuns (DeRef mr) =
     -- We take the references while holding the MVar to make sure the MergingRun
     -- does not get completed concurrently before we are done.
     withMVar (mergeState mr) $ \case
-      CompletedMerge r    -> V.singleton <$> dupRun r
-      OngoingMerge rs _ _ -> V.mapM dupRun rs
-  where
-    dupRun r = allocateTemp reg (dupRef r) releaseRef
+      CompletedMerge r    -> V.singleton <$> dupRef r
+      OngoingMerge rs _ _ -> withTempRegistry $ \reg ->
+        V.mapM (\r -> allocateTemp reg (dupRef r) releaseRef) rs
 
 {-------------------------------------------------------------------------------
   Credits
@@ -331,19 +386,19 @@ tryTakeUnspentCredits ::
 tryTakeUnspentCredits
     unspentCreditsVar@(UnspentCreditsVar !var)
     thresh@(CreditThreshold !creditsThresh)
-    (Credits !prev)
-  | prev < creditsThresh = pure Nothing
+    (Credits !before)
+  | before < creditsThresh = pure Nothing
   | otherwise = do
       -- numThresholds is guaranteed to be >= 1
-      let !numThresholds = prev `div` creditsThresh
+      let !numThresholds = before `div` creditsThresh
           !creditsToTake = numThresholds * creditsThresh
-          !new = prev - creditsToTake
-      assert (new < creditsThresh) $ pure ()
-      prev' <- casInt var prev new
-      if prev' == prev then
+          !after = before - creditsToTake
+      assert (after < creditsThresh) $ pure ()
+      before' <- casInt var before after
+      if before' == before then
         pure (Just (Credits creditsToTake))
       else
-        tryTakeUnspentCredits unspentCreditsVar thresh (Credits prev')
+        tryTakeUnspentCredits unspentCreditsVar thresh (Credits before')
 
 {-# SPECIALISE putBackUnspentCredits ::
      UnspentCreditsVar RealWorld
@@ -446,13 +501,14 @@ completeMerge mergeVar mergeKnownCompletedVar = do
         pure $! CompletedMerge r
 
 {-# SPECIALISE expectCompleted ::
-     TempRegistry IO
-  -> Ref (MergingRun IO h)
+     Ref (MergingRun IO h)
   -> IO (Ref (Run IO h)) #-}
+-- | This does /not/ release the reference, but allocates a new reference for
+-- the returned run, which must be released at some point.
 expectCompleted ::
      (MonadMVar m, MonadSTM m, MonadST m, MonadMask m)
-  => TempRegistry m -> Ref (MergingRun m h) -> m (Ref (Run m h))
-expectCompleted reg mr@(DeRef MergingRun {..}) = do
+  => Ref (MergingRun m h) -> m (Ref (Run m h))
+expectCompleted (DeRef MergingRun {..}) = do
     knownCompleted <- readMutVar mergeKnownCompleted
     -- The merge is not guaranteed to be complete, so we do the remaining steps
     when (knownCompleted == MergeMaybeCompleted) $ do
@@ -462,13 +518,9 @@ expectCompleted reg mr@(DeRef MergingRun {..}) = do
       when isMergeDone $ completeMerge mergeState mergeKnownCompleted
       -- TODO: can we think of a check to see if we did not do too much work
       -- here?
-    r <- withMVar mergeState $ \case
-      CompletedMerge r -> pure r
+    withMVar mergeState $ \case
+      CompletedMerge r -> dupRef r  -- return a fresh reference to the run
       OngoingMerge{} -> do
         -- If the algorithm finds an ongoing merge here, then it is a bug in
         -- our merge sceduling algorithm. As such, we throw a pure error.
         error "expectCompleted: expected a completed merge, but found an ongoing merge"
-    -- return a fresh reference to the run
-    r' <- allocateTemp reg (dupRef r) releaseRef
-    freeTemp reg (releaseRef mr)
-    pure r'