Merge pull request #464 from IntersectMBO/mheinzel/prototype-preparation

Prototype improvements, support mupsert

Author: mheinzel

lsm-tree.cabal

@@ -689,6 +689,7 @@ library prototypes
   exposed-modules:
     FormatPage
     ScheduledMerges
+    ScheduledMergesTest
     ScheduledMergesTestQLS
 
   build-depends:
@@ -698,7 +699,6 @@ library prototypes
     , constraints
     , containers
     , contra-tracer
-    , lsm-tree
     , QuickCheck
     , quickcheck-dynamic
     , quickcheck-lockstep

prototypes/ScheduledMerges.hs

@@ -19,14 +19,15 @@
 module ScheduledMerges (
     -- * Main API
     LSM,
-    Key, Value, Blob,
+    Key (K), Value (V), resolveValue, Blob (B),
     new,
     LookupResult (..),
     lookup, lookups,
     Update (..),
     update, updates,
     insert, inserts,
     delete, deletes,
+    mupsert, mupserts,
     supply,
     duplicate,
 
@@ -52,8 +53,6 @@ import           Control.Monad.ST
 import           Control.Tracer (Tracer, contramap, traceWith)
 import           GHC.Stack (HasCallStack, callStack)
 
-import           Database.LSMTree.Normal (LookupResult (..), Update (..))
-
 
 data LSM s  = LSMHandle !(STRef s Counter)
                         !(STRef s (LSMContent s))
@@ -122,12 +121,20 @@ runSize = Map.size
 bufferSize :: Buffer -> Int
 bufferSize = Map.size
 
-type Op     = Update Value Blob
+type Op = Update Value Blob
+
+newtype Key = K Int
+  deriving stock (Eq, Ord, Show)
+  deriving newtype Enum
+
+newtype Value  = V Int
+  deriving stock (Eq, Show)
 
-type Key    = Int
-type Value  = Int
-type Blob   = Int
+resolveValue :: Value -> Value -> Value
+resolveValue (V x) (V y) = V (x + y)
 
+newtype Blob = B Int
+  deriving stock (Eq, Show)
 
 -- | The size of the 4 tiering runs at each level are allowed to be:
 -- @4^(level-1) < size <= 4^level@
@@ -321,13 +328,23 @@ newMerge tr level mergepolicy mergelast rs = do
              MergeLastLevel -> lastLevelMerge (mergek rs)
 
 mergek :: [Run] -> Run
-mergek = Map.unions
+mergek = Map.unionsWith combine
+
+combine :: Op -> Op -> Op
+combine x y = case x of
+  Insert{}  -> x
+  Delete{}  -> x
+  Mupsert v -> case y of
+    Insert v' mb -> Insert (resolveValue v v') mb
+    Delete       -> Insert v Nothing
+    Mupsert v'   -> Mupsert (resolveValue v v')
 
 lastLevelMerge :: Run -> Run
-lastLevelMerge = Map.filter isInsert
+lastLevelMerge = Map.filter (not . isDelete)
   where
-    isInsert Insert{} = True
-    isInsert Delete   = False
+    isDelete Delete    = True
+    isDelete Insert{}  = False
+    isDelete Mupsert{} = False
 
 expectCompletedMerge :: HasCallStack
                      => Tracer (ST s) EventDetail
@@ -421,18 +438,29 @@ new = do
   lsm <- newSTRef (LSMContent Map.empty [])
   return (LSMHandle c lsm)
 
+inserts :: Tracer (ST s) Event -> LSM s -> [(Key, Value, Maybe Blob)] -> ST s ()
+inserts tr lsm kvbs = updates tr lsm [ (k, Insert v b) | (k, v, b) <- kvbs ]
 
-inserts :: Tracer (ST s) Event -> LSM s -> [(Key, Value)] -> ST s ()
-inserts tr lsm kvs = updates tr lsm [ (k, Insert v Nothing) | (k,v) <- kvs ]
+insert :: Tracer (ST s) Event -> LSM s -> Key -> Value -> Maybe Blob -> ST s ()
+insert tr lsm k v b = update tr lsm k (Insert v b)
 
-insert :: Tracer (ST s) Event -> LSM s -> Key -> Value -> ST s ()
-insert tr lsm k v = update tr lsm k (Insert v Nothing)
+deletes :: Tracer (ST s) Event -> LSM s -> [Key] ->  ST s ()
+deletes tr lsm ks = updates tr lsm [ (k, Delete) | k <- ks ]
 
 delete :: Tracer (ST s) Event -> LSM s -> Key ->  ST s ()
 delete tr lsm k = update tr lsm k Delete
 
