Merge pull request #647 from IntersectMBO/mheinzel/union-improvements

Small improvements related to table unions

Author: mheinzel

bench/macro/lsm-tree-bench-lookups.hs

@@ -486,7 +486,8 @@ benchLookupsIO !hbio !arenaManager !resolve !wb !wbblobs !rs !bs !ics !hs =
       | n <= 0    = pure ()
       | otherwise = do
           let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize
-          !_ <- lookupsIO hbio arenaManager resolve wb wbblobs rs bs ics hs ks
+          !_ <- lookupsIOWithWriteBuffer
+                  hbio arenaManager resolve wb wbblobs rs bs ics hs ks
           go keyRng' (n-benchmarkGenBatchSize)
 
 {-------------------------------------------------------------------------------

bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs

@@ -23,7 +23,8 @@ import           Database.LSMTree.Extras.UTxO
 import           Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))
 import           Database.LSMTree.Internal.Index as Index
 import           Database.LSMTree.Internal.Lookup (bloomQueries, indexSearches,
-                     intraPageLookups, lookupsIO, prepLookups)
+                     intraPageLookupsWithWriteBuffer, lookupsIOWithWriteBuffer,
+                     prepLookups)
 import           Database.LSMTree.Internal.Page (getNumPages)
 import           Database.LSMTree.Internal.Paths (RunFsPaths (..))
 import           Database.LSMTree.Internal.Run (Run)
@@ -132,17 +133,18 @@ benchLookups conf@Config{name} =
                 )
                 (\(_, _, _, arena) -> closeArena arenaManager arena)
                 (\ ~(rkixs, ioops, ioress, _) -> do
-                  !_ <- intraPageLookups resolveV WB.empty wbblobs
-                                         rs ks rkixs ioops ioress
+                  !_ <- intraPageLookupsWithWriteBuffer
+                          resolveV WB.empty wbblobs rs ks rkixs ioops ioress
                   pure ())
             -- The whole shebang: lookup preparation, doing the IO, and then
             -- performing intra-page-lookups. Again, we evaluate the result to
-            -- WHNF because it is the same result that intraPageLookups produces
-            -- (see above).
+            -- WHNF because it is the same result that
+            -- intraPageLookupsWithWriteBuffer produces (see above).
           , bench "Lookups in IO" $
-              whnfAppIO (\ks' -> lookupsIO hasBlockIO arenaManager resolveV
-                                           WB.empty wbblobs
-                                           rs blooms indexes kopsFiles ks') ks
+              whnfAppIO (\ks' -> lookupsIOWithWriteBuffer
+                                   hasBlockIO arenaManager resolveV
+                                   WB.empty wbblobs
+                                   rs blooms indexes kopsFiles ks') ks
           ]
           -- TODO: consider adding benchmarks that also use the write buffer
           -- (then we can't just use 'WB.empty', but must take it from the env)

src/Database/LSMTree/Internal.hs

@@ -127,7 +127,7 @@ import           Database.LSMTree.Internal.Entry (Entry, NumEntries (..))
 import           Database.LSMTree.Internal.IncomingRun (IncomingRun (..))
 import           Database.LSMTree.Internal.Lookup (ResolveSerialisedValue,
                      TableCorruptedError (..), lookupsIO,
-                     lookupsIOWithoutWriteBuffer)
+                     lookupsIOWithWriteBuffer)
 import           Database.LSMTree.Internal.MergeSchedule
 import           Database.LSMTree.Internal.MergingRun (TableTooLargeError (..))
 import qualified Database.LSMTree.Internal.MergingRun as MR
@@ -838,7 +838,7 @@ lookups resolve ks t = do
   where
     regularLevelLookups tEnv tableContent = do
         let !cache = tableCache tableContent
-        lookupsIO
+        lookupsIOWithWriteBuffer
           (tableHasBlockIO tEnv)
           (tableArenaManager t)
           resolve
@@ -851,7 +851,7 @@ lookups resolve ks t = do
           ks
 
     treeBatchLookups tEnv runs =
-        lookupsIOWithoutWriteBuffer
+        lookupsIO
           (tableHasBlockIO tEnv)
           (tableArenaManager t)
           resolve

src/Database/LSMTree/Internal/Lookup.hs

@@ -3,8 +3,8 @@
 module Database.LSMTree.Internal.Lookup (
     ResolveSerialisedValue
   , LookupAcc
+  , lookupsIOWithWriteBuffer
   , lookupsIO
-  , lookupsIOWithoutWriteBuffer
     -- * Errors
   , TableCorruptedError (..)
     -- * Internal: exposed for tests and benchmarks
@@ -14,7 +14,8 @@ module Database.LSMTree.Internal.Lookup (
   , prepLookups
   , bloomQueries
   , indexSearches
-  , intraPageLookups
+  , intraPageLookupsWithWriteBuffer
+  , intraPageLookupsOn
   ) where
 
 import           Data.Arena (Arena, ArenaManager, allocateFromArena, withArena)
@@ -114,7 +115,7 @@ type ResolveSerialisedValue = SerialisedValue -> SerialisedValue -> SerialisedVa
 
 type LookupAcc m h = V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m h)))
 
-{-# SPECIALIZE lookupsIO ::
+{-# SPECIALIZE lookupsIOWithWriteBuffer ::
        HasBlockIO IO h
     -> ArenaManager RealWorld
     -> ResolveSerialisedValue
@@ -127,13 +128,8 @@ type LookupAcc m h = V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m 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.
-lookupsIO ::
+-- | Like 'lookupsIO', but takes a write buffer into account.
+lookupsIOWithWriteBuffer ::
      forall m h. (MonadThrow m, MonadST m)
   => HasBlockIO m h
   -> ArenaManager (PrimState m)
@@ -146,12 +142,12 @@ lookupsIO ::
   -> V.Vector (Handle h) -- ^ The file handles to the key\/value files inside @rs@
   -> V.Vector SerialisedKey
   -> m (LookupAcc m h)
-lookupsIO !hbio !mgr !resolveV !wb !wbblobs !rs !blooms !indexes !kopsFiles !ks =
+lookupsIOWithWriteBuffer !hbio !mgr !resolveV !wb !wbblobs !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
-      intraPageLookups resolveV wb wbblobs rs ks rkixs ioops ioress
+      intraPageLookupsWithWriteBuffer resolveV wb wbblobs rs ks rkixs ioops ioress
   where
     -- we check only that the lengths match, because checking the contents is
     -- too expensive.
@@ -161,7 +157,7 @@ lookupsIO !hbio !mgr !resolveV !wb !wbblobs !rs !blooms !indexes !kopsFiles !ks
       assert (V.length rs == V.length kopsFiles) $
       True
 
-{-# SPECIALIZE lookupsIOWithoutWriteBuffer ::
+{-# SPECIALIZE lookupsIO ::
        HasBlockIO IO h
     -> ArenaManager RealWorld
     -> ResolveSerialisedValue
@@ -174,11 +170,9 @@ lookupsIO !hbio !mgr !resolveV !wb !wbblobs !rs !blooms !indexes !kopsFiles !ks
   #-}
 -- | 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 ::
+lookupsIO ::
      forall m h. (MonadThrow m, MonadST m)
   => HasBlockIO m h
   -> ArenaManager (PrimState m)
@@ -189,7 +183,7 @@ lookupsIOWithoutWriteBuffer ::
   -> 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 =
+lookupsIO !hbio !mgr !resolveV !rs !blooms !indexes !kopsFiles !ks =
     assert precondition $
     withArena mgr $ \arena -> do
       (rkixs, ioops) <- ST.stToIO $ prepLookups arena blooms indexes kopsFiles ks
@@ -204,7 +198,7 @@ lookupsIOWithoutWriteBuffer !hbio !mgr !resolveV !rs !blooms !indexes !kopsFiles
       assert (V.length rs == V.length kopsFiles) $
       True
 
-{-# SPECIALIZE intraPageLookups ::
+{-# SPECIALIZE intraPageLookupsWithWriteBuffer ::
        ResolveSerialisedValue
     -> WB.WriteBuffer
     -> Ref (WBB.WriteBufferBlobs IO h)
@@ -215,14 +209,10 @@ lookupsIOWithoutWriteBuffer !hbio !mgr !resolveV !rs !blooms !indexes !kopsFiles
     -> 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.
+-- | Like 'intraPageLookupsOn', but uses the write buffer as the initial
+-- accumulator.
 --
-intraPageLookups ::
+intraPageLookupsWithWriteBuffer ::
      forall m h. (PrimMonad m, MonadThrow m)
   => ResolveSerialisedValue
   -> WB.WriteBuffer
@@ -233,21 +223,9 @@ intraPageLookups ::
   -> V.Vector (IOOp (PrimState m) h)
   -> VU.Vector IOResult
   -> 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
-    -- 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.
-    --
-
+intraPageLookupsWithWriteBuffer !resolveV !wb !wbblobs !rs !ks !rkixs !ioops !ioress = do
+    -- The most recent values are in the write buffer, so we use it to
+    -- initialise the accumulator.
     acc0 <-
       V.generateM (V.length ks) $ \ki ->
         case WB.lookup wb (V.unsafeIndex ks ki) of
@@ -278,7 +256,7 @@ data TableCorruptedError
     -> IO (LookupAcc IO h)
   #-}
 -- | Intra-page lookups, and combining lookup results from multiple runs and
--- the write buffer.
+-- a potential initial accumulator (e.g. from 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
@@ -287,7 +265,7 @@ data TableCorruptedError
 intraPageLookupsOn ::
      forall m h. (PrimMonad m, MonadThrow m)
   => ResolveSerialisedValue
-  -> LookupAcc m h
+  -> LookupAcc m h  -- initial acc
   -> V.Vector (Ref (Run m h))
   -> V.Vector SerialisedKey
   -> VP.Vector RunIxKeyIx

src/Database/LSMTree/Internal/MergingRun.hs

@@ -258,6 +258,10 @@ unsafeNew 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.
+--
+-- TODO: This interface doesn't work well with the action registry. Just doing
+-- @withRollback reg (duplicateRuns mr) (mapM_ releaseRef)@ isn't exception-safe
+-- since if one of the @releaseRef@ calls fails, the following ones aren't run.
 {-# SPECIALISE duplicateRuns ::
      Ref (MergingRun t IO h) -> IO (V.Vector (Ref (Run IO h))) #-}
 duplicateRuns ::
@@ -757,8 +761,8 @@ atomicSpendCredits (CreditsVar var) spend =
 
 {-# SPECIALISE remainingMergeDebt ::
      Ref (MergingRun t IO h) -> IO (MergeDebt, NumEntries) #-}
--- | Calculate an upper bound on the merge credits required to complete the
--- merge, as well as an upper bound on the size of the resulting run.
+-- | Calculate the merge credits required to complete the merge, as well as an
+-- upper bound on the size of the resulting run.
 remainingMergeDebt ::
      (MonadMVar m, PrimMonad m)
   => Ref (MergingRun t m h) -> m (MergeDebt, NumEntries)

src/Database/LSMTree/Internal/MergingTree/Lookup.hs

@@ -81,9 +81,6 @@ mergeLookupAcc resolve mt accs =
 --
 -- This function duplicates the references to all the tree's runs.
 -- These references later need to be released using 'releaseLookupTree'.
---
--- This function should be run with asynchronous exceptions masked to prevent
--- failing after internal resources have already been created.
 {-# SPECIALISE buildLookupTree ::
      ActionRegistry IO
   -> Ref (MT.MergingTree IO h)

test/Test/Database/LSMTree/Internal/Lookup.hs

@@ -319,7 +319,7 @@ prop_roundtripFromWriteBufferLookupIO (SmallList dats) =
     arenaManager <- newArenaManager
     realres <-
       fetchBlobs hfs =<< -- retrieve blobs to match type of model result
-      lookupsIO
+      lookupsIOWithWriteBuffer
         hbio
         arenaManager
         resolveV
@@ -339,8 +339,9 @@ prop_roundtripFromWriteBufferLookupIO (SmallList dats) =
     fetchBlobs hfs = traverse (traverse (traverse (readWeakBlobRef hfs)))
 
 -- | Given a bunch of 'InMemLookupData', prepare the data into the form needed
--- for 'lookupsIO': a write buffer (and blobs) and a vector of on-disk runs.
--- Also passes the model and the keys to look up to the inner action.
+-- for 'lookupsIOWithWriteBuffer': a write buffer (and blobs) and a vector of
+-- on-disk runs. Also passes the model and the keys to look up to the inner
+-- action.
 --
 withWbAndRuns :: FS.HasFS IO h
          -> FS.HasBlockIO IO h

test/Test/Database/LSMTree/Internal/MergingTree.hs

@@ -14,6 +14,7 @@ import           Data.Map (Map)
 import qualified Data.Map as Map
 import           Data.Traversable (for)
 import qualified Data.Vector as V
+import           Database.LSMTree.Extras (showPowersOf10)
 import           Database.LSMTree.Extras.MergingRunData
 import           Database.LSMTree.Extras.MergingTreeData
 import           Database.LSMTree.Extras.RunData
@@ -138,8 +139,9 @@ prop_lookupTree hfs hbio keys mtd = do
       -> IO (V.Vector (Maybe (Entry v SerialisedBlob)))
     fetchBlobs = traverse (traverse (traverse (readWeakBlobRef hfs)))
 
-    -- trees are always in the last level, there is no distinction between
-    -- (Nothing and Just Delete), (Insert and Mupsert)
+    -- the lookup accs might be different between implementation and model
+    -- (Nothing vs. Just Delete, Insert vs. Mupsert), but this doesn't matter
+    -- for the final result of the lookup
     normalise = V.map toLookupResult
 
     toLookupResult Nothing  = Nothing
@@ -156,13 +158,14 @@ prop_lookupTree hfs hbio keys mtd = do
             return $ V.map (const Nothing) keys
           False -> do
             batches <- buildLookupTree reg tree
-            releaseLookupTree reg batches  -- only happens at the end
             results <- traverse (performLookups mgr) batches
-            foldLookupTree resolveVal results
+            acc <- foldLookupTree resolveVal results
+            releaseLookupTree reg batches
+            return acc
 
     performLookups mgr runs =
         Async.async $
-          Lookup.lookupsIOWithoutWriteBuffer
+          Lookup.lookupsIO
             hbio
             mgr
             resolveVal
@@ -227,22 +230,38 @@ prop_supplyCredits hfs hbio threshold credits mtd = do
     FS.createDirectory hfs (FS.mkFsPath ["active"])
     counter <- newUniqCounter 0
     withMergingTree hfs hbio resolveVal runParams setupPath counter mtd $ \tree -> do
+      (MR.MergeDebt initialDebt, _) <- remainingMergeDebt tree
       props <- for credits $ \c -> do
         (MR.MergeDebt debt, _) <- remainingMergeDebt tree
-        leftovers <-
-          supplyCredits hfs hbio resolveVal runParams threshold root counter tree c
-        (MR.MergeDebt debt', _) <- remainingMergeDebt tree
-        return $
-          counterexample (show (debt, leftovers, debt')) $ conjoin [
-              counterexample "negative values" $
-                debt >= 0 && leftovers >= 0 && debt' >= 0
-            , counterexample "did not reduce debt sufficiently" $
-                debt' <= debt - (c - leftovers)
-            ]
-      return (conjoin (toList props))
+        if debt <= 0
+          then
+            return $ property True
+          else do
+            leftovers <-
+              supplyCredits hfs hbio resolveVal runParams threshold root counter tree c
+            (MR.MergeDebt debt', _) <- remainingMergeDebt tree
+            return $
+              -- semi-useful, but mainly tells us in how many steps we supplied
+              tabulate "supplied credits" [showPowersOf10 (fromIntegral c)] $
+              counterexample (show (debt, leftovers, debt')) $ conjoin [
+                  counterexample "negative values" $
+                    debt >= 0 && leftovers >= 0 && debt' >= 0
+                , counterexample "did not reduce debt sufficiently" $
+                    debt' <= debt - (c - leftovers)
+                ]
+      (MR.MergeDebt finalDebt, _) <- remainingMergeDebt tree
+      return $
+        labelDebt initialDebt finalDebt $
+          conjoin (toList props)
   where
     root = Paths.SessionRoot (FS.mkFsPath [])
-    setupPath = FS.mkFsPath ["setup"]  -- separate dir, so it doesn't clash
+    setupPath = FS.mkFsPath ["setup"]  -- separate dir, so file paths in errors
+                                       -- are identifiable as created in setup
+                                       --
+    labelDebt initial final
+      | initial == 0 = label "trivial"
+      | final   == 0 = label "completed"
+      | otherwise    = label "incomplete"
 
 instance Arbitrary MR.MergeCredits where
   arbitrary = MR.MergeCredits . getPositive <$> arbitrary