prototype: merge the union level into regular levels

Author: mheinzel

prototypes/ScheduledMerges.hs

@@ -65,6 +65,7 @@ module ScheduledMerges (
     MergeDebt(..),
     NominalCredit(..),
     NominalDebt(..),
+    maxBufferSize,
     Run,
     runSize,
     UnionCredits (..),
@@ -85,6 +86,7 @@ import           Prelude hiding (lookup)
 
 import           Data.Bits
 import           Data.Foldable (for_, toList, traverse_)
+import           Data.Functor ((<&>))
 import           Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import           Data.Maybe (catMaybes)
@@ -344,19 +346,25 @@ invariant (LSMContent _ levels ul) = do
     expectedRunLengths :: Int -> [Run] -> [Level s] -> ST s ()
     expectedRunLengths ln rs ls =
       case mergePolicyForLevel ln ls ul of
-        -- Levels using levelling have only one (incoming) run, which almost
-        -- always consists of an ongoing merge. The exception is when a
-        -- levelling run becomes too large and is promoted, in that case
-        -- initially there's no merge, but it is still represented as an
-        -- 'IncomingRun', using 'Single'. Thus there are no other resident runs.
-        MergePolicyLevelling -> assertST $ null rs
-        -- Runs in tiering levels usually fit that size, but they can be one
-        -- larger, if a run has been held back (creating a 5-way merge).
-        MergePolicyTiering   -> assertST $ all (\r -> tieringRunSizeToLevel r `elem` [ln, ln+1]) rs
-        -- (This is actually still not really true, but will hold in practice.
-        -- In the pathological case, all runs passed to the next level can be
-        -- factor (5/4) too large, and there the same holding back can lead to
-        -- factor (6/4) etc., until at level 12 a run is two levels too large.
+        MergePolicyLevelling ->
+          -- Levels using levelling have only one (incoming) run, which almost
+          -- always consists of an ongoing merge. The exception is when a
+          -- levelling run becomes too large and is promoted, in that case
+          -- initially there's no merge, but it is still represented as an
+          -- 'IncomingRun', using 'Single'. Thus there are no other resident
+          -- runs.
+          assertST $ null rs
+        MergePolicyTiering -> do
+          -- Runs in tiering levels usually fit that size, but they can be one
+          -- larger, if a run has been held back (creating a 5-way merge).
+          --
+          -- TODO: This is actually still not really true, but will hold in
+          -- practice. In the pathological case, all runs passed to the next
+          -- level can be factor (5/4) too large, and there the same holding
+          -- back can lead to factor (6/4) etc., until at level 12 a run is two
+          -- levels too large.
+          assertST $ all (\r -> runSize r > 0) rs
+          assertST $ all (\r -> tieringRunSizeToLevel r `elem` [ln, ln+1]) rs
 
     -- Incoming runs being merged also need to be of the right size, but the
     -- conditions are more complicated.
@@ -367,11 +375,12 @@ invariant (LSMContent _ levels ul) = do
         MergePolicyLevelling -> do
           case (ir, mrs) of
             -- A single incoming run (which thus didn't need merging) must be
-            -- of the expected size range already
+            -- of the expected size range already, but it could also be smaller
+            -- if it comes from a union level.
             (Single r, m) -> do
               assertST $ case m of CompletedMerge{} -> True
                                    OngoingMerge{}   -> False
-              assertST $ levellingRunSizeToLevel r == ln
+              assertST $ levellingRunSizeToLevel r <= ln
 
             -- A completed merge for levelling can be of almost any size at all!
             -- It can be smaller, due to deletions in the last level. But it
@@ -496,6 +505,11 @@ isCompletedMergingTree (MergingTree ref) = do
       OngoingTreeMerge mr  -> isCompletedMergingRun mr
       PendingTreeMerge _   -> failI $ "not completed: PendingTreeMerge"
 
+getCompletedMergingTree :: MergingTree s -> ST s (Maybe Run)
+getCompletedMergingTree t =
+    either (const Nothing) Just
+      <$> evalInvariant (isCompletedMergingTree t)
+
 type Invariant s = E.ExceptT String (ST s)
 
 assertI :: String -> Bool -> Invariant s ()
@@ -781,8 +795,11 @@ update tr (LSMHandle scr lsmr) k op = do
     let wb' = Map.insertWith combine k op wb
     if bufferSize wb' >= maxBufferSize
       then do
-        ls' <- increment tr sc (bufferToRun wb') ls unionLevel
-        let content' = LSMContent Map.empty ls' unionLevel
+        -- Before adding the run to the regular levels, we check if we can get
+        -- rid of the union level (by moving it into into the regular ones).
+        (ls', ul') <- migrateUnionLevel tr sc ls unionLevel
+        ls'' <- increment tr sc (bufferToRun wb') ls' ul'
+        let content' = LSMContent Map.empty ls'' ul'
         invariant content'
         writeSTRef lsmr content'
       else
@@ -1158,9 +1175,44 @@ depositNominalCredit (NominalDebt nominalDebt)
 -- Updates
 --
 
+-- | At some point, we want to merge the union level with the regular levels.
+-- We achieve this by moving it into a new last regular level once it is both
+-- completed (merged down to a single run) and fits into such a new level.
+--
+-- Our representation doesn't allow for empty levels, so we can only put the
+-- run directly after the pre-existing regular levels. If it is too large for
+-- that, we don't want to move it yet to avoid violating run size invariants
+-- and doing inefficient merges of runs with very different sizes.
+migrateUnionLevel :: forall s. Tracer (ST s) Event
+                  -> Counter -> Levels s -> UnionLevel s
+                  -> ST s (Levels s, UnionLevel s)
+migrateUnionLevel _ _ ls NoUnion = do
+    -- nothing to do
+    return (ls, NoUnion)
+migrateUnionLevel _tr _sc ls ul@(Union t _) =
+    -- TODO: tracing
+    getCompletedMergingTree t <&> \case
+      Just r
+        | null r ->
+          -- If the union level is empty, we can just drop it.
+          (ls, NoUnion)
+        | levellingRunSizeToLevel r <= length ls + 1 ->
+          -- If it fits into a hypothetical new last level, put it there.
+          --
+          -- TODO: In some cases it seems desirable to even add it to the
+          -- existing last regular level (so it becomes part of a merge
+          -- sooner), but that would lead to additional merging work that was
+          -- not accounted for. We'd need to be careful to ensure the merge
+          -- completes in time, without doing a lot of work in a short time.
+          (ls ++ [Level (Single r) []], NoUnion)
+      _ ->
+        -- Otherwise, just leave it for now.
+        (ls, ul)
+
 increment :: forall s. Tracer (ST s) Event
-          -> Counter -> Run -> Levels s -> UnionLevel s -> ST s (Levels s)
-increment tr sc run0 ls0 ul = do
+          -> Counter -> Run -> Levels s -> UnionLevel s
+          -> ST s (Levels s)
+increment tr sc run0 ls0 ul =
     go 1 [run0] ls0
   where
     mergeTypeFor :: Levels s -> LevelMergeType

prototypes/ScheduledMergesTest.hs

@@ -24,7 +24,8 @@ tests :: TestTree
 tests = testGroup "Unit and property tests"
     [ testCase "test_regression_empty_run" test_regression_empty_run
     , testCase "test_merge_again_with_incoming" test_merge_again_with_incoming
-    , testProperty "prop_union" prop_union
+    , testProperty "prop_union_supply_all" prop_union_supply_all
+    , testProperty "prop_union_merge_into_levels" prop_union_merge_into_levels
     , testGroup "T"
         [ localOption (QuickCheckTests 1000) $  -- super quick, run more
             testProperty "Arbitrary satisfies invariant" prop_arbitrarySatisfiesInvariant
@@ -176,29 +177,72 @@ test_merge_again_with_incoming =
 -- properties
 --
 
+-- TODO: also generate nested unions?
+
 -- | Supplying enough credits for the remaining debt completes the union merge.
-prop_union :: [[(LSM.Key, LSM.Op)]] -> Property
-prop_union kopss = length (filter (not . null) kopss) > 1 QC.==>
+prop_union_supply_all :: [[(LSM.Key, LSM.Op)]] -> Property
+prop_union_supply_all kopss = length (filter (not . null) kopss) > 1 QC.==>
     QC.ioProperty $ runWithTracer $ \tr ->
       stToIO $ do
         ts <- traverse (mkTable tr) kopss
         t <- LSM.unions ts
 
         debt@(UnionDebt x) <- LSM.remainingUnionDebt t
-        _ <- LSM.supplyUnionCredits t (UnionCredits x)
+        leftovers <- LSM.supplyUnionCredits t (UnionCredits x)
         debt' <- LSM.remainingUnionDebt t
 
         rep <- dumpRepresentation t
         return $ QC.counterexample (show (debt, debt')) $ QC.conjoin
-          [ debt =/= UnionDebt 0
-          , debt' === UnionDebt 0
+          [ QC.counterexample "debt" $ debt =/= UnionDebt 0
+          , QC.counterexample "debt'" $ debt' === UnionDebt 0
+          , QC.counterexample "leftovers" $ leftovers >= 0
           , hasUnionWith isCompleted rep
           ]
   where
     isCompleted = \case
         MLeaf{} -> True
         MNode{} -> False
 
+-- | The union level will get merged into the last regular level once the union
+-- merge is completed and sufficient new entries have been inserted.
+prop_union_merge_into_levels :: [[(LSM.Key, LSM.Op)]] -> Property
+prop_union_merge_into_levels kopss = length (filter (not . null) kopss) > 1 QC.==>
+    QC.forAll arbitrary $ \firstPay ->
+    QC.ioProperty $ runWithTracer $ \tr ->
+      stToIO $ do
+        ts <- traverse (mkTable tr) kopss
+        t <- LSM.unions ts
+
+        -- pay off the union and insert enough that it fits into
+        -- the last level
+        let payOffDebt = do
+              UnionDebt d <- LSM.remainingUnionDebt t
+              _ <- LSM.supplyUnionCredits t (UnionCredits d)
+              return ()
+
+        -- insert as many new entries as there are in the completed
+        -- union level
+        let fillTable = do
+              unionRunSize <- length <$> LSM.logicalValue t
+              LSM.inserts tr t
+                [(K k, V 0, Nothing) | k <- [1 .. unionRunSize]]
+
+        -- we can do these in any order
+        if firstPay
+          then payOffDebt >> fillTable
+          else fillTable >> payOffDebt
+
+        -- then flush the write buffer
+        LSM.inserts tr t
+          [(K k, V 0, Nothing) | k <- [1 .. maxBufferSize]]
+
+        (_, _, mtree) <- representationShape <$> dumpRepresentation t
+
+        -- the union level is gone
+        return $ QC.conjoin
+          [ mtree === Nothing
+          ]
+
 mkTable :: Tracer (ST s) Event -> [(LSM.Key, LSM.Op)] -> ST s (LSM s)
 mkTable tr ks = do
     t <- LSM.new