Refactor: move spent credits var up from MergingRunState to MergingRun

The MergingRun contains the other two credit tracking vars. It makes
more sense for all three to live together in once place. And there's no
need for one of the vars to disappear when we move from the OngoingMerge
state to the CompletedMerge state.

Author: dcoutts

src/Database/LSMTree/Internal/MergingRun.hs

@@ -55,6 +55,8 @@ data MergingRun m h = MergingRun {
     , mergeNumEntries     :: !NumEntries
       -- | The number of currently /unspent/ credits
     , mergeUnspentCredits :: !(UnspentCreditsVar (PrimState m))
+      -- | The total number of spent credits.
+    , mergeSpentCredits   :: !(SpentCreditsVar (PrimState m))
       -- | The total number of performed merging steps.
     , mergeStepsPerformed :: !(TotalStepsVar (PrimState m))
       -- | A variable that caches knowledge about whether the merge has been
@@ -77,8 +79,12 @@ newtype UnspentCreditsVar s = UnspentCreditsVar {
     getUnspentCreditsVar :: PrimVar s Int
   }
 
+newtype SpentCreditsVar s = SpentCreditsVar {
+    getSpentCreditsVar :: PrimVar s Int
+  }
+
 newtype TotalStepsVar s = TotalStepsVar {
-    getTotalStepsVar ::  PrimVar s Int
+    getTotalStepsVar :: PrimVar s Int
   }
 
 data MergingRunState m h =
@@ -88,14 +94,8 @@ data MergingRunState m h =
   | OngoingMerge
       !(V.Vector (Ref (Run m h)))
       -- ^ Input runs
-      !(SpentCreditsVar (PrimState m))
-      -- ^ The total number of spent credits.
       !(Merge m h)
 
-newtype SpentCreditsVar s = SpentCreditsVar {
-    getSpentCreditsVar :: PrimVar s Int
-  }
-
 data MergeKnownCompleted = MergeKnownCompleted | MergeMaybeCompleted
   deriving stock (Show, Eq, Read)
 
@@ -139,9 +139,8 @@ new hfs hbio resolve caching alloc mergeType runPaths inputRuns =
         <$> Merge.new hfs hbio caching alloc mergeType resolve runPaths runs
       let numInputRuns = NumRuns $ V.length runs
       let numInputEntries = V.foldMap' Run.size runs
-      spentCreditsVar <- SpentCreditsVar <$> newPrimVar 0
       unsafeNew numInputRuns numInputEntries MergeMaybeCompleted $
-        OngoingMerge runs spentCreditsVar merge
+        OngoingMerge runs merge
 
 {-# SPECIALISE newCompleted ::
      NumRuns
@@ -176,6 +175,7 @@ unsafeNew ::
   -> m (Ref (MergingRun m h))
 unsafeNew mergeNumRuns mergeNumEntries knownCompleted state = do
     mergeUnspentCredits <- UnspentCreditsVar <$> newPrimVar 0
+    mergeSpentCredits   <- SpentCreditsVar <$> newPrimVar 0
     mergeStepsPerformed <- TotalStepsVar <$> newPrimVar 0
     case state of
       OngoingMerge{}   -> assert (knownCompleted == MergeMaybeCompleted) (pure ())
@@ -187,6 +187,7 @@ unsafeNew mergeNumRuns mergeNumEntries knownCompleted state = do
         mergeNumRuns
       , mergeNumEntries
       , mergeUnspentCredits
+      , mergeSpentCredits
       , mergeStepsPerformed
       , mergeKnownCompleted
       , mergeState
@@ -196,7 +197,7 @@ unsafeNew mergeNumRuns mergeNumEntries knownCompleted state = do
     finalise var = withMVar var $ \case
         CompletedMerge r ->
           releaseRef r
-        OngoingMerge rs _ m -> do
+        OngoingMerge rs m -> do
           V.forM_ rs releaseRef
           Merge.abort m
 
@@ -211,8 +212,8 @@ duplicateRuns (DeRef mr) =
     -- We take the references while holding the MVar to make sure the MergingRun
     -- does not get completed concurrently before we are done.
     withMVar (mergeState mr) $ \case
-      CompletedMerge r    -> V.singleton <$> dupRef r
-      OngoingMerge rs _ _ -> withActionRegistry $ \reg ->
+      CompletedMerge r  -> V.singleton <$> dupRef r
+      OngoingMerge rs _ -> withActionRegistry $ \reg ->
         V.mapM (\r -> withRollback reg (dupRef r) releaseRef) rs
 
 {-------------------------------------------------------------------------------
@@ -314,7 +315,8 @@ supplyCredits (Credits c) creditsThresh (DeRef MergingRun {..}) = do
         isMergeDone <-
           bracketOnError (takeAllUnspentCredits mergeUnspentCredits)
                          (putBackUnspentCredits mergeUnspentCredits)
-                         (stepMerge mergeState mergeStepsPerformed)
+                         (stepMerge mergeSpentCredits mergeStepsPerformed
+                                    mergeState)
         when isMergeDone $ completeMerge mergeState mergeKnownCompleted
       else if unspentCredits' >= getCreditThreshold creditsThresh then do
         -- We can do some merging work without finishing the merge immediately
@@ -329,7 +331,8 @@ supplyCredits (Credits c) creditsThresh (DeRef MergingRun {..}) = do
             (tryTakeUnspentCredits mergeUnspentCredits creditsThresh (Credits unspentCredits'))
             (mapM_ (putBackUnspentCredits mergeUnspentCredits)) $ \case
               Nothing -> pure False
-              Just c' -> stepMerge mergeState mergeStepsPerformed c'
+              Just c' -> stepMerge mergeSpentCredits mergeStepsPerformed
+                                   mergeState c'
 
         -- If we just finished the merge, then we convert the output of the
         -- merge into a new run. i.e., we complete the merge.
@@ -427,20 +430,24 @@ takeAllUnspentCredits (UnspentCreditsVar !unspentCreditsVar) = do
         casLoop prev'
 
 {-# SPECIALISE stepMerge ::
-     StrictMVar IO (MergingRunState IO h)
+     SpentCreditsVar RealWorld
   -> TotalStepsVar RealWorld
+  -> StrictMVar IO (MergingRunState IO h)
   -> Credits
   -> IO Bool #-}
 stepMerge ::
      (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)
-  => StrictMVar m (MergingRunState m h)
+  => SpentCreditsVar (PrimState m)
   -> TotalStepsVar (PrimState m)
+  -> StrictMVar m (MergingRunState m h)
   -> Credits
   -> m Bool
-stepMerge mergeVar (TotalStepsVar totalStepsVar) (Credits c) =
+stepMerge (SpentCreditsVar spentCreditsVar)
+          (TotalStepsVar totalStepsVar)
+          mergeVar (Credits c) =
     withMVar mergeVar $ \case
       CompletedMerge{} -> pure False
-      (OngoingMerge _rs (SpentCreditsVar spentCreditsVar) m) -> do
+      (OngoingMerge _rs m) -> do
         totalSteps <- readPrimVar totalStepsVar
         spentCredits <- readPrimVar spentCreditsVar
 
@@ -486,7 +493,7 @@ completeMerge ::
 completeMerge mergeVar mergeKnownCompletedVar = do
     modifyMVarMasked_ mergeVar $ \case
       mrs@CompletedMerge{} -> pure $! mrs
-      (OngoingMerge rs _spentCreditsVar m) -> do
+      (OngoingMerge rs m) -> do
         -- first try to complete the merge before performing other side effects,
         -- in case the completion fails
         r <- Merge.complete m
@@ -509,7 +516,8 @@ expectCompleted (DeRef MergingRun {..}) = do
     when (knownCompleted == MergeMaybeCompleted) $ do
       totalSteps <- readPrimVar (getTotalStepsVar mergeStepsPerformed)
       let !credits = Credits (unNumEntries mergeNumEntries - totalSteps)
-      isMergeDone <- stepMerge mergeState mergeStepsPerformed credits
+      isMergeDone <- stepMerge mergeSpentCredits mergeStepsPerformed
+                               mergeState credits
       when isMergeDone $ completeMerge mergeState mergeKnownCompleted
       -- TODO: can we think of a check to see if we did not do too much work
       -- here?

src/Database/LSMTree/Internal/Snapshot.hs

@@ -138,13 +138,13 @@ instance NFData r => NFData (SnapLevel r) where
   rnf (SnapLevel a b) = rnf a `seq` rnf b
 
 data SnapIncomingRun r =
-    SnapMergingRun !MergePolicyForLevel !NumRuns !NumEntries !UnspentCredits !(SnapMergingRunState r)
+    SnapMergingRun !MergePolicyForLevel !NumRuns !NumEntries !UnspentCredits !SpentCredits !(SnapMergingRunState r)
   | SnapSingleRun !r
   deriving stock (Show, Eq, Functor, Foldable, Traversable)
 
 instance NFData r => NFData (SnapIncomingRun r) where
-  rnf (SnapMergingRun a b c d e) =
-      rnf a `seq` rnf b `seq` rnf c `seq` rnf d `seq` rnf e
+  rnf (SnapMergingRun a b c d e f) =
+      rnf a `seq` rnf b `seq` rnf c `seq` rnf d `seq` rnf e `seq` rnf f
   rnf (SnapSingleRun a) = rnf a
 
 -- | The total number of unspent credits. This total is used in combination with
@@ -156,12 +156,12 @@ newtype UnspentCredits = UnspentCredits { getUnspentCredits :: Int }
 
 data SnapMergingRunState r =
     SnapCompletedMerge !r
-  | SnapOngoingMerge !(V.Vector r) !SpentCredits !MergeType
+  | SnapOngoingMerge !(V.Vector r) !MergeType
   deriving stock (Show, Eq, Functor, Foldable, Traversable)
 
 instance NFData r => NFData (SnapMergingRunState r) where
-  rnf (SnapCompletedMerge a)   = rnf a
-  rnf (SnapOngoingMerge a b c) = rnf a `seq` rnf b `seq` rnf c
+  rnf (SnapCompletedMerge a) = rnf a
+  rnf (SnapOngoingMerge a b) = rnf a `seq` rnf b
 
 -- | The total number of spent credits. This total is used in combination with
 -- 'UnspentCedits' on snapshot load to restore merging work that was lost when
@@ -196,33 +196,28 @@ toSnapIncomingRun ::
   => IncomingRun m h
   -> m (SnapIncomingRun (Ref (Run m h)))
 toSnapIncomingRun (Single r) = pure (SnapSingleRun r)
--- We need to know how many credits were yet unspent so we can restore merge
--- work on snapshot load. No need to snapshot the contents of totalStepsVar
--- here, since we still start counting from 0 again when loading the snapshot.
 toSnapIncomingRun (Merging mergePolicy (DeRef MR.MergingRun {..})) = do
+    -- We need to know how many credits were spend and yet unspent so we can
+    -- restore merge work on snapshot load. No need to snapshot the contents
+    -- of totalStepsVar here, since we still start counting from 0 again when
+    -- loading the snapshot.
     unspentCredits <- readPrimVar (MR.getUnspentCreditsVar mergeUnspentCredits)
-    smrs <- withMVar mergeState $ \mrs -> toSnapMergingRunState mrs
+    spentCredits   <- readPrimVar (MR.getSpentCreditsVar mergeSpentCredits)
+    smrs <- toSnapMergingRunState <$> readMVar mergeState
     pure $
       SnapMergingRun
         mergePolicy
         mergeNumRuns
         mergeNumEntries
         (UnspentCredits unspentCredits)
+        (SpentCredits spentCredits)
         smrs
 
-{-# SPECIALISE toSnapMergingRunState ::
-     MR.MergingRunState IO h
-  -> IO (SnapMergingRunState (Ref (Run IO h))) #-}
 toSnapMergingRunState ::
-     PrimMonad m
-  => MR.MergingRunState m h
-  -> m (SnapMergingRunState (Ref (Run m h)))
-toSnapMergingRunState (MR.CompletedMerge r) = pure (SnapCompletedMerge r)
--- We need to know how many credits were spent already so we can restore merge
--- work on snapshot load.
-toSnapMergingRunState (MR.OngoingMerge rs (MR.SpentCreditsVar spentCreditsVar) m) = do
-    spentCredits <- readPrimVar spentCreditsVar
-    pure (SnapOngoingMerge rs (SpentCredits spentCredits) (Merge.mergeType m))
+     MR.MergingRunState m h
+  -> SnapMergingRunState (Ref (Run m h))
+toSnapMergingRunState (MR.CompletedMerge r)  = SnapCompletedMerge r
+toSnapMergingRunState (MR.OngoingMerge rs m) = SnapOngoingMerge rs (Merge.mergeType m)
 
 {-------------------------------------------------------------------------------
   Write Buffer
@@ -454,12 +449,13 @@ fromSnapLevels reg hfs hbio conf@TableConfig{..} uc resolve dir (SnapLevels leve
           -> m (IncomingRun m h)
         fromSnapIncomingRun (SnapSingleRun run) = do
             Single <$> dupRun run
-        fromSnapIncomingRun (SnapMergingRun mpfl nr ne unspentCredits smrs) = do
+        fromSnapIncomingRun (SnapMergingRun mpfl nr ne unspentCredits
+                                            spentCredits smrs) = do
             Merging mpfl <$> case smrs of
               SnapCompletedMerge run ->
                 withRollback reg (MR.newCompleted nr ne run) releaseRef
 
-              SnapOngoingMerge runs spentCredits mt -> do
+              SnapOngoingMerge runs mt -> do
                 rn <- uniqueToRunNumber <$> incrUniqCounter uc
                 mr <- withRollback reg
                   (MR.new hfs hbio resolve caching alloc mt (mkPath rn) runs)

src/Database/LSMTree/Internal/Snapshot/Codec.hs

@@ -459,13 +459,14 @@ instance DecodeVersioned RunNumber where
 -- SnapIncomingRun
 
 instance Encode (SnapIncomingRun RunNumber) where
-  encode (SnapMergingRun mpfl nr ne uc smrs) =
-       encodeListLen 6
+  encode (SnapMergingRun mpfl nr ne uc sc smrs) =
+       encodeListLen 7
     <> encodeWord 0
     <> encode mpfl
     <> encode nr
     <> encode ne
     <> encode uc
+    <> encode sc
     <> encode smrs
   encode (SnapSingleRun x) =
        encodeListLen 2
@@ -477,9 +478,9 @@ instance DecodeVersioned (SnapIncomingRun RunNumber) where
       n <- decodeListLen
       tag <- decodeWord
       case (n, tag) of
-        (6, 0) -> SnapMergingRun <$>
+        (7, 0) -> SnapMergingRun <$>
           decodeVersioned v <*> decodeVersioned v <*> decodeVersioned v <*>
-          decodeVersioned v <*> decodeVersioned v
+          decodeVersioned v <*> decodeVersioned v <*> decodeVersioned v
         (2, 1) -> SnapSingleRun <$> decodeVersioned v
         _ -> fail ("[SnapMergingRun] Unexpected combination of list length and tag: " <> show (n, tag))
 
@@ -520,11 +521,10 @@ instance Encode (SnapMergingRunState RunNumber) where
          encodeListLen 2
       <> encodeWord 0
       <> encode x
-  encode (SnapOngoingMerge rs tc l) =
-         encodeListLen 4
+  encode (SnapOngoingMerge rs l) =
+         encodeListLen 3
       <> encodeWord 1
       <> encode rs
-      <> encode tc
       <> encode l
 
 instance DecodeVersioned (SnapMergingRunState RunNumber) where
@@ -533,8 +533,8 @@ instance DecodeVersioned (SnapMergingRunState RunNumber) where
       tag <- decodeWord
       case (n, tag) of
         (2, 0) -> SnapCompletedMerge <$> decodeVersioned v
-        (4, 1) -> SnapOngoingMerge <$>
-          decodeVersioned v <*> decodeVersioned v <*> decodeVersioned v
+        (3, 1) -> SnapOngoingMerge <$> decodeVersioned v
+                                   <*> decodeVersioned v
         _ -> fail ("[SnapMergingRunState] Unexpected combination of list length and tag: " <> show (n, tag))
 
 -- SpentCredits

test/Test/Database/LSMTree/Internal/Snapshot/Codec.hs

@@ -278,12 +278,12 @@ deriving newtype instance Arbitrary RunNumber
 instance Arbitrary (SnapIncomingRun RunNumber) where
   arbitrary = oneof [
         SnapMergingRun <$> arbitrary <*> arbitrary <*> arbitrary
-                       <*> arbitrary <*> arbitrary
+                       <*> arbitrary <*> arbitrary <*> arbitrary
       , SnapSingleRun <$> arbitrary
       ]
-  shrink (SnapMergingRun a b c d e) =
-      [ SnapMergingRun a' b' c' d' e'
-      | (a', b', c', d', e') <- shrink (a, b, c, d, e) ]
+  shrink (SnapMergingRun a b c d e f) =
+      [ SnapMergingRun a' b' c' d' e' f'
+      | (a', b', c', d', e', f') <- shrink (a, b, c, d, e, f) ]
   shrink (SnapSingleRun a)  = SnapSingleRun <$> shrink a
 
 deriving newtype instance Arbitrary NumRuns
@@ -297,11 +297,11 @@ deriving newtype instance Arbitrary UnspentCredits
 instance Arbitrary (SnapMergingRunState RunNumber) where
   arbitrary = oneof [
         SnapCompletedMerge <$> arbitrary
-      , SnapOngoingMerge <$> arbitrary <*> arbitrary <*> arbitrary
+      , SnapOngoingMerge <$> arbitrary <*> arbitrary
       ]
   shrink (SnapCompletedMerge x) = SnapCompletedMerge <$> shrink x
-  shrink (SnapOngoingMerge x y z)   =
-      [ SnapOngoingMerge x' y' z' | (x', y', z') <- shrink (x, y, z) ]
+  shrink (SnapOngoingMerge x y) =
+      [ SnapOngoingMerge x' y' | (x', y') <- shrink (x, y) ]
 
 deriving newtype instance Arbitrary SpentCredits