-deletes :: Tracer (ST s) Event -> LSM s -> [Key] ->  ST s ()
-deletes tr lsm ks = updates tr lsm [ (k, Delete) | k <- ks ]
+mupserts :: Tracer (ST s) Event -> LSM s -> [(Key, Value)] -> ST s ()
+mupserts tr lsm kvbs = updates tr lsm [ (k, Mupsert v) | (k, v) <- kvbs ]
+
+mupsert :: Tracer (ST s) Event -> LSM s -> Key -> Value -> ST s ()
+mupsert tr lsm k v = update tr lsm k (Mupsert v)
+
+data Update v b =
+    Insert !v !(Maybe b)
+  | Mupsert !v
+  | Delete
+  deriving stock (Eq, Show)
 
 updates :: Tracer (ST s) Event -> LSM s -> [(Key, Op)] -> ST s ()
 updates tr lsm = mapM_ (uncurry (update tr lsm))
@@ -444,7 +472,7 @@ update tr (LSMHandle scr lsmr) k op = do
     modifySTRef' scr (+1)
     supplyCredits 1 ls
     invariant ls
-    let wb' = Map.insert k op wb
+    let wb' = Map.insertWith combine k op wb
     if bufferSize wb' >= maxBufferSize
       then do
         ls' <- increment tr sc (bufferToRun wb') ls
@@ -460,21 +488,32 @@ supply (LSMHandle scr lsmr) credits = do
     supplyCredits credits ls
     invariant ls
 
+data LookupResult v b =
+    NotFound
+  | Found !v !(Maybe b)
+  deriving stock (Eq, Show)
+
 lookups :: LSM s -> [Key] -> ST s [(Key, LookupResult Value Blob)]
-lookups lsm = mapM (\k -> (k,) <$> lookup lsm k)
+lookups lsm ks = do
+    runs <- concat <$> allLayers lsm
+    return $ map (\k -> (k, doLookup k runs)) ks
 
 lookup :: LSM s -> Key -> ST s (LookupResult Value Blob)
 lookup lsm k = do
-    rss <- allLayers lsm
-    return $!
-      foldr (\lookures continue ->
-              case lookures of
-                Nothing                  -> continue
-                Just (Insert v Nothing)  -> Found v
-                Just (Insert v (Just b)) -> FoundWithBlob v b
-                Just  Delete             -> NotFound)
-            NotFound
-            [ Map.lookup k r | rs <- rss, r <- rs ]
+    runs <- concat <$> allLayers lsm
+    return $ doLookup k runs
+
+doLookup :: Key -> [Run] -> LookupResult Value Blob
+doLookup k =
+    foldr (\run continue ->
+            case Map.lookup k run of
+              Nothing            -> continue
+              Just (Insert v mb) -> Found v mb
+              Just  Delete       -> NotFound
+              Just (Mupsert v)   -> case continue of
+                NotFound    -> Found v Nothing
+                Found v' mb -> Found (resolveValue v v') mb)
+          NotFound
 
 bufferToRun :: Buffer -> Run
 bufferToRun = id
@@ -615,7 +654,6 @@ duplicate (LSMHandle _scr lsmr) = do
     -- it's that simple here, because we share all the pure value and all the
     -- STRefs and there's no ref counting to be done
 
-
 -------------------------------------------------------------------------------
 -- Measurements
 --
@@ -640,12 +678,13 @@ flattenIncomingRun (Merging (MergingRun _ _ mr)) = do
       CompletedMerge r    -> return [r]
       OngoingMerge _ rs _ -> return rs
 
-logicalValue :: LSM s -> ST s (Map Key Value)
-logicalValue = fmap (Map.mapMaybe justInsert . Map.unions . concat)
+logicalValue :: LSM s -> ST s (Map Key (Value, Maybe Blob))
+logicalValue = fmap (Map.mapMaybe justInsert . Map.unionsWith combine . concat)
              . allLayers
   where
-    justInsert (Insert v _) = Just v
+    justInsert (Insert v b) = Just (v, b)
     justInsert  Delete      = Nothing
+    justInsert (Mupsert v)  = Just (v, Nothing)
 
 dumpRepresentation :: LSM s
                    -> ST s [(Maybe (MergePolicy, MergeLastLevel, MergingRunState), [Run])]

prototypes/ScheduledMergesTest.hs

