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

@@ -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 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 ->
@@ -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 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,15 +647,15 @@ supplyCredits (DeRef MergingRun {
     assert (credits >= 0) $ do
     mergeCompleted <- readMutVar mergeKnownCompleted
     case mergeCompleted of
-      MergeKnownCompleted -> pure ()
+      MergeKnownCompleted -> pure credits
       MergeMaybeCompleted -> do
         -- Atomically add credits to the unspent credits (but not allowing
         -- 'suppliedCredits' to exceed the total debt), determine which case
         -- we're in and thus how many credits we should try to spend now on
         -- performing merge steps. Return the credits to spend now and any
         -- leftover credits that would exceed the debt limit.
         (spendCredits,
-         _leftoverCredits) <- atomicDepositAndSpendCredits
+         leftoverCredits) <- atomicDepositAndSpendCredits
                                mergeCreditsVar
                                (numEntriesToTotalDebt mergeNumEntries)
                                creditBatchThreshold credits
@@ -673,6 +673,8 @@ supplyCredits (DeRef MergingRun {
           -- complete the merge.
           when weFinishedMerge $ completeMerge mergeState mergeKnownCompleted
 
+        return leftoverCredits
+
 {-# SPECIALISE performMergeSteps ::
      StrictMVar IO (MergingRunState IO h)
   -> CreditsVar RealWorld

src/Database/LSMTree/Internal/Snapshot.hs

@@ -29,6 +29,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, when)
 import           Control.Monad.Class.MonadST (MonadST)
 import           Control.Monad.Class.MonadThrow (MonadMask)
@@ -469,9 +470,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