Move impl of retrieveBlobs to new BlobRef.readWeakBlobRefs

Instead of open-coding the block IO stuff in the top level Internal
module in retrieveBlobs, it is moved next to BlobRef.readWeakBlobRef, so
we have the singular and bulk versions next to each other.

Author: dcoutts

src/Database/LSMTree/Internal.hs

@@ -87,7 +87,6 @@ import           Data.Kind
 import           Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import           Data.Maybe (catMaybes)
-import qualified Data.Primitive.ByteArray as P
 import qualified Data.Set as Set
 import           Data.Typeable
 import qualified Data.Vector as V
@@ -104,7 +103,6 @@ import           Database.LSMTree.Internal.Paths (SessionRoot (..),
                      SnapshotName)
 import qualified Database.LSMTree.Internal.Paths as Paths
 import           Database.LSMTree.Internal.Range (Range (..))
-import qualified Database.LSMTree.Internal.RawBytes as RB
 import           Database.LSMTree.Internal.Run (Run)
 import qualified Database.LSMTree.Internal.Run as Run
 import           Database.LSMTree.Internal.RunReaders (OffsetKey (..))
@@ -836,35 +834,10 @@ retrieveBlobs ::
   -> m (V.Vector SerialisedBlob)
 retrieveBlobs sesh wrefs =
     withOpenSession sesh $ \seshEnv ->
