Pass BlobFile from WriteBufferBlobs to WriteBufferReader

src/Database/LSMTree/Internal/Snapshot.hs

@@ -40,7 +40,6 @@ import           Data.Primitive.PrimVar
 import           Data.Text (Text)
 import           Data.Traversable (for)
 import qualified Data.Vector as V
-import           Database.LSMTree.Internal.BlobFile (openBlobFile)
 import           Database.LSMTree.Internal.Config
 import           Database.LSMTree.Internal.Entry
 import           Database.LSMTree.Internal.Lookup (ResolveSerialisedValue)
@@ -303,14 +302,15 @@ openWriteBuffer reg resolve hfs hbio uc activeDir snapWriteBufferPaths = do
   let activeWriteBufferBlobPath =
         getActiveDir activeDir </> FS.mkFsPath [show activeWriteBufferNumber] <.> "wbblobs"
   copyFile reg hfs hbio (writeBufferBlobPath snapWriteBufferPaths) activeWriteBufferBlobPath
-  blobFile <-
+  writeBufferBlobs <-
     allocateTemp reg
-      (openBlobFile hfs activeWriteBufferBlobPath (FS.ReadWriteMode FS.AllowExisting))
+      (WBB.open hfs activeWriteBufferBlobPath FS.AllowExisting)
       releaseRef
   -- Read write buffer key/ops
   let kOpsPath = ForKOps (writeBufferKOpsPath snapWriteBufferPaths)
-  writeBuffer <- WBR.readWriteBuffer resolve hfs hbio kOpsPath =<< dupRef blobFile
-  writeBufferBlobs <- allocateTemp reg (WBB.fromBlobFile hfs blobFile) releaseRef
+  writeBuffer <-
+    withRef writeBufferBlobs $ \wbb ->
+      WBR.readWriteBuffer resolve hfs hbio kOpsPath (WBB.blobFile wbb)
   pure (writeBuffer, writeBufferBlobs)
 
 {-------------------------------------------------------------------------------

src/Database/LSMTree/Internal/WriteBufferReader.hs

@@ -9,7 +9,7 @@ import           Control.Monad.Class.MonadST (MonadST (..))
 import           Control.Monad.Class.MonadSTM (MonadSTM (..))
 import           Control.Monad.Class.MonadThrow (MonadMask, MonadThrow (..))
 import           Control.Monad.Primitive (PrimMonad (..))
-import           Control.RefCount (Ref, releaseRef)
+import           Control.RefCount (Ref, dupRef, releaseRef)
 import           Data.Primitive.MutVar (MutVar, newMutVar, readMutVar,
                      writeMutVar)
 import           Data.Primitive.PrimVar
@@ -95,10 +95,11 @@ new :: forall m h.
   -> ForKOps FS.FsPath
   -> Ref (BlobFile m h)
   -> m (WriteBufferReader m h)
-new readerHasFS readerHasBlockIO kOpsPath readerBlobFile = do
+new readerHasFS readerHasBlockIO kOpsPath blobFile = do
   readerKOpsHandle <- FS.hOpen readerHasFS (unForKOps kOpsPath) FS.ReadMode
   -- Double the file readahead window (only applies to this file descriptor)
   FS.hAdviseAll readerHasBlockIO readerKOpsHandle FS.AdviceSequential
+  readerBlobFile <- dupRef blobFile
   -- Load first page from disk, if it exists.
   readerCurrentEntryNo <- newPrimVar (0 :: Word16)
   firstPage <- readDiskPage readerHasFS readerKOpsHandle

test/Test/Database/LSMTree/Internal/Run.hs

@@ -237,9 +237,9 @@ prop_WriteAndOpenWriteBuffer hfs hbio rd = do
     withSerialisedWriteBuffer hfs hbio wbPaths wb wbb $ do
       -- Read write buffer from disk:
       let kOpsPath = Paths.ForKOps (Paths.writeBufferKOpsPath wbPaths)
-      let DeRef wbb' = wbb
-      wb' <- readWriteBuffer resolve hfs hbio kOpsPath =<< dupRef (WBB.blobFile wbb')
-      assertEqual "k/ops" wb wb'
+      withRef wbb $ \wbb' -> do
+        wb' <- readWriteBuffer resolve hfs hbio kOpsPath (WBB.blobFile wbb')
+        assertEqual "k/ops" wb wb'
   -- TODO: return a proper Property instead of using assertEqual etc.
   return (property True)