Use more consistent reference handling for newIncomingMergingRun

It should duplicate its input reference for consistency, since it
retains it.

Author: dcoutts

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -420,14 +420,14 @@ newIncomingSingleRun r = Single <$> dupRef r
 
 {-# INLINE newIncomingMergingRun #-}
 newIncomingMergingRun ::
-     PrimMonad m
+     (PrimMonad m, MonadThrow m)
   => MergePolicyForLevel
   -> NominalDebt
   -> Ref (MergingRun MR.LevelMergeType m h)
   -> m (IncomingRun m h)
 newIncomingMergingRun mergePolicy nominalDebt mr = do
     nominalCreditsVar <- newPrimVar (NominalCredits 0)
-    return (Merging mergePolicy nominalDebt nominalCreditsVar mr)
+    Merging mergePolicy nominalDebt nominalCreditsVar <$> dupRef mr
 
 {-# SPECIALISE supplyCreditsIncomingRun ::
      TableConfig
@@ -1029,12 +1029,12 @@ newIncomingRunAtLevel tr hfs hbio
       TraceNewMerge (V.map Run.size rs) (runNumber runPaths)
                     runParams mergePolicy mergeType
 
-    mr <- MR.new hfs hbio resolve runParams mergeType runPaths rs
-
-    assert (MR.totalMergeDebt mr <= maxMergeDebt conf mergePolicy ln) $ pure ()
-
-    let nominalDebt = nominalDebtForLevel conf ln
-    newIncomingMergingRun mergePolicy nominalDebt mr
+    bracket
+      (MR.new hfs hbio resolve runParams mergeType runPaths rs)
+      releaseRef $ \mr ->
+        assert (MR.totalMergeDebt mr <= maxMergeDebt conf mergePolicy ln) $
+        let nominalDebt = nominalDebtForLevel conf ln in
+        newIncomingMergingRun mergePolicy nominalDebt mr
 
 mergingRunParamsForLevel ::
      ActiveDir

src/Database/LSMTree/Internal/Snapshot.hs

@@ -39,7 +39,8 @@ import           Control.Concurrent.Class.MonadSTM (MonadSTM)
 import           Control.DeepSeq (NFData (..))
 import           Control.Monad (void)
 import           Control.Monad.Class.MonadST (MonadST)
-import           Control.Monad.Class.MonadThrow (MonadMask, bracketOnError)
+import           Control.Monad.Class.MonadThrow (MonadMask, bracket,
+                     bracketOnError)
 import           Control.Monad.Primitive (PrimMonad)
 import           Control.RefCount
 import           Data.Foldable (sequenceA_)
@@ -652,14 +653,17 @@ fromSnapLevels hfs hbio uc conf resolve reg dir (SnapLevels levels) =
         newIncomingSingleRun run
 
     fromSnapIncomingRun ln (SnapMergingRun mergePolicy nominalDebt
-                                           nominalCredits smrs) = do
-      mr <- fromSnapMergingRunState hfs hbio uc resolve dir smrs
-      ir <- newIncomingMergingRun mergePolicy nominalDebt mr
-      -- This will set the correct nominal credits, but it will not do any more
-      -- merging work because fromSnapMergingRunState already supplies all the
-      -- merging credits already.
-      supplyCreditsIncomingRun conf ln ir nominalCredits
-      return ir
+                                           nominalCredits smrs) =
+      bracket
+        (fromSnapMergingRunState hfs hbio uc resolve dir smrs)
+        releaseRef $ \mr -> do
+
+        ir <- newIncomingMergingRun mergePolicy nominalDebt mr
+        -- This will set the correct nominal credits, but it will not do any
+        -- more merging work because fromSnapMergingRunState already supplies
+        -- all the merging credits already.
+        supplyCreditsIncomingRun conf ln ir nominalCredits
+        return ir
 
 {-# SPECIALISE fromSnapMergingRunState ::
      MR.IsMergeType t