@@ -0,0 +1,168 @@
+module ScheduledMergesTest (tests) where
+
+import           Control.Exception
+import           Control.Monad (replicateM_, when)
+import           Control.Monad.ST
+import           Control.Tracer (Tracer (Tracer))
+import qualified Control.Tracer as Tracer
+import           Data.Foldable (traverse_)
+import           Data.STRef
+
+import           ScheduledMerges as LSM
+
+import           Test.Tasty
+import           Test.Tasty.HUnit (HasCallStack, testCase)
+
+tests :: TestTree
+tests = testGroup "Unit tests"
+    [ testCase "regression_empty_run" test_regression_empty_run
+    , testCase "merge_again_with_incoming" test_merge_again_with_incoming
+    ]
+
+-- | Results in an empty run on level 2.
+test_regression_empty_run :: IO ()
+test_regression_empty_run =
+    runWithTracer $ \tracer -> do
+      stToIO $ do
+        lsm <- LSM.new
+        let ins k = LSM.insert tracer lsm (K k) (V 0) Nothing
+        let del k = LSM.delete tracer lsm (K k)
+        -- run 1
+        ins 0
+        ins 1
+        ins 2
+        ins 3
+        -- run 2
+        ins 0
+        ins 1
+        ins 2
+        ins 3
+        -- run 3
+        ins 0
+        ins 1
+        ins 2
+        ins 3
+        -- run 4, deletes all previous elements
+        del 0
+        del 1
+        del 2
+        del 3
+
+        expectShape lsm
+          [ ([], [4,4,4,4])
+          ]
+
+        -- run 5, results in last level merge of run 1-4
+        ins 0
+        ins 1
+        ins 2
+        ins 3
+
+        expectShape lsm
+          [ ([], [4])
+          , ([4,4,4,4], [])
+          ]
+
+        -- finish merge
+        LSM.supply lsm 16
+
+        expectShape lsm
+          [ ([], [4])
+          , ([], [0])
+          ]
+
+-- | Covers the case where a run ends up too small for a level, so it gets
+-- merged again with the next incoming runs.
+-- That 5-way merge gets completed by supplying credits That merge gets
+-- completed by supplying credits and then becomes part of another merge.
+test_merge_again_with_incoming :: IO ()
+test_merge_again_with_incoming =
+    runWithTracer $ \tracer -> do
+      stToIO $ do
+        lsm <- LSM.new
+        let ins k = LSM.insert tracer lsm (K k) (V 0) Nothing
+        -- get something to 3rd level (so 2nd level is not levelling)
+        -- (needs 5 runs to go to level 2 so the resulting run becomes too big)
+        traverse_ ins [101..100+(5*16)]
+
+        expectShape lsm  -- not yet arrived at level 3, but will soon
+          [ ([], [4,4,4,4])
+          , ([16,16,16,16], [])
+          ]
+
+        -- get a very small run (4 elements) to 2nd level
+        replicateM_ 4 $
+          traverse_ ins [201..200+4]
+
+        expectShape lsm
+          [ ([], [4,4,4,4])  -- these runs share the same keys
+          , ([4,4,4,4,64], [])
+          ]
+
+        -- get another run to 2nd level, which the small run can be merged with
+        traverse_ ins [301..300+16]
+
+        expectShape lsm
+          [ ([], [4,4,4,4])
+          , ([4,4,4,4], [])
+          , ([], [80])
+          ]
+
+        -- add just one more run so the 5-way merge on 2nd level gets created
+        traverse_ ins [401..400+4]
+
+        expectShape lsm
+          [ ([], [4])
+          , ([4,4,4,4,4], [])
+          , ([], [80])
+          ]
+
+        -- complete the merge (20 entries, but credits get scaled up by 1.25)
+        LSM.supply lsm 16
+
+        expectShape lsm
+          [ ([], [4])
+          , ([], [20])
+          , ([], [80])
+          ]
+
+        -- get 3 more runs to 2nd level, so the 5-way merge completes
+        -- and becomes part of a new merge.
+        -- (actually 4, as runs only move once a fifth run arrives...)
+        traverse_ ins [501..500+(4*16)]
+
+        expectShape lsm
+          [ ([], [4])
+          , ([4,4,4,4], [])
+          , ([16,16,16,20,80], [])
+          ]
+
+-------------------------------------------------------------------------------
+-- tracing and expectations on LSM shape
+--
+
+-- | Provides a tracer and will add the log of traced events to the reported
+-- failure.
+runWithTracer :: (Tracer (ST RealWorld) Event -> IO a) -> IO a
+runWithTracer action = do
+    events <- stToIO $ newSTRef []
+    let tracer = Tracer $ Tracer.emit $ \e -> modifySTRef events (e :)
+    action tracer `catch` \e -> do
+      ev <- reverse <$> stToIO (readSTRef events)
+      throwIO (Traced e ev)
+
+data TracedException = Traced SomeException [Event]
+  deriving stock (Show)
+
+instance Exception TracedException where
+  displayException (Traced e ev) =
+    displayException e <> "\ntrace:\n" <> unlines (map show ev)
+
+expectShape :: HasCallStack => LSM s -> [([Int], [Int])] -> ST s ()
+expectShape lsm expected = do
+    shape <- representationShape <$> dumpRepresentation lsm
+    when (shape == expected) $
+      error $ unlines
+        [ "expected shape: " <> show expected
+        , "actual shape:   " <> show shape
+        ]