Make MergingRun.supplyCredits return the leftover credits

This will be needed for supplying credits to merging trees.

Author: dcoutts

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -673,7 +673,7 @@ addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root uc r0 reg levels =
           OneShot -> do
             let !required = MR.Credits (unNumEntries (V.foldMap' Run.size rs))
             let !thresh = creditThresholdForLevel conf ln
-            MR.supplyCredits required thresh mr
+            _leftoverCredits <- MR.supplyCredits required thresh mr
             -- This ensures the merge is really completed. However, we don't
             -- release the merge yet and only briefly inspect the resulting run.
             bracket (MR.expectCompleted mr) releaseRef $ \r ->
@@ -802,7 +802,8 @@ supplyCredits conf c levels =
         Merging mp mr -> do
           let !c' = scaleCreditsForMerge mp mr c
           let !thresh = creditThresholdForLevel conf ln
-          MR.supplyCredits c' thresh mr
+          _leftoverCredits <- MR.supplyCredits c' thresh mr
+          return ()
 
 -- | Scale a number of credits to a number of merge steps to be performed, based
 -- on the merging run.

src/Database/LSMTree/Internal/MergingRun.hs

@@ -353,24 +353,27 @@ writeStepsPerformed (StepsPerformedVar v) x = writePrimVar v x
      Credits
   -> CreditThreshold
   -> Ref (MergingRun IO h)
-  -> IO () #-}
+  -> IO Credits #-}
 -- | Supply the given amount of credits to a merging run. This /may/ cause an
 -- ongoing merge to progress.
 supplyCredits ::
      forall m h. (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)
   => Credits
   -> CreditThreshold
   -> Ref (MergingRun m h)
-  -> m ()
+  -> m Credits
 supplyCredits (Credits c) creditsThresh (DeRef MergingRun {..}) = do
     mergeCompleted <- readMutVar mergeKnownCompleted
 
     -- The merge is already finished
     if mergeCompleted == MergeKnownCompleted then
-      pure ()
+      -- we're already done so all supplied credits are leftovers
+      pure (Credits c)
     else do
       -- unspentCredits' is our /estimate/ of what the new total of unspent
-      -- credits is.
+      -- credits is. It is an estimate due to potential concurrency: other
+      -- threads may also by supplying credits. We only check the credit
+      -- balance once, which is the result of addUnspentCredits here:
       Credits unspentCredits' <- addUnspentCredits mergeUnspentCreditsVar (Credits c)
       stepsPerformed <- readStepsPerformed mergeStepsPerformedVar
 
@@ -382,6 +385,9 @@ supplyCredits (Credits c) creditsThresh (DeRef MergingRun {..}) = do
                          (stepMerge mergeSpentCreditsVar mergeStepsPerformedVar
                                     mergeState)
         when isMergeDone $ completeMerge mergeState mergeKnownCompleted
+        let !leftoverCredits = stepsPerformed + unspentCredits'
+                             - unNumEntries mergeNumEntries
+        return (Credits leftoverCredits)
       else if unspentCredits' >= getCreditThreshold creditsThresh then do
         -- We can do some merging work without finishing the merge immediately
         isMergeDone <-
@@ -405,9 +411,14 @@ supplyCredits (Credits c) creditsThresh (DeRef MergingRun {..}) = do
         -- completion, then that is fine. The next supplyCredits will
         -- complete the merge.
         when isMergeDone $ completeMerge mergeState mergeKnownCompleted
+        -- We return 0 credits here, even if we completed the merge. This is
+        -- because if we did complete the merge, it was due to other threads
+        -- supplying credits, not due to the credits we supplied here.
+        return (Credits 0)
       else
-        -- Just accumulate credits, because we are not over the threshold yet
-        pure ()
+        -- Just accumulate credits, because we are not over the threshold yet.
+        -- Thus no leftover credits.
+        pure (Credits 0)
 
 {-# SPECIALISE addUnspentCredits ::
      UnspentCreditsVar RealWorld
@@ -512,11 +523,12 @@ stepMerge spentCreditsVar stepsPerformedVar mergeVar (Credits c) =
       CompletedMerge{} -> pure False
       (OngoingMerge _rs m) -> do
         stepsPerformed <- readStepsPerformed stepsPerformedVar
-        spentCredits <- readSpentCredits spentCreditsVar
+        spentCredits   <- readSpentCredits   spentCreditsVar
 
         -- If we previously performed too many merge steps, then we perform
         -- fewer now.
-        let stepsToDo = max 0 (spentCredits + c - stepsPerformed)
+        let stepsSurplus = stepsPerformed - spentCredits
+            stepsToDo    = max 0 (c - stepsSurplus)
         -- Merge.steps guarantees that @stepsDone >= stepsToDo@ /unless/ the
         -- merge was just now finished.
         (stepsDone, stepResult) <- Merge.steps m stepsToDo

src/Database/LSMTree/Internal/Snapshot.hs

@@ -465,7 +465,11 @@ fromSnapLevels reg hfs hbio conf@TableConfig{..} uc resolve dir (SnapLevels leve
                 -- snapshot was taken.
                 let c = getUnspentCredits unspentCredits
                       + getSpentCredits spentCredits
-                MR.supplyCredits (MR.Credits c) (creditThresholdForLevel conf ln) mr
+                _c' <- MR.supplyCredits (MR.Credits c)
+                                        (creditThresholdForLevel conf ln) mr
+                -- We would typically expect the leftover credits c' to be 0,
+                -- but it's hard to reason that this is guaranteed, so we do
+                -- not assert it here.
                 return mr
 
     dupRun r = withRollback reg (dupRef r) releaseRef