Merge pull request #633 from IntersectMBO/mheinzel/union-debt

Implement remainingUnionDebt and supplyUnionCredits

Author: jorisdral

bench/macro/lsm-tree-bench-lookups.hs

@@ -377,7 +377,7 @@ lookupsEnv runSizes keyRng0 hfs hbio caching = do
     putStr "DONE"
 
     -- return runs
-    runs <- V.fromList <$> mapM Run.fromMutable rbs
+    runs <- V.fromList <$> mapM Run.fromBuilder rbs
     let blooms  = V.map (\(DeRef r) -> Run.runFilter   r) runs
         indexes = V.map (\(DeRef r) -> Run.runIndex    r) runs
         handles = V.map (\(DeRef r) -> Run.runKOpsFile r) runs

bench/micro/Bench/Database/LSMTree/Internal/Index.hs

@@ -7,8 +7,9 @@ import           Control.DeepSeq (rnf)
 import           Control.Monad.ST.Strict (runST)
 import           Criterion.Main (Benchmark, Benchmarkable, bench, bgroup, env,
                      whnf)
-#if __GLASGOW_HASKELL__ < 910
+#if !MIN_VERSION_base(4,20,0)
 import           Data.List (foldl')
+                 -- foldl' is included in the Prelude from base 4.20 onwards
 #endif
 import           Database.LSMTree.Extras.Generators (getKeyForIndexCompact,
                      mkPages, toAppends)

prototypes/ScheduledMerges.hs

@@ -1371,8 +1371,17 @@ remainingDebtMergingTree :: MergingTree s -> ST s (Debt, Size)
 remainingDebtMergingTree (MergingTree ref) =
     readSTRef ref >>= \case
       CompletedTreeMerge r -> return (0, runSize r)
-      OngoingTreeMerge mr  -> remainingDebtMergingRun mr
-      PendingTreeMerge pm  -> remainingDebtPendingMerge pm
+      OngoingTreeMerge mr  -> addDebtOne <$> remainingDebtMergingRun mr
+      PendingTreeMerge pm  -> addDebtOne <$> remainingDebtPendingMerge pm
+  where
+    -- An ongoing merge should never have 0 debt, even if the 'MergingRun' in it
+    -- says it is completed. We still need to update it to 'CompletedTreeMerge'.
+    -- Similarly, a pending merge needs some work to complete it, even if all
+    -- its inputs are empty.
+    --
+    -- Note that we can't use @max 1@, as this would violate the property that
+    -- supplying N credits reduces the remaining debt by at least N.
+    addDebtOne (debt, size) = (debt + 1, size)
 
 remainingDebtPendingMerge :: PendingMerge s -> ST s (Debt, Size)
 remainingDebtPendingMerge (PendingMerge _ prs trees) = do
@@ -1381,11 +1390,7 @@ remainingDebtPendingMerge (PendingMerge _ prs trees) = do
         , traverse remainingDebtMergingTree trees
         ]
     let totalSize = sum sizes
-    -- A pending merge should never have 0 remaining debt. It needs some work to
-    -- complete it, even if all its inputs are empty. It's not enought to use
-    -- @max 1@, as this would violate the property that supplying N credits
-    -- reduces the remaining debt by at least N.
-    let totalDebt = sum debts + totalSize + 1
+    let totalDebt = sum debts + totalSize
     return (totalDebt, totalSize)
   where
     remainingDebtPreExistingRun = \case
@@ -1478,8 +1483,9 @@ supplyCreditsPendingMerge = checked remainingDebtPendingMerge $ \credits -> \cas
     -- approximately equal, being more precise would require more iterations
     splitEqually :: (Credit -> a -> ST s Credit) -> [a] -> Credit -> ST s Credit
     splitEqually f xs credits =
-        -- first give each tree k = ceil(1/n) credits (last ones might get less)
-        -- any remainders go left to right
+        -- first give each tree k = ceil(1/n) credits (last ones might get less).
+        -- it's important we fold here to collect leftovers.
+        -- any remainders go left to right.
         foldM supply credits xs >>= leftToRight f xs
       where
         !n = length xs

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

@@ -149,6 +149,9 @@ unsafeCreateRunAt ::
   -> SerialisedRunData
   -> IO (Ref (Run IO h))
 unsafeCreateRunAt fs hbio runParams fsPaths (RunData m) = do
