TOSQUASH

Author: jorisdral

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -915,18 +915,33 @@ supplyMergeCredits (ScaledCredits c) (MergingRun _ _ var) = do
 
         -- If we previously performed too many merge steps, then we perform
         -- fewer now.
-        let stepsToDo = totalCredits - totalSteps + c
-        -- Merge.steps guarantees that stepsDone >= stepsToDo
+        let stepsToDo = max 0 (totalCredits + c - totalSteps)
+        -- Merge.steps guarantees that stepsDone >= stepsToDo /unless/ the merge
+        -- was just now finished in fewer than stepsToDo steps, in which case
+        -- stepsDone <= stepsToDo.
         (stepsDone, stepResult) <- Merge.steps m stepsToDo
+        assert (case stepResult of
+                  MergeInProgress -> stepsDone >= stepsToDo
+                  MergeComplete   -> stepsDone <= stepsToDo
+               ) $ pure ()
 
         -- This should be the only point at which we write to these variables.
-        -- Since we we never perform fewer merge steps than we requested, the
-        -- invariant is that @readPrimVar totalStepsVar >= readPrimVar
-        -- totalCreditsVar@. That is also why we update totalStepsVar before
-        -- totalCreditsVar, in case an asynchronous exception is thrown in
-        -- between the writes.
+        --
+        -- Unless the merge was finished in fewer than stepsToDo steps, we can
+        -- guarantee that totalSteps' >= totalCredits'
+        let totalSteps' = totalSteps + stepsDone
+        let totalCredits' = totalCredits + c
+        -- If an exception happens between the next two writes, then only
+        -- totalCreditsVar will be outdated, which is okay because we will
+        -- resupply credits. It also means we can maintain that @readPrimVar
+        -- totalStepsVar >= readPrimVar totalCreditsVar@, unless the merge
+        -- finished in fewer than stepsToDo steps.
         writePrimVar totalStepsVar $! totalSteps + stepsDone
         writePrimVar totalCreditsVar $! totalCredits + c
+        assert (case stepResult of
+                  MergeInProgress -> totalSteps' >= totalCredits'
+                  MergeComplete   -> totalSteps' <= totalCredits'
+               ) $ pure ()
 
         pure $ stepResult == MergeComplete
     when mergeIsDone $