-      handle (\(BlobRef.WeakBlobRefInvalid i) -> throwIO (ErrBlobRefInvalid i)) $
-      BlobRef.withWeakBlobRefs wrefs $ \refs -> do
-
-        -- Prepare the IOOps:
-        -- We use a single large memory buffer, with appropriate offsets within
-        -- the buffer.
-        let bufSize :: Int
-            !bufSize = V.sum (V.map BlobRef.blobRefSpanSize refs)
-
-            {-# INLINE bufOffs #-}
-            bufOffs :: V.Vector Int
-            bufOffs = V.scanl (+) 0 (V.map BlobRef.blobRefSpanSize refs)
-        buf <- P.newPinnedByteArray bufSize
-        let ioops = V.zipWith (BlobRef.readBlobIOOp buf) bufOffs refs
-            hbio  = sessionHasBlockIO seshEnv
-
-        -- Submit the IOOps all in one go:
-        _ <- FS.submitIO hbio ioops
-        -- We do not need to inspect the results because IO errors are
-        -- thrown as exceptions, and the result is just the read length
-        -- which is already known. Short reads can't happen here.
-
-        -- Construct the SerialisedBlobs results:
-        -- This is just the different offsets within the shared buffer.
-        ba <- P.unsafeFreezeByteArray buf
-        pure $! V.zipWith
-                  (\off len -> SerialisedBlob (RB.fromByteArray off len ba))
-                  bufOffs
-                  (V.map BlobRef.blobRefSpanSize refs)
+      let hbio = sessionHasBlockIO seshEnv in
+      handle (\(BlobRef.WeakBlobRefInvalid i) ->
+                throwIO (ErrBlobRefInvalid i)) $
+      BlobRef.readWeakBlobRefs hbio wrefs
 
 {-------------------------------------------------------------------------------
   Cursors

src/Database/LSMTree/Internal/BlobRef.hs

@@ -18,7 +18,7 @@ module Database.LSMTree.Internal.BlobRef (
   , removeReferences
   , readRawBlobRef
   , readWeakBlobRef
-  , readBlobIOOp
+  , readWeakBlobRefs
   ) where
 
 import           Control.DeepSeq (NFData (..))
@@ -27,15 +27,18 @@ import           Control.Monad.Class.MonadThrow (Exception, MonadMask,
                      MonadThrow (..), bracket, throwIO)
 import           Control.Monad.Primitive
 import qualified Control.RefCount as RC
-import qualified Data.Primitive.ByteArray as P (MutableByteArray)
+import qualified Data.Primitive.ByteArray as P (newPinnedByteArray,
+                     unsafeFreezeByteArray)
 import qualified Data.Vector as V
 import qualified Data.Vector.Mutable as VM
 import           Database.LSMTree.Internal.BlobFile (BlobFile (..), BlobSpan (..))
 import qualified Database.LSMTree.Internal.BlobFile as BlobFile
+import qualified Database.LSMTree.Internal.RawBytes as RB
 import           Database.LSMTree.Internal.Serialise (SerialisedBlob (..))
 import qualified System.FS.API as FS
 import           System.FS.API (HasFS)
 import qualified System.FS.BlockIO.API as FS
+import           System.FS.BlockIO.API (HasBlockIO)
 
 
 -- | A raw blob reference is a reference to a blob within a blob file.
@@ -199,17 +202,47 @@ readWeakBlobRef fs wref =
       \StrongBlobRef {strongBlobRefFile, strongBlobRefSpan} ->
         BlobFile.readBlobFile fs strongBlobRefFile strongBlobRefSpan
 
-readBlobIOOp ::
-     P.MutableByteArray s -> Int
-  -> StrongBlobRef m h
-  -> FS.IOOp s h
-readBlobIOOp buf bufoff
-             StrongBlobRef {
-               strongBlobRefFile = BlobFile {blobFileHandle},
-               strongBlobRefSpan = BlobSpan {blobSpanOffset, blobSpanSize}
-             } =
-    FS.IOOpRead
-      blobFileHandle
-      (fromIntegral blobSpanOffset :: FS.FileOffset)
-      buf (FS.BufferOffset bufoff)
-      (fromIntegral blobSpanSize :: FS.ByteCount)
+{-# SPECIALISE readWeakBlobRefs :: HasBlockIO IO h -> V.Vector (WeakBlobRef IO h) -> IO (V.Vector SerialisedBlob) #-}
+readWeakBlobRefs ::
+     (MonadMask m, PrimMonad m)
+  => HasBlockIO m h
+  -> V.Vector (WeakBlobRef m h)
+  -> m (V.Vector SerialisedBlob)
+readWeakBlobRefs hbio wrefs =
+    bracket (deRefWeakBlobRefs wrefs) (V.mapM_ removeReference) $ \refs -> do
+      -- Prepare the IOOps:
+      -- We use a single large memory buffer, with appropriate offsets within
+      -- the buffer.
+      let bufSize :: Int
+          !bufSize = V.sum (V.map blobRefSpanSize refs)
+
+          {-# INLINE bufOffs #-}
+          bufOffs :: V.Vector Int
+          bufOffs = V.scanl (+) 0 (V.map blobRefSpanSize refs)
+      buf <- P.newPinnedByteArray bufSize
+
+      -- Submit the IOOps all in one go:
+      _ <- FS.submitIO hbio $
+             V.zipWith
+               (\bufoff
+                 StrongBlobRef {
+                   strongBlobRefFile = BlobFile {blobFileHandle},
+                   strongBlobRefSpan = BlobSpan {blobSpanOffset, blobSpanSize}
+                 } ->
+                 FS.IOOpRead
+                   blobFileHandle
+                   (fromIntegral blobSpanOffset :: FS.FileOffset)
+                   buf (FS.BufferOffset bufoff)
+                   (fromIntegral blobSpanSize :: FS.ByteCount))
+               bufOffs refs
+      -- We do not need to inspect the results because IO errors are
+      -- thrown as exceptions, and the result is just the read length
+      -- which is already known. Short reads can't happen here.
+
+      -- Construct the SerialisedBlobs results:
+      -- This is just the different offsets within the shared buffer.
+      ba <- P.unsafeFreezeByteArray buf
+      pure $! V.zipWith
+                (\off len -> SerialisedBlob (RB.fromByteArray off len ba))
+                bufOffs
+                (V.map blobRefSpanSize refs)