Merge pull request #609 from IntersectMBO/mheinzel/merging-tree-lookups

Lookups in merging tree

Author: mheinzel

lsm-tree.cabal

@@ -141,6 +141,7 @@ library
     Database.LSMTree.Internal.MergeSchedule
     Database.LSMTree.Internal.MergingRun
     Database.LSMTree.Internal.MergingTree
+    Database.LSMTree.Internal.MergingTree.Lookup
     Database.LSMTree.Internal.Page
     Database.LSMTree.Internal.PageAcc
     Database.LSMTree.Internal.PageAcc1

prototypes/ScheduledMerges.hs

@@ -935,16 +935,17 @@ checkedUnionDebt tree debtRef = do
 
 type LookupAcc = Maybe Op
 
-mergeAcc :: [LookupAcc] -> LookupAcc
-mergeAcc = foldl (updateAcc combine) Nothing . catMaybes
-
-unionAcc :: [LookupAcc] -> LookupAcc
-unionAcc = foldl (updateAcc combineUnion) Nothing . catMaybes
-
 updateAcc :: (Op -> Op -> Op) -> LookupAcc -> Op -> LookupAcc
 updateAcc _ Nothing     old = Just old
 updateAcc f (Just new_) old = Just (f new_ old)  -- acc has more recent Op
 
+mergeAcc :: TreeMergeType -> [LookupAcc] -> LookupAcc
+mergeAcc mt = foldl (updateAcc com) Nothing . catMaybes
+  where
+    com = case mt of
+      MergeLevel -> combine
+      MergeUnion -> combineUnion
+
 -- | We handle lookups by accumulating results by going through the runs from
 -- most recent to least recent, starting with the write buffer.
 --
@@ -957,8 +958,12 @@ doLookup wb runs ul k = do
       NoUnion ->
         return (convertAcc acc0)
       Union tree _ -> do
-        accTree <- lookupsTree k tree
-        return (convertAcc (mergeAcc [acc0, accTree]))
+        treeBatches <- buildLookupTree tree
+        let treeResults = lookupBatch Nothing k <$> treeBatches
+        return $ convertAcc $ foldLookupTree $
+          if null wb && null runs
+          then treeResults
+          else LookupNode MergeLevel [LookupBatch acc0, treeResults ]
   where
     convertAcc :: LookupAcc -> LookupResult Value Blob
     convertAcc = \case
@@ -976,6 +981,10 @@ lookupBatch acc k rs =
     let ops = [op | r <- rs, Just op <- [Map.lookup k r]]
     in foldl (updateAcc combine) acc ops
 
+data LookupTree a = LookupBatch a
+                  | LookupNode TreeMergeType [LookupTree a]
+  deriving stock Functor
+
 -- | Do lookups on runs at the leaves and recursively combine the resulting
 -- 'LookupAcc's, either using 'mergeAcc' or 'unionAcc' depending on the merge
 -- type.
@@ -989,32 +998,38 @@ lookupBatch acc k rs =
 -- have a union level) and then do lookups, two batches of lookups have to be
 -- performed (plus a batch for the table's regular levels if it has been updated
 -- after the union).
-lookupsTree :: Key -> MergingTree s -> ST s LookupAcc
-lookupsTree k = go
+--
+-- TODO: we can still improve the batching, for example combining the child of
+-- PendingLevelMerge with the pre-existing runs when it is already completed.
+buildLookupTree :: MergingTree s -> ST s (LookupTree [Run])
+buildLookupTree = go
   where
-    go :: MergingTree s -> ST s LookupAcc
+    go :: MergingTree s -> ST s (LookupTree [Run])
     go (MergingTree treeState) = readSTRef treeState >>= \case
         CompletedTreeMerge r ->
-          return $ lookupBatch' [r]
+          return $ LookupBatch [r]
         OngoingTreeMerge (MergingRun mt _ mergeState) ->
           readSTRef mergeState >>= \case
             CompletedMerge r ->
