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

@@ -687,33 +687,33 @@ atomicUnspendCredits (CreditsVar var) unspend =
      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 mrun@MergingRun {mergeKnownCompleted})
               !creditBatchThreshold !credits = do
     mergeCompleted <- readMutVar mergeKnownCompleted
     case mergeCompleted of
-      MergeKnownCompleted -> pure ()
+      MergeKnownCompleted -> pure credits
       MergeMaybeCompleted -> supplyCredits' mrun creditBatchThreshold credits
 
 {-# SPECIALISE supplyCredits' ::
      MergingRun IO h
   -> CreditThreshold
   -> Credits
-  -> IO () #-}
+  -> IO Credits #-}
 supplyCredits' ::
      forall m h. (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)
   => MergingRun m h
   -> CreditThreshold
   -> Credits
-  -> m ()
+  -> m Credits
 supplyCredits' MergingRun {
                  mergeNumEntries,
                  mergeCreditsVar,
@@ -763,6 +763,8 @@ supplyCredits' MergingRun {
         -- completion, then that is fine. The next supplyCredits will
         -- complete the merge.
         when weFinishedMerge $ completeMerge mergeState mergeKnownCompleted
+        -- We reliably know the credits we supplied that were left over.
+        return leftoverCredits
 
       | unspentCredits >= batchThreshold -> do
         -- If the unspent credits have reached the threshold then we will try
@@ -784,10 +786,14 @@ supplyCredits' MergingRun {
         -- completion, then that is fine. The next supplyCredits will
         -- complete the merge.
         when weFinishedMerge $ completeMerge mergeState mergeKnownCompleted
+        -- We didn't finish, so can be no leftover credits
+        assert (leftoverCredits == 0) $ pure (Credits 0)
 
       -- Otherwise just accumulate credits (which we did already above),
       -- because we are not over the threshold yet.
-      | otherwise -> pure ()
+      | otherwise ->
+        -- We didn't finish, so can be no leftover credits
+        assert (leftoverCredits == 0) $ pure (Credits 0)
 
 {-# SPECIALISE performMergeSteps ::
      StrictMVar IO (MergingRunState IO h)

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