From 8f1fa56f4d857ffd6769f18b31eb0ddee898981b Mon Sep 17 00:00:00 2001
From: Duncan Coutts <duncan@well-typed.com>
Date: Wed, 18 Sep 2024 11:11:00 +0100
Subject: [PATCH] FIXUP: ref count cursor blobs properly

---
 src/Database/LSMTree/Internal.hs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/Database/LSMTree/Internal.hs b/src/Database/LSMTree/Internal.hs
index 352c48660..174e56af5 100644
--- a/src/Database/LSMTree/Internal.hs
+++ b/src/Database/LSMTree/Internal.hs
@@ -825,6 +825,10 @@ data CursorEnv m h = CursorEnv {
     -- | The runs held open by the cursor. We must remove a reference when the
     -- cursor gets closed.
   , cursorRuns       :: !(V.Vector (Run m (Handle h)))
+
+    -- | The write buffer blobs, which like the runs, we have to keep open
+    -- untile the cursor is closed.
+  , cursorWBB        :: WBB.WriteBufferBlobs m h
   }
 
 {-# SPECIALISE withCursor :: OffsetKey -> TableHandle IO h -> (Cursor IO h -> IO a) -> IO a #-}
@@ -859,12 +863,13 @@ newCursor !offsetKey th = withOpenTable th $ \thEnv -> do
     -- 'sessionOpenTables'.
     withOpenSession cursorSession $ \_ -> do
       withTempRegistry $ \reg -> do
-        (writeBuffer, cursorRuns) <-
+        (writeBuffer, writeBufferBlobs, cursorRuns) <-
           allocTableContent reg (tableContent thEnv)
         cursorReaders <-
           allocateMaybeTemp reg
             (Readers.new hfs hbio offsetKey (Just writeBuffer) cursorRuns)
             (Readers.close hfs hbio)
+        let cursorWBB = writeBufferBlobs
         cursorState <- newMVar (CursorOpen CursorEnv {..})
         let !cursor = Cursor {cursorState, cursorTracer}
         -- Track cursor, but careful: If now an exception is raised, all
@@ -891,7 +896,7 @@ newCursor !offsetKey th = withOpenTable th $ \thEnv -> do
             allocateTemp reg
               (Run.addReference r)
               (\_ -> Run.removeReference r)
-          pure (wb, runs)
+          pure (wb, wbblobs, runs)
 
 {-# SPECIALISE closeCursor :: Cursor IO h -> IO () #-}
 -- | See 'Database.LSMTree.Normal.closeCursor'.
@@ -917,6 +922,7 @@ closeCursor Cursor {..} = do
 
         forM_ cursorReaders $ freeTemp reg . Readers.close hfs hbio
         V.forM_ cursorRuns $ freeTemp reg . Run.removeReference
+        freeTemp reg (WBB.removeReference cursorWBB)
         return CursorClosed
 
 {-# SPECIALISE readCursor :: ResolveSerialisedValue -> Int -> Cursor IO h -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef IO (Handle h)) -> res) -> IO (V.Vector res) #-}