-              return $ lookupBatch' [r]
+              return $ LookupBatch [r]
             OngoingMerge _ rs _ -> case mt of
-              MergeLevel -> return $ lookupBatch' rs  -- combine into batch
-              MergeUnion -> return $ unionAcc (map (\r -> lookupBatch' [r]) rs)
-        PendingTreeMerge (PendingUnionMerge trees) -> do
-          unionAcc <$> traverse go trees
+              MergeLevel -> return $ LookupBatch rs  -- combine into batch
+              MergeUnion -> return $ LookupNode MergeUnion $ map (\r -> LookupBatch [r]) rs
         PendingTreeMerge (PendingLevelMerge prs tree) -> do
-          runs <- concat <$> traverse flattenPreExistingRun prs -- combine into batch
-          let acc0 = lookupBatch' runs
+          preExisting <- LookupBatch . concat <$>
+            traverse flattenPreExistingRun prs -- combine into batch
           case tree of
-            Nothing -> return acc0  -- only runs and merging level runs, done
+            Nothing -> return preExisting
             Just t  -> do
-              accTree <- go t
-              return (mergeAcc [acc0, accTree])
+              lTree <- go t
+              return (LookupNode MergeLevel [preExisting, lTree])
+        PendingTreeMerge (PendingUnionMerge trees) -> do
+          LookupNode MergeUnion <$> traverse go trees
 
-    lookupBatch' = lookupBatch Nothing k
+foldLookupTree :: LookupTree LookupAcc -> LookupAcc
+foldLookupTree = \case
+    LookupBatch acc        -> acc
+    LookupNode mt children -> mergeAcc mt (map foldLookupTree children)
 
 -------------------------------------------------------------------------------
 -- Nominal credits

src/Database/LSMTree/Common.hs

@@ -49,7 +49,8 @@ module Database.LSMTree.Common (
   ) where
 
 import           Control.Concurrent.Class.MonadMVar.Strict
-import           Control.Concurrent.Class.MonadSTM (MonadSTM, STM)
+import           Control.Concurrent.Class.MonadSTM (STM)
+import           Control.Monad.Class.MonadAsync
 import           Control.Monad.Class.MonadST
 import           Control.Monad.Class.MonadThrow
 import           Control.Monad.Primitive (PrimMonad)
@@ -74,8 +75,8 @@ import           System.FS.IO (HandleIO)
 -------------------------------------------------------------------------------}
 
 -- | Utility class for grouping @io-classes@ constraints.
-class ( MonadMVar m, MonadSTM m, MonadThrow (STM m), MonadThrow m, MonadCatch m
-      , MonadMask m, PrimMonad m, MonadST m
+class ( MonadAsync m, MonadMVar m, MonadThrow (STM m), MonadThrow m
+      , MonadCatch m , MonadMask m, PrimMonad m, MonadST m
       ) => IOLike m
 
 instance IOLike IO

src/Database/LSMTree/Internal.hs

@@ -82,6 +82,7 @@ import           Control.Concurrent.Class.MonadSTM.RWVar (RWVar)
 import qualified Control.Concurrent.Class.MonadSTM.RWVar as RW
 import           Control.DeepSeq
 import           Control.Monad (forM, unless, void)
+import           Control.Monad.Class.MonadAsync as Async
 import           Control.Monad.Class.MonadST (MonadST (..))
 import           Control.Monad.Class.MonadThrow
 import           Control.Monad.Primitive
@@ -105,9 +106,12 @@ import           Database.LSMTree.Internal.Config
 import qualified Database.LSMTree.Internal.Cursor as Cursor
 import           Database.LSMTree.Internal.Entry (Entry)
 import           Database.LSMTree.Internal.Lookup (ByteCountDiscrepancy,
-                     ResolveSerialisedValue, lookupsIO)
+                     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.Lookup as MT
 import           Database.LSMTree.Internal.Paths (SessionRoot (..),
                      SnapshotMetaDataChecksumFile (..),
                      SnapshotMetaDataFile (..), SnapshotName)
@@ -123,6 +127,7 @@ import           Database.LSMTree.Internal.Serialise (SerialisedBlob (..),
 import           Database.LSMTree.Internal.Snapshot
 import           Database.LSMTree.Internal.Snapshot.Codec
 import           Database.LSMTree.Internal.UniqCounter
+import qualified Database.LSMTree.Internal.Vector as V
 import qualified Database.LSMTree.Internal.WriteBuffer as WB
 import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB
 import qualified System.FS.API as FS
@@ -799,16 +804,48 @@ close t = do
   -> IO (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef IO h)))) #-}
 -- | See 'Database.LSMTree.Normal.lookups'.
 lookups ::