+    -- the WBB file path doesn't have to be at a specific place relative to
+    -- the run we want to create, but fsPaths should already point to a unique
+    -- location, so we just append something to not conflict with that.
     let blobpath = FS.addExtension (runBlobPath fsPaths) ".wb"
     bracket (WBB.new fs blobpath) releaseRef $ \wbblobs -> do
       wb <- WB.fromMap <$> traverse (traverse (WBB.addBlob fs wbblobs)) m

src/Database/LSMTree.hs

@@ -564,15 +564,19 @@ remainingUnionDebt (Internal.Table' t) =
     (\(Internal.UnionDebt x) -> UnionDebt x) <$>
       Internal.remainingUnionDebt t
 
-{-# SPECIALISE supplyUnionCredits :: Table IO k v b -> UnionCredits -> IO UnionCredits #-}
+{-# SPECIALISE supplyUnionCredits ::
+     ResolveValue v => Table IO k v b -> UnionCredits -> IO UnionCredits #-}
 supplyUnionCredits ::
-     IOLike m
+     forall m k v b. (IOLike m, ResolveValue v)
   => Table m k v b
   -> UnionCredits
   -> m UnionCredits
 supplyUnionCredits (Internal.Table' t) (UnionCredits credits) =
     (\(Internal.UnionCredits x) -> UnionCredits x) <$>
-      Internal.supplyUnionCredits t (Internal.UnionCredits credits)
+      Internal.supplyUnionCredits
+        (resolve (Proxy @v))
+        t
+        (Internal.UnionCredits credits)
 
 {-------------------------------------------------------------------------------
   Monoidal value resolution

src/Database/LSMTree/Internal.hs

@@ -105,14 +105,15 @@ import           Database.LSMTree.Internal.BlobRef (WeakBlobRef (..))
 import qualified Database.LSMTree.Internal.BlobRef as BlobRef
 import           Database.LSMTree.Internal.Config
 import qualified Database.LSMTree.Internal.Cursor as Cursor
-import           Database.LSMTree.Internal.Entry (Entry)
+import           Database.LSMTree.Internal.Entry (Entry, NumEntries (..))
 import           Database.LSMTree.Internal.IncomingRun (IncomingRun (..))
 import           Database.LSMTree.Internal.Lookup (ByteCountDiscrepancy,
                      ResolveSerialisedValue, lookupsIO,
                      lookupsIOWithoutWriteBuffer)
 import           Database.LSMTree.Internal.MergeSchedule
 import qualified Database.LSMTree.Internal.MergingRun as MR
 import           Database.LSMTree.Internal.MergingTree
+import qualified Database.LSMTree.Internal.MergingTree as MT
 import qualified Database.LSMTree.Internal.MergingTree.Lookup as MT
 import           Database.LSMTree.Internal.Paths (SessionRoot (..),
                      SnapshotMetaDataChecksumFile (..),
@@ -1626,32 +1627,56 @@ newtype UnionDebt = UnionDebt Int
 
 {-# SPECIALISE remainingUnionDebt :: Table IO h -> IO UnionDebt #-}
 -- | See 'Database.LSMTree.Normal.remainingUnionDebt'.
-remainingUnionDebt :: (MonadSTM m, MonadThrow m) => Table m h -> m UnionDebt
+remainingUnionDebt ::
+     (MonadSTM m, MonadMVar m, MonadThrow m, PrimMonad m)
+  => Table m h -> m UnionDebt
 remainingUnionDebt t = do
     traceWith (tableTracer t) TraceRemainingUnionDebt
     withOpenTable t $ \tEnv -> do
-      RW.withReadAccess (tableContent tEnv) $ \tableContent ->
+      RW.withReadAccess (tableContent tEnv) $ \tableContent -> do
         case tableUnionLevel tableContent of
-          NoUnion -> pure (UnionDebt 0)
-          Union{} -> error "remainingUnionDebt: not yet implemented"
+          NoUnion ->
+            pure (UnionDebt 0)
+          Union mt -> do
+            (MergeDebt (MergeCredits c), _) <- MT.remainingMergeDebt mt
+            pure (UnionDebt c)
 
 -- | See 'Database.LSMTree.Normal.UnionCredits'.
 newtype UnionCredits = UnionCredits Int
   deriving newtype (Show, Eq, Ord, Num)
 
-{-# SPECIALISE supplyUnionCredits :: Table IO h -> UnionCredits -> IO UnionCredits #-}
+{-# SPECIALISE supplyUnionCredits ::
+     ResolveSerialisedValue -> Table IO h -> UnionCredits -> IO UnionCredits #-}
 -- | See 'Database.LSMTree.Normal.supplyUnionCredits'.
-supplyUnionCredits :: (MonadSTM m, MonadCatch m) => Table m h -> UnionCredits -> m UnionCredits
-supplyUnionCredits t credits = do
+supplyUnionCredits ::
+     (MonadST m, MonadSTM m, MonadMVar m, MonadMask m)
+  => ResolveSerialisedValue -> Table m h -> UnionCredits -> m UnionCredits
+supplyUnionCredits resolve t credits = do
     traceWith (tableTracer t) $ TraceSupplyUnionCredits credits
     withOpenTable t $ \tEnv -> do
-      -- TODO: should this be acquiring read or write access?
-      RW.withWriteAccess (tableContent tEnv) $ \tableContent ->
+      -- No need to mutate the table content here. In the rare case that we want
+      -- to move a completed union level into the regular levels, we can still
+      -- take the write lock for that.
+      RW.withReadAccess (tableContent tEnv) $ \tableContent -> do
         case tableUnionLevel tableContent of
-          NoUnion -> pure (tableContent, credits) -- all leftovers
-          Union{}
-           | credits <= UnionCredits 0 -> pure (tableContent, UnionCredits 0)
-           --TODO: remove this 0 special case once the general case covers it.
-           -- We do not need to optimise the 0 case. It is just here to
-           -- simplify test coverage.
-           | otherwise -> error "supplyUnionCredits: not yet implemented"
+          NoUnion ->
+            pure (max 0 credits)  -- all leftovers (but never negative)
+          Union mt -> do
+            let conf = tableConfig t
+            let AllocNumEntries (NumEntries x) = confWriteBufferAlloc conf
+            -- We simply use the write buffer size as merge credit threshold, as
+            -- the regular level merges also do.
+            -- TODO: pick a more suitable threshold or make configurable?
+            let thresh = MR.CreditThreshold (MR.UnspentCredits (MergeCredits x))
+            MergeCredits leftovers <-
+              MT.supplyCredits
+                (tableHasFS tEnv)
+                (tableHasBlockIO tEnv)
+                resolve
+                (runParamsForLevel conf UnionLevel)
+                thresh
+                (tableSessionRoot tEnv)
+                (tableSessionUniqCounter tEnv)
+                mt
+                (let UnionCredits c = credits in MergeCredits c)
+            pure (UnionCredits leftovers)

src/Database/LSMTree/Internal/Config.hs

@@ -3,6 +3,8 @@ module Database.LSMTree.Internal.Config (
     -- * Table configuration
   , TableConfig (..)
   , defaultTableConfig
+  , RunLevelNo (..)
+  , runParamsForLevel
     -- * Table configuration override
   , TableConfigOverride
   , applyOverride
@@ -34,16 +36,16 @@ import           Control.DeepSeq (NFData (..))
 import           Data.Maybe (fromMaybe)
 import           Data.Monoid (Last (..))
 import           Data.Word (Word64)
-import           Database.LSMTree.Internal.Assertions (assert)
 import           Database.LSMTree.Internal.Entry (NumEntries (..))
 import           Database.LSMTree.Internal.Index (IndexType)
 import qualified Database.LSMTree.Internal.Index as Index
                      (IndexType (Compact, Ordinary))
 import           Database.LSMTree.Internal.Run (RunDataCaching (..))
 import           Database.LSMTree.Internal.RunAcc (RunBloomFilterAlloc (..))
+import           Database.LSMTree.Internal.RunBuilder (RunParams (..))
 
 newtype LevelNo = LevelNo Int
-  deriving stock (Show, Eq)
+  deriving stock (Show, Eq, Ord)
   deriving newtype (Enum, NFData)
 
 {-------------------------------------------------------------------------------
@@ -95,6 +97,16 @@ defaultTableConfig =
       , confMergeSchedule     = defaultMergeSchedule
       }
 
+data RunLevelNo = RegularLevel LevelNo | UnionLevel
+
+runParamsForLevel :: TableConfig -> RunLevelNo -> RunParams
+runParamsForLevel conf@TableConfig {..} levelNo =
+    RunParams
+      { runParamCaching = diskCachePolicyForLevel confDiskCachePolicy levelNo
+      , runParamAlloc   = bloomFilterAllocForLevel conf levelNo
+      , runParamIndex   = indexTypeForRun confFencePointerIndex
+      }
+
 {-------------------------------------------------------------------------------
   Table configuration override
 -------------------------------------------------------------------------------}
@@ -227,9 +239,8 @@ instance NFData BloomFilterAlloc where
 defaultBloomFilterAlloc :: BloomFilterAlloc
 defaultBloomFilterAlloc = AllocFixed 10
 
-bloomFilterAllocForLevel :: TableConfig -> LevelNo -> RunBloomFilterAlloc
-bloomFilterAllocForLevel conf (LevelNo l) =
-    assert (l > 0) $
+bloomFilterAllocForLevel :: TableConfig -> RunLevelNo -> RunBloomFilterAlloc
+bloomFilterAllocForLevel conf _levelNo =
     case confBloomFilterAlloc conf of
       AllocFixed n        -> RunAllocFixed n
       AllocRequestFPR fpr -> RunAllocRequestFPR fpr
@@ -326,15 +337,16 @@ instance NFData DiskCachePolicy where
 -- | Interpret the 'DiskCachePolicy' for a level: should we cache data in runs
 -- at this level.
 --
-diskCachePolicyForLevel :: DiskCachePolicy -> LevelNo -> RunDataCaching
-diskCachePolicyForLevel policy (LevelNo ln) =
+diskCachePolicyForLevel :: DiskCachePolicy -> RunLevelNo -> RunDataCaching
+diskCachePolicyForLevel policy levelNo =
   case policy of
-    DiskCacheAll               -> CacheRunData
-    DiskCacheNone              -> NoCacheRunData
-    DiskCacheLevelsAtOrBelow n
-      | ln <= n                -> CacheRunData
-      | otherwise              -> NoCacheRunData
-
+    DiskCacheAll  -> CacheRunData
+    DiskCacheNone -> NoCacheRunData
+    DiskCacheLevelsAtOrBelow n ->
+      case levelNo of
+        RegularLevel l | l <= LevelNo n -> CacheRunData
+                       | otherwise      -> NoCacheRunData
+        UnionLevel                      -> NoCacheRunData
 
 {-------------------------------------------------------------------------------
   Merge schedule

src/Database/LSMTree/Internal/Merge.hs

@@ -174,6 +174,9 @@ new ::
 new hfs hbio runParams mergeType mergeMappend targetPaths runs = do
     -- no offset, no write buffer
     mreaders <- Readers.new Readers.NoOffsetKey Nothing runs
+    -- TODO: Exception safety! If Readers.new fails after already creating some
+    -- run readers, or Builder.new fails, the run readers will stay open,
+    -- holding handles of the input runs' files.
     for mreaders $ \mergeReaders -> do
       -- calculate upper bounds based on input runs
       let numEntries = V.foldMap' Run.size runs
@@ -236,7 +239,7 @@ complete Merge{..} = do
       Merging -> error "complete: Merge is not done"
       MergingDone -> do
         -- the readers are already drained, therefore closed
-        r <- Run.fromMutable mergeBuilder
+        r <- Run.fromBuilder mergeBuilder
         writeMutVar mergeState $! Completed
         pure r
       Completed -> error "complete: Merge is already completed"

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -762,18 +762,10 @@ mergingRunParamsForLevel ::
   -> Unique
   -> LevelNo
   -> (RunParams, RunFsPaths)
-mergingRunParamsForLevel dir
-                         conf@TableConfig {
-                                confDiskCachePolicy,
-                                confFencePointerIndex
-                         }
-                         unique ln =
-    (RunParams {..}, runPaths)
+mergingRunParamsForLevel dir conf unique ln =
+    (runParamsForLevel conf (RegularLevel ln), runPaths)
   where
-    !runParamCaching = diskCachePolicyForLevel confDiskCachePolicy ln
-    !runParamAlloc   = bloomFilterAllocForLevel conf ln
-    !runParamIndex   = indexTypeForRun confFencePointerIndex
-    !runPaths        = Paths.runPath dir (uniqueToRunNumber unique)
+    !runPaths = Paths.runPath dir (uniqueToRunNumber unique)
 
 -- | We use levelling on the last level, unless that is also the first level.
 mergePolicyForLevel ::