Return leftover credits from supplyCredits

This is now easy because it's reported as part of the credit accounting
in a reliable way.

Author: dcoutts

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -746,7 +746,7 @@ addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root uc r0 reg levels ul
           OneShot -> do
             let !required = MR.Credits (unNumEntries (V.foldMap' Run.size rs))
             let !thresh = creditThresholdForLevel conf ln
-            MR.supplyCredits mr thresh required
+            _leftoverCredits <- MR.supplyCredits mr thresh required
             -- 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 ->
@@ -877,7 +877,8 @@ supplyCredits conf c levels =
         Merging mp mr -> do
           let !c' = scaleCreditsForMerge mp mr c
           let !thresh = creditThresholdForLevel conf ln
-          MR.supplyCredits mr thresh c'
+          _leftoverCredits <- MR.supplyCredits mr thresh c'
+          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

@@ -628,15 +628,15 @@ atomicSpendCredits (CreditsVar var) spend =
      Ref (MergingRun IO h)
   -> CreditThreshold
   -> Credits
-  -> 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)
   => Ref (MergingRun m h)
   -> CreditThreshold
   -> Credits
-  -> m ()
+  -> m Credits
 supplyCredits (DeRef MergingRun {
                  mergeKnownCompleted,
                  mergeNumEntries,
@@ -647,7 +647,7 @@ supplyCredits (DeRef MergingRun {
     assert (credits >= 0) $ do
     mergeCompleted <- readMutVar mergeKnownCompleted
     case mergeCompleted of
-      MergeKnownCompleted -> pure ()
+      MergeKnownCompleted -> pure credits
       MergeMaybeCompleted ->
         bracketOnError
           -- Atomically add credits to the unspent credits (but not allowing
@@ -665,7 +665,7 @@ supplyCredits (DeRef MergingRun {
           (\(spendCredits, _leftoverCredits) ->
             atomicSpendCredits mergeCreditsVar (-spendCredits))
 
-          (\(spendCredits, _leftoverCredits) ->
+          (\(spendCredits, leftoverCredits) -> do
             when (spendCredits > 0) $ do
               weFinishedMerge <-
                 performMergeSteps mergeState mergeCreditsVar spendCredits
@@ -674,7 +674,9 @@ supplyCredits (DeRef MergingRun {
               -- completion, then that is fine. The next supplyCredits will
               -- complete the merge.
               when weFinishedMerge $
-                completeMerge mergeState mergeKnownCompleted)
+                completeMerge mergeState mergeKnownCompleted
+
+            return leftoverCredits)
 
 {-# SPECIALISE performMergeSteps ::
      StrictMVar IO (MergingRunState IO h)

src/Database/LSMTree/Internal/Snapshot.hs

@@ -28,6 +28,7 @@ import           Control.ActionRegistry
 import           Control.Concurrent.Class.MonadMVar.Strict
 import           Control.Concurrent.Class.MonadSTM (MonadSTM)
 import           Control.DeepSeq (NFData (..))
+import           Control.Exception (assert)
 import           Control.Monad (void)
 import           Control.Monad.Class.MonadST (MonadST)
 import           Control.Monad.Class.MonadThrow (MonadMask)
@@ -462,9 +463,10 @@ fromSnapLevels reg hfs hbio conf@TableConfig{..} uc resolve dir (SnapLevels leve
                 -- When a snapshot is created, merge progress is lost, so we
                 -- have to redo merging work here. SuppliedCredits tracks how
                 -- many credits were supplied before the snapshot was taken.
-                MR.supplyCredits mr (creditThresholdForLevel conf ln)
-                                 (MR.Credits sc)
-                return mr
+                leftoverCredits <- MR.supplyCredits
+                                     mr (creditThresholdForLevel conf ln)
+                                     (MR.Credits sc)
+                assert (leftoverCredits == 0) $ return mr
 
     dupRun r = withRollback reg (dupRef r) releaseRef