-     (MonadST m, MonadSTM m, MonadThrow m)
+     (MonadAsync m, MonadMask m, MonadMVar m, MonadST m)
   => ResolveSerialisedValue
   -> V.Vector SerialisedKey
   -> Table m h
   -> m (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m h))))
 lookups resolve ks t = do
     traceWith (tableTracer t) $ TraceLookups (V.length ks)
     withOpenTable t $ \tEnv ->
-      RW.withReadAccess (tableContent tEnv) $ \tableContent ->
-        let !cache = tableCache tableContent in
+      RW.withReadAccess (tableContent tEnv) $ \tableContent -> do
+        case tableUnionLevel tableContent of
+          NoUnion -> regularLevelLookups tEnv tableContent
+          Union tree -> do
+            isStructurallyEmpty tree >>= \case
+              True  -> regularLevelLookups tEnv tableContent
+              False ->
+                -- TODO: the blob refs returned from the tree can be invalidated
+                -- by supplyUnionCredits or other operations on any table that
+                -- shares merging runs or trees. We need to keep open the runs!
+                -- This could be achieved by storing the LookupTree and only
+                -- calling MT.releaseLookupTree later, when we are okay with
+                -- invalidating the blob refs (similar to the LevelsCache).
+                -- Lookups then use the cached tree, but when should we rebuild
+                -- the tree? On each call to supplyUnionCredits?
+                withActionRegistry $ \reg -> do
+                  regularResult <-
+                    -- asynchronously, so tree lookup batches can already be
+                    -- submitted without waiting for the result.
+                    Async.async $ regularLevelLookups tEnv tableContent
+                  treeBatches <- MT.buildLookupTree reg tree
+                  treeResults <- forM treeBatches $ \runs ->
+                    Async.async $ treeBatchLookups tEnv runs
+                  -- TODO: if regular levels are empty, don't add them to tree
+                  res <- MT.foldLookupTree resolve $
+                    MT.mkLookupNode MR.MergeLevel $ V.fromList
+                      [ MT.LookupBatch regularResult
+                      , treeResults
+                      ]
+                  MT.releaseLookupTree reg treeBatches
+                  return res
+  where
+    regularLevelLookups tEnv tableContent = do
+        let !cache = tableCache tableContent
         lookupsIO
           (tableHasBlockIO tEnv)
           (tableArenaManager t)
@@ -821,6 +858,17 @@ lookups resolve ks t = do
           (cachedKOpsFiles cache)
           ks
 
+    treeBatchLookups tEnv runs =
+        lookupsIOWithoutWriteBuffer
+          (tableHasBlockIO tEnv)
+          (tableArenaManager t)
+          resolve
+          runs
+          (V.mapStrict (\(DeRef r) -> Run.runFilter   r) runs)
+          (V.mapStrict (\(DeRef r) -> Run.runIndex    r) runs)
+          (V.mapStrict (\(DeRef r) -> Run.runKOpsFile r) runs)
+          ks
+
 {-# SPECIALISE rangeLookup ::
      ResolveSerialisedValue
   -> Range SerialisedKey