Merge pull request #426 from IntersectMBO/jdral/scheduled-merges

Base implementation of scheduled merges

Author: jorisdral

src-control/Control/RefCount.hs

@@ -10,6 +10,7 @@ module Control.RefCount (
   , mkRefCounter1
   , addReference
   , removeReference
+  , removeReferenceN
   , upgradeWeakReference
   , readRefCount
   ) where
@@ -21,7 +22,9 @@ import           Control.Monad.Class.MonadThrow
 import           Control.Monad.Primitive
 import           Data.Maybe
 import           Data.Primitive.PrimVar
+import           Data.Word
 import           GHC.Stack
+import           Text.Printf
 
 -- | A reference counter with an optional finaliser action. Once the reference
 -- count reaches @0@, the finaliser will be run.
@@ -89,6 +92,43 @@ removeReference RefCounter{countVar, finaliser} = mask_ $ do
     assertWithCallStack (prevCount > 0) $ pure ()
     when (prevCount == 1) $ sequence_ finaliser
 
+-- TODO: remove uses of this API. Eventually all references should be singular,
+-- and not use patterns where if A contains B then N references on A becomes N
+-- references on B. Instead this should be a single reference from A to B,
+-- irrespective of the number of references to A.
+{-# SPECIALISE removeReferenceN :: HasCallStack => RefCounter IO -> Word64 -> IO () #-}
+-- | Decrease the reference counter by @n@. @n@ must be a positive number.
+--
+-- The count must be known (from context) to be non-zero and at least as large
+-- as @n@. Typically this will be because the caller has @n@ references already
+-- (that they took out themselves or were given).
+removeReferenceN :: (HasCallStack, PrimMonad m, MonadMask m) => RefCounter m -> Word64 -> m ()
+removeReferenceN RefCounter{countVar, finaliser} n = mask_ $ do
+    -- n should be positive
+    assert (n > 0) $ pure ()
+    let !n' = fromIntegralChecked n
+    prevCount <- fetchSubInt countVar n'
+    -- the reference count must not already be 0, because then the finaliser
+    -- will have run already
+    assertWithCallStack (prevCount > 0) $ pure ()
+    -- the reference count can not go below zero
+    assertWithCallStack (prevCount >= n') $ pure ()
+    when (prevCount <= n') $ sequence_ finaliser
+
+-- TODO: remove when removeReferenceN is removed
+{-# INLINABLE fromIntegralChecked #-}
+-- | Like 'fromIntegral', but throws an error when @(x :: a) /= fromIntegral
+-- (fromIntegral x :: b)@.
+fromIntegralChecked :: (HasCallStack, Integral a, Integral b, Show a) => a -> b
+fromIntegralChecked x
+  | x'' == x
+  = x'
+  | otherwise
+  = error $ printf "fromIntegralChecked: conversion failed, %s /= %s" (show x) (show x'')
+  where
+    x' = fromIntegral x
+    x'' = fromIntegral x'
+
 -- | Try to turn a \"weak\" reference on something into a proper reference.
 -- This is by analogy with @deRefWeak :: Weak v -> IO (Maybe v)@, but for
 -- reference counts.
@@ -120,6 +160,7 @@ upgradeWeakReference RefCounter{countVar} = do
             then return True
             else casLoop prevCount'
 
+-- TODO: remove when removeRefenceN is removed
 {-# SPECIALISE readRefCount :: RefCounter IO -> IO RefCount #-}
 -- | Warning: reading the current reference count is inherently racy as there is
 -- no way to reliably act on the information. It can be useful for debugging.

src-extras/Database/LSMTree/Extras/NoThunks.hs

@@ -277,6 +277,12 @@ deriving stock instance Generic (MergingRunState m h)
 deriving anyclass instance (Typeable m, Typeable (PrimState m), Typeable h)
                         => NoThunks (MergingRunState m h)
 
+deriving stock instance Generic MergePolicyForLevel
+deriving anyclass instance NoThunks MergePolicyForLevel
+
+deriving stock instance Generic NumRuns
+deriving anyclass instance NoThunks NumRuns
+
 {-------------------------------------------------------------------------------
   Entry
 -------------------------------------------------------------------------------}

src/Database/LSMTree/Internal.hs

@@ -96,7 +96,7 @@ import           Database.LSMTree.Internal.BlobRef (WeakBlobRef (..))
 import qualified Database.LSMTree.Internal.BlobRef as BlobRef
 import           Database.LSMTree.Internal.Config
 import qualified Database.LSMTree.Internal.Cursor as Cursor
-import           Database.LSMTree.Internal.Entry (Entry)
+import           Database.LSMTree.Internal.Entry (Entry, unNumEntries)
 import           Database.LSMTree.Internal.Lookup (ByteCountDiscrepancy,
                      ResolveSerialisedValue, lookupsIO)
 import           Database.LSMTree.Internal.MergeSchedule
@@ -1108,6 +1108,12 @@ snapshot resolve snap label th = do
                     (RW.unsafeAcquireWriteAccess (tableContent thEnv))
                     (atomically . RW.unsafeReleaseWriteAccess (tableContent thEnv))
                     $ \reg content -> do
+        -- TODO: When we flush the buffer here, it might be underfull, which
+        -- could mess up the scheduling. The conservative approach is to supply
+        -- credits as if the buffer was full, and then flush the (possibly)
+        -- underfull buffer. However, note that this bit of code
+        -- here is probably going to change anyway because of #392
+        supplyCredits (unNumEntries $ case confWriteBufferAlloc conf of AllocNumEntries x -> x) (tableLevels content)
         content' <- flushWriteBuffer
               (TraceMerge `contramap` tableTracer th)
               conf
@@ -1140,6 +1146,7 @@ snapshot resolve snap label th = do
   -> SnapshotLabel
   -> TableConfigOverride
   -> SnapshotName
+  -> ResolveSerialisedValue
   -> IO (TableHandle IO h) #-}
 -- |  See 'Database.LSMTree.Normal.open'.
 open ::
@@ -1148,8 +1155,9 @@ open ::
   -> SnapshotLabel -- ^ Expected label
   -> TableConfigOverride -- ^ Optional config override
   -> SnapshotName
+  -> ResolveSerialisedValue
   -> m (TableHandle m h)
-open sesh label override snap = do
+open sesh label override snap resolve = do
     traceWith (sessionTracer sesh) $ TraceOpenSnapshot snap override
     withOpenSession sesh $ \seshEnv -> do
       withTempRegistry $ \reg -> do
@@ -1173,7 +1181,7 @@ open sesh label override snap = do
           <- allocateTemp reg
                (WBB.new hfs blobpath)
                WBB.removeReference
-        tableLevels <- openLevels reg hfs hbio conf (sessionRoot seshEnv) snappedLevels
+        tableLevels <- openLevels reg hfs hbio conf (sessionUniqCounter seshEnv) (sessionRoot seshEnv) resolve snappedLevels
         tableCache <- mkLevelsCache reg tableLevels
         newWith reg sesh seshEnv conf' am $! TableContent {
             tableWriteBuffer = WB.empty

src/Database/LSMTree/Internal/Config.hs

@@ -409,8 +409,6 @@ instance NFData MergeSchedule where
 -- | The default 'MergeSchedule'.
 --
 -- >>> defaultMergeSchedule
--- OneShot
---
--- TODO: replace by 'Incremental'
+-- Incremental
 defaultMergeSchedule :: MergeSchedule
-defaultMergeSchedule = OneShot
+defaultMergeSchedule = Incremental

src/Database/LSMTree/Internal/Merge.hs

@@ -9,6 +9,8 @@ module Database.LSMTree.Internal.Merge (
   , new
   , addReference
   , removeReference
+  , removeReferenceN
+  , readRefCount
   , complete
   , stepsToCompletion
   , stepsToCompletionCounted
@@ -24,12 +26,13 @@ import           Control.Monad.Class.MonadThrow (MonadCatch, MonadMask (..),
                      MonadThrow (..))
 import           Control.Monad.Fix (MonadFix)
 import           Control.Monad.Primitive (PrimMonad, PrimState, RealWorld)
-import           Control.RefCount (RefCounter)
+import           Control.RefCount (RefCount (..), RefCounter)
 import qualified Control.RefCount as RC
 import           Data.Coerce (coerce)
 import           Data.Primitive.MutVar
 import           Data.Traversable (for)
 import qualified Data.Vector as V
+import           Data.Word
 import           Database.LSMTree.Internal.BlobRef (BlobRef)
 import           Database.LSMTree.Internal.Entry
 import           Database.LSMTree.Internal.Run (Run, RunDataCaching)
@@ -130,6 +133,14 @@ addReference Merge{..} = RC.addReference mergeRefCounter
 removeReference :: (HasCallStack, PrimMonad m, MonadMask m) => Merge m h -> m ()
 removeReference Merge{..} = RC.removeReference mergeRefCounter
 
+{-# SPECIALISE removeReferenceN :: Merge IO h -> Word64 -> IO () #-}
+removeReferenceN :: (HasCallStack, PrimMonad m, MonadMask m) => Merge m h -> Word64 -> m ()
+removeReferenceN r = RC.removeReferenceN (mergeRefCounter r)
+
+{-# SPECIALISE readRefCount :: Merge IO h -> IO RefCount #-}
+readRefCount :: PrimMonad m => Merge m h -> m RefCount
+readRefCount Merge{..} = RC.readRefCount mergeRefCounter
+
 {-# SPECIALISE finaliser ::
      MutVar RealWorld MergeState
   -> RunBuilder IO h
@@ -239,6 +250,7 @@ stepsToCompletionCounted m stepBatchSize = go 0
                                 in (stepsSum',) <$> complete m
 
 data StepResult = MergeInProgress | MergeComplete
+  deriving stock Eq
 
 stepsInvariant :: Int -> (Int, StepResult) -> Bool
 stepsInvariant requestedSteps = \case