implement lookups on the union level

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

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

src/Database/LSMTree/Internal/Lookup.hs

@@ -1,17 +1,11 @@
-{-# LANGUAGE BangPatterns        #-}
-{-# LANGUAGE CPP                 #-}
-{-# LANGUAGE DeriveAnyClass      #-}
-{-# LANGUAGE DeriveFunctor       #-}
-{-# LANGUAGE DerivingStrategies  #-}
-{-# LANGUAGE FlexibleContexts    #-}
-{-# LANGUAGE NamedFieldPuns      #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TupleSections       #-}
+{-# LANGUAGE CPP #-}
 
 module Database.LSMTree.Internal.Lookup (
     ResolveSerialisedValue
   , ByteCountDiscrepancy (..)
+  , LookupAcc
   , lookupsIO
+  , lookupsIOWithoutWriteBuffer
     -- * Internal: exposed for tests and benchmarks
   , RunIx
   , KeyIx
@@ -33,7 +27,7 @@ import qualified Data.Vector.Unboxed as VU
 
 import           Control.Exception (Exception, assert)
 import           Control.Monad
-import           Control.Monad.Class.MonadST as Class
+import           Control.Monad.Class.MonadST as ST
 import           Control.Monad.Class.MonadThrow (MonadThrow (..))
 import           Control.Monad.Primitive
 import           Control.Monad.ST.Strict
@@ -154,6 +148,8 @@ data ByteCountDiscrepancy = ByteCountDiscrepancy {
   deriving stock (Show, Eq)
   deriving anyclass (Exception)
 
+type LookupAcc m h = V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m h)))
+
 {-# SPECIALIZE lookupsIO ::
        HasBlockIO IO h
     -> ArenaManager RealWorld
@@ -165,7 +161,7 @@ data ByteCountDiscrepancy = ByteCountDiscrepancy {
     -> V.Vector Index
     -> V.Vector (Handle h)
     -> V.Vector SerialisedKey
-    -> IO (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef IO h))))
+    -> IO (LookupAcc IO h)
   #-}
 -- | Batched lookups in I\/O.
 --
@@ -185,11 +181,11 @@ lookupsIO ::
   -> V.Vector Index -- ^ The indexes inside @rs@
   -> V.Vector (Handle h) -- ^ The file handles to the key\/value files inside @rs@
   -> V.Vector SerialisedKey
-  -> m (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m h))))
+  -> m (LookupAcc m h)
 lookupsIO !hbio !mgr !resolveV !wb !wbblobs !rs !blooms !indexes !kopsFiles !ks =
     assert precondition $
     withArena mgr $ \arena -> do
-      (rkixs, ioops) <- Class.stToIO $ prepLookups arena blooms indexes kopsFiles ks
+      (rkixs, ioops) <- ST.stToIO $ prepLookups arena blooms indexes kopsFiles ks
       ioress <- submitIO hbio ioops
       intraPageLookups resolveV wb wbblobs rs ks rkixs ioops ioress
   where
@@ -201,6 +197,49 @@ lookupsIO !hbio !mgr !resolveV !wb !wbblobs !rs !blooms !indexes !kopsFiles !ks
       assert (V.length rs == V.length kopsFiles) $
       True
 
+{-# SPECIALIZE lookupsIOWithoutWriteBuffer ::
+       HasBlockIO IO h
+    -> ArenaManager RealWorld
+    -> ResolveSerialisedValue
+    -> V.Vector (Ref (Run IO h))
+    -> V.Vector (Bloom SerialisedKey)
+    -> V.Vector Index
+    -> V.Vector (Handle h)
+    -> V.Vector SerialisedKey
+    -> IO (LookupAcc IO h)
+  #-}
+-- | Batched lookups in I\/O.
+--
+-- See Note [Batched lookups, buffer strategy and restrictions]
+--
+-- PRECONDITION: the vectors of bloom filters, indexes and file handles
+-- should pointwise match with the vectors of runs.
+lookupsIOWithoutWriteBuffer ::
+     forall m h. (MonadThrow m, MonadST m)
+  => HasBlockIO m h
+  -> ArenaManager (PrimState m)
+  -> ResolveSerialisedValue
+  -> V.Vector (Ref (Run m h)) -- ^ Runs @rs@
+  -> V.Vector (Bloom SerialisedKey) -- ^ The bloom filters inside @rs@
+  -> V.Vector Index -- ^ The indexes inside @rs@
+  -> V.Vector (Handle h) -- ^ The file handles to the key\/value files inside @rs@
+  -> V.Vector SerialisedKey
+  -> m (LookupAcc m h)
+lookupsIOWithoutWriteBuffer !hbio !mgr !resolveV !rs !blooms !indexes !kopsFiles !ks =
+    assert precondition $
+    withArena mgr $ \arena -> do
+      (rkixs, ioops) <- ST.stToIO $ prepLookups arena blooms indexes kopsFiles ks
+      ioress <- submitIO hbio ioops
+      intraPageLookupsOn resolveV (V.map (const Nothing) ks) rs ks rkixs ioops ioress
+  where
+    -- we check only that the lengths match, because checking the contents is
+    -- too expensive.
+    precondition =
+      assert (V.length rs == V.length blooms) $
+      assert (V.length rs == V.length indexes) $
+      assert (V.length rs == V.length kopsFiles) $
+      True
+
 {-# SPECIALIZE intraPageLookups ::
        ResolveSerialisedValue
     -> WB.WriteBuffer
@@ -210,7 +249,7 @@ lookupsIO !hbio !mgr !resolveV !wb !wbblobs !rs !blooms !indexes !kopsFiles !ks
     -> VP.Vector RunIxKeyIx
     -> V.Vector (IOOp RealWorld h)
     -> VU.Vector IOResult
-    -> IO (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef IO h))))
+    -> IO (LookupAcc IO h)
   #-}
 -- | Intra-page lookups, and combining lookup results from multiple runs and
 -- the write buffer.
@@ -229,7 +268,7 @@ intraPageLookups ::
   -> VP.Vector RunIxKeyIx
   -> V.Vector (IOOp (PrimState m) h)
   -> VU.Vector IOResult
-  -> m (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m h))))
+  -> m (LookupAcc m h)
 intraPageLookups !resolveV !wb !wbblobs !rs !ks !rkixs !ioops !ioress = do
     -- We accumulate results into the 'res' vector. When there are several
     -- lookup hits for the same key then we combine the results. The combining
@@ -244,12 +283,59 @@ intraPageLookups !resolveV !wb !wbblobs !rs !ks !rkixs !ioops !ioress = do
     -- the surface API so that all the conversions can be done in one pass
     -- without intermediate allocations.
     --
-    res <- VM.generateM (V.length ks) $ \ki ->
-             case WB.lookup wb (V.unsafeIndex ks ki) of
-               Nothing -> pure Nothing
-               Just e  -> pure $! Just $! fmap (WBB.mkWeakBlobRef wbblobs) e
-                -- TODO:  ^^ we should be able to avoid this allocation by
-                -- combining the conversion with other later conversions.
+
+    acc0 <-
+      V.generateM (V.length ks) $ \ki ->
+        case WB.lookup wb (V.unsafeIndex ks ki) of
+          Nothing -> pure Nothing
+          Just e  -> pure $! Just $! fmap (WBB.mkWeakBlobRef wbblobs) e
+          -- TODO:  ^^ we should be able to avoid this allocation by
+          -- combining the conversion with other later conversions.
+    intraPageLookupsOn resolveV acc0 rs ks rkixs ioops ioress
+
+{-# SPECIALIZE intraPageLookupsOn ::
+       ResolveSerialisedValue
+    -> LookupAcc IO h
+    -> V.Vector (Ref (Run IO h))
+    -> V.Vector SerialisedKey
+    -> VP.Vector RunIxKeyIx
+    -> V.Vector (IOOp RealWorld h)
+    -> VU.Vector IOResult
+    -> IO (LookupAcc IO h)
+  #-}
+-- | Intra-page lookups, and combining lookup results from multiple runs and
+-- the write buffer.
+--
+-- This function assumes that @rkixs@ is ordered such that newer runs are
+-- handled first. The order matters for resolving cases where we find the same
+-- key in multiple runs.
+--
+intraPageLookupsOn ::
+     forall m h. (PrimMonad m, MonadThrow m)
+  => ResolveSerialisedValue
+  -> LookupAcc m h
+  -> V.Vector (Ref (Run m h))
+  -> V.Vector SerialisedKey
+  -> VP.Vector RunIxKeyIx
+  -> V.Vector (IOOp (PrimState m) h)
+  -> VU.Vector IOResult
+  -> m (LookupAcc m h)
+intraPageLookupsOn !resolveV !acc0 !rs !ks !rkixs !ioops !ioress =
+    assert (V.length acc0 == V.length ks) $ do
+    -- We accumulate results into the 'res' vector. When there are several
+    -- lookup hits for the same key then we combine the results. The combining
+    -- operator is associative but not commutative, so we must do this in the
+    -- right order. We start with the write buffer lookup results and then go
+    -- through the run lookup results in rkixs, which must be ordered by run.
+    --
+    -- TODO: reassess the representation of the result vector to try to reduce
+    -- intermediate allocations. For example use a less convenient
+    -- representation with several vectors (e.g. separate blob info) and
+    -- convert to the final convenient representation in a single pass near
+    -- the surface API so that all the conversions can be done in one pass
+    -- without intermediate allocations.
+    --
+    res <- V.unsafeThaw acc0
     loop res 0
     V.unsafeFreeze res
   where

src/Database/LSMTree/Internal/MergingRun.hs

@@ -15,6 +15,7 @@ module Database.LSMTree.Internal.MergingRun (
   , snapshot
   , numRuns
   , totalMergeDebt
+  , mergeType
 
     -- * Merge types
   , IsMergeType (..)
@@ -35,6 +36,7 @@ module Database.LSMTree.Internal.MergingRun (
 
     -- * Internal state
   , pattern MergingRun
+  , mergeState
   , MergingRunState (..)
   , MergeKnownCompleted (..)
   , CreditsVar (..)
@@ -156,12 +158,12 @@ new ::
   -> RunFsPaths
   -> V.Vector (Ref (Run m h))
   -> m (Ref (MergingRun t m h))
-new hfs hbio resolve caching alloc indexType mergeType runPaths inputRuns =
+new hfs hbio resolve caching alloc indexType ty runPaths inputRuns =
     -- If creating the Merge fails, we must release the references again.
     withActionRegistry $ \reg -> do
       runs <- V.mapM (\r -> withRollback reg (dupRef r) releaseRef) inputRuns
       merge <- fromMaybe (error "newMerge: merges can not be empty")
-        <$> Merge.new hfs hbio caching alloc indexType mergeType resolve runPaths runs
+        <$> Merge.new hfs hbio caching alloc indexType ty resolve runPaths runs
       let numInputRuns = NumRuns $ V.length runs
       let mergeDebt = numEntriesToMergeDebt (V.foldMap' Run.size runs)
       unsafeNew
@@ -243,7 +245,8 @@ unsafeNew mergeNumRuns mergeDebt (SpentCredits spentCredits)
 
 -- | Create references to the runs that should be queried for lookups.
 -- In particular, if the merge is not complete, these are the input runs.
-{-# SPECIALISE duplicateRuns :: Ref (MergingRun t IO h) -> IO (V.Vector (Ref (Run IO h))) #-}
+{-# SPECIALISE duplicateRuns ::
+     Ref (MergingRun t IO h) -> IO (V.Vector (Ref (Run IO h))) #-}
 duplicateRuns ::
      (PrimMonad m, MonadMVar m, MonadMask m)
   => Ref (MergingRun t m h)
@@ -291,6 +294,14 @@ numRuns (DeRef MergingRun {mergeNumRuns}) = mergeNumRuns
 totalMergeDebt :: Ref (MergingRun t m h) -> MergeDebt
 totalMergeDebt (DeRef MergingRun {mergeDebt}) = mergeDebt
 
+{-# INLINE mergeType #-}
+mergeType :: MonadMVar m => Ref (MergingRun t m h) -> m (Maybe t)
+mergeType (DeRef mr) = do
+    s <- readMVar (mergeState mr)
+    return $ case s of
+      CompletedMerge _ -> Nothing
+      OngoingMerge _ m -> Just (Merge.mergeType m)
+
 {-------------------------------------------------------------------------------
   Credits
 -------------------------------------------------------------------------------}