From 606c3534f4d1efa715caf428996345917afb4bc6 Mon Sep 17 00:00:00 2001
From: Joris Dral <joris@well-typed.com>
Date: Mon, 17 Mar 2025 13:22:42 +0100
Subject: [PATCH 1/2] Don't pass an `ActionRegistry` to `duplicateIncomingRun`

---
 src/Database/LSMTree/Internal/IncomingRun.hs   | 14 ++++++--------
 src/Database/LSMTree/Internal/MergeSchedule.hs |  2 +-
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/src/Database/LSMTree/Internal/IncomingRun.hs b/src/Database/LSMTree/Internal/IncomingRun.hs
index a21357592..141bc78fc 100644
--- a/src/Database/LSMTree/Internal/IncomingRun.hs
+++ b/src/Database/LSMTree/Internal/IncomingRun.hs
@@ -25,7 +25,6 @@ module Database.LSMTree.Internal.IncomingRun (
   , immediatelyCompleteIncomingRun
   ) where
 
-import           Control.ActionRegistry
 import           Control.Concurrent.Class.MonadMVar.Strict
 import           Control.DeepSeq (NFData (..))
 import           Control.Monad.Class.MonadST (MonadST)
@@ -64,18 +63,17 @@ instance NFData MergePolicyForLevel where
   rnf LevelTiering   = ()
   rnf LevelLevelling = ()
 
-{-# SPECIALISE duplicateIncomingRun :: ActionRegistry IO -> IncomingRun IO h -> IO (IncomingRun IO h) #-}
+{-# SPECIALISE duplicateIncomingRun :: IncomingRun IO h -> IO (IncomingRun IO h) #-}
 duplicateIncomingRun ::
      (PrimMonad m, MonadMask m)
-  => ActionRegistry m
-  -> IncomingRun m h
+  => IncomingRun m h
   -> m (IncomingRun m h)
-duplicateIncomingRun reg (Single r) =
-    Single <$> withRollback reg (dupRef r) releaseRef
+duplicateIncomingRun (Single r) =
+    Single <$> dupRef r
 
-duplicateIncomingRun reg (Merging mp md mcv mr) =
+duplicateIncomingRun (Merging mp md mcv mr) =
     Merging mp md <$> (newPrimVar =<< readPrimVar mcv)
-                  <*> withRollback reg (dupRef mr) releaseRef
+                  <*> dupRef mr
 
 {-# SPECIALISE releaseIncomingRun :: IncomingRun IO h -> IO () #-}
 releaseIncomingRun ::
diff --git a/src/Database/LSMTree/Internal/MergeSchedule.hs b/src/Database/LSMTree/Internal/MergeSchedule.hs
index b69f03cd2..9d952bef3 100644
--- a/src/Database/LSMTree/Internal/MergeSchedule.hs
+++ b/src/Database/LSMTree/Internal/MergeSchedule.hs
@@ -326,7 +326,7 @@ duplicateLevels ::
   -> m (Levels m h)
 duplicateLevels reg levels =
     forMStrict levels $ \Level {incomingRun, residentRuns} -> do
-      incomingRun'  <- duplicateIncomingRun reg incomingRun
+      incomingRun'  <- withRollback reg (duplicateIncomingRun incomingRun) releaseIncomingRun
       residentRuns' <- forMStrict residentRuns $ \r ->
                          withRollback reg (dupRef r) releaseRef
       return $! Level {

From e4c2a0df71a28583d86a9b46073e71167f22f3c7 Mon Sep 17 00:00:00 2001
From: Joris Dral <joris@well-typed.com>
Date: Mon, 17 Mar 2025 13:32:11 +0100
Subject: [PATCH 2/2] Don't re-export `IncomingRun` from `MergeSchedule`

---
 src-extras/Database/LSMTree/Extras/NoThunks.hs | 1 +
 src/Database/LSMTree/Internal.hs               | 1 +
 src/Database/LSMTree/Internal/MergeSchedule.hs | 7 +------
 src/Database/LSMTree/Internal/Snapshot.hs      | 1 +
 4 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/src-extras/Database/LSMTree/Extras/NoThunks.hs b/src-extras/Database/LSMTree/Extras/NoThunks.hs
index e8cf06bac..d81eb7730 100644
--- a/src-extras/Database/LSMTree/Extras/NoThunks.hs
+++ b/src-extras/Database/LSMTree/Extras/NoThunks.hs
@@ -41,6 +41,7 @@ import           Database.LSMTree.Internal.Chunk
 import           Database.LSMTree.Internal.Config
 import           Database.LSMTree.Internal.CRC32C
 import           Database.LSMTree.Internal.Entry
+import           Database.LSMTree.Internal.IncomingRun
 import           Database.LSMTree.Internal.Index
 import           Database.LSMTree.Internal.Index.Compact
 import           Database.LSMTree.Internal.Index.CompactAcc
diff --git a/src/Database/LSMTree/Internal.hs b/src/Database/LSMTree/Internal.hs
index c92186c1d..a234e7fbe 100644
--- a/src/Database/LSMTree/Internal.hs
+++ b/src/Database/LSMTree/Internal.hs
@@ -105,6 +105,7 @@ import qualified Database.LSMTree.Internal.BlobRef as BlobRef
 import           Database.LSMTree.Internal.Config
 import qualified Database.LSMTree.Internal.Cursor as Cursor
 import           Database.LSMTree.Internal.Entry (Entry)
+import           Database.LSMTree.Internal.IncomingRun (IncomingRun (..))
 import           Database.LSMTree.Internal.Lookup (ByteCountDiscrepancy,
                      ResolveSerialisedValue, lookupsIO,
                      lookupsIOWithoutWriteBuffer)
diff --git a/src/Database/LSMTree/Internal/MergeSchedule.hs b/src/Database/LSMTree/Internal/MergeSchedule.hs
index 9d952bef3..c0d6da270 100644
--- a/src/Database/LSMTree/Internal/MergeSchedule.hs
+++ b/src/Database/LSMTree/Internal/MergeSchedule.hs
@@ -23,13 +23,7 @@ module Database.LSMTree.Internal.MergeSchedule (
     -- * Levels, runs and ongoing merges
   , Levels
   , Level (..)
-  , IncomingRun (..)
   , MergePolicyForLevel (..)
-  , newIncomingSingleRun
-  , newIncomingMergingRun
-  , releaseIncomingRun
-  , supplyCreditsIncomingRun
-  , snapshotIncomingRun
   , mergingRunParamsForLevel
     -- * Union level
   , UnionLevel (..)
@@ -710,6 +704,7 @@ addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root uc r0 reg levels ul
               TraceCompletedMerge (Run.size r) (Run.runFsPathsNumber r)
 
       return ir
+
 {-# SPECIALISE newIncomingRunAtLevel ::
      Tracer IO (AtLevel MergeTrace)
   -> HasFS IO h
diff --git a/src/Database/LSMTree/Internal/Snapshot.hs b/src/Database/LSMTree/Internal/Snapshot.hs
index 9347fa0fc..5604035d9 100644
--- a/src/Database/LSMTree/Internal/Snapshot.hs
+++ b/src/Database/LSMTree/Internal/Snapshot.hs
@@ -49,6 +49,7 @@ import qualified Data.Vector as V
 import           Database.LSMTree.Internal.Config
 import           Database.LSMTree.Internal.CRC32C (checkCRC)
 import qualified Database.LSMTree.Internal.CRC32C as CRC
+import           Database.LSMTree.Internal.IncomingRun
 import           Database.LSMTree.Internal.Lookup (ResolveSerialisedValue)
 import qualified Database.LSMTree.Internal.Merge as Merge
 import           Database.LSMTree.Internal.MergeSchedule