implement supplyUnionCredits

Author: mheinzel

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

@@ -1483,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/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,7 +105,7 @@ 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,
@@ -1645,19 +1645,38 @@ remainingUnionDebt t = do
 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/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 ::