Introduce proper distinction between {Raw,Weak,Strong}BlobRef

Previously a RawBlobRef served dual purpose for raw and strong. Keep
that distinction clear.

For the moment, all three have equivalent representations, but this is
likely to change in a later refactoring of the reference counting API.

Also split the functions for making blob refs from runs or the write
buffer into the two variants that we use: raw (internally) and weak
(externally).

Author: dcoutts

src-extras/Database/LSMTree/Extras/NoThunks.hs

@@ -417,17 +417,21 @@ deriving anyclass instance NoThunks RawOverflowPage
   BlobRef
 -------------------------------------------------------------------------------}
 
-deriving stock instance Generic (RawBlobRef m h)
-deriving anyclass instance (Typeable h, Typeable (PrimState m))
-                        => NoThunks (RawBlobRef m h)
-
 deriving stock instance Generic BlobSpan
 deriving anyclass instance NoThunks BlobSpan
 
 deriving stock instance Generic (BlobFile m h)
 deriving anyclass instance (Typeable h, Typeable (PrimState m))
                         => NoThunks (BlobFile m h)
 
+deriving stock instance Generic (RawBlobRef m h)
+deriving anyclass instance (Typeable h, Typeable (PrimState m))
+                        => NoThunks (RawBlobRef m h)
+
+deriving stock instance Generic (WeakBlobRef m h)
+deriving anyclass instance (Typeable h, Typeable (PrimState m))
+                        => NoThunks (WeakBlobRef m h)
+
 {-------------------------------------------------------------------------------
   Arena
 -------------------------------------------------------------------------------}

src/Database/LSMTree/Internal/BlobRef.hs

@@ -13,8 +13,10 @@ module Database.LSMTree.Internal.BlobRef (
   , deRefWeakBlobRef
   , deRefWeakBlobRefs
   , WeakBlobRefInvalid (..)
+  , rawToWeakBlobRef
   , removeReference
   , removeReferences
+  , readRawBlobRef
   , readBlob
   , readBlobIOOp
   ) where
@@ -25,9 +27,9 @@ import           Control.Monad.Class.MonadThrow (Exception, MonadMask,
                      MonadThrow (..), bracket, throwIO)
 import           Control.Monad.Primitive
 import qualified Control.RefCount as RC
-import           Data.Coerce (coerce)
 import qualified Data.Primitive.ByteArray as P (MutableByteArray)
 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           Database.LSMTree.Internal.Serialise (SerialisedBlob (..))
@@ -54,8 +56,8 @@ data RawBlobRef m h = RawBlobRef {
 instance NFData h => NFData (RawBlobRef m h) where
   rnf (RawBlobRef a b) = rnf a `seq` rnf b
 
-blobRefSpanSize :: RawBlobRef m h -> Int
-blobRefSpanSize = fromIntegral . blobSpanSize . rawBlobRefSpan
+blobRefSpanSize :: StrongBlobRef m h -> Int
+blobRefSpanSize = fromIntegral . blobSpanSize . strongBlobRefSpan
 
 -- | A \"weak\" reference to a blob within a blob file. These are the ones we
 -- can return in the public API and can outlive their parent table.
@@ -67,8 +69,33 @@ blobRefSpanSize = fromIntegral . blobSpanSize . rawBlobRefSpan
 --
 -- See 'Database.LSMTree.Common.BlobRef' for more info.
 --
-newtype WeakBlobRef m h = WeakBlobRef (RawBlobRef m h)
-  deriving newtype (Show, NFData)
+data WeakBlobRef m h = WeakBlobRef {
+      weakBlobRefFile :: {-# NOUNPACK #-} !(BlobFile m h)
+    , weakBlobRefSpan :: {-# UNPACK #-}   !BlobSpan
+    }
+  deriving stock (Show)
+
+-- | A \"strong\" reference to a blob within a blob file. The blob file remains
+-- open while the strong reference is live. Thus it is safe to do I\/O to
+-- retrieve the blob based on the reference. Strong references must be released
+-- using 'releaseBlobRef' when no longer in use (e.g. after completing I\/O).
+--
+data StrongBlobRef m h = StrongBlobRef {
+      strongBlobRefFile :: {-# NOUNPACK #-} !(BlobFile m h)
+    , strongBlobRefSpan :: {-# UNPACK #-}   !BlobSpan
+    }
+  deriving stock (Show)
+
+-- | Convert a 'RawBlobRef' to a 'WeakBlobRef'.
+rawToWeakBlobRef :: RawBlobRef m h -> WeakBlobRef m h
+rawToWeakBlobRef RawBlobRef {rawBlobRefFile, rawBlobRefSpan} =
+    -- This doesn't need to really do anything, becuase the raw version
+    -- does not maintain an independent ref count, and the weak one does
+    -- not either.
+    WeakBlobRef {
+      weakBlobRefFile = rawBlobRefFile,
+      weakBlobRefSpan = rawBlobRefSpan
+    }
 
 -- | The 'WeakBlobRef' now points to a blob that is no longer available.
 newtype WeakBlobRefInvalid = WeakBlobRefInvalid Int
@@ -77,7 +104,7 @@ newtype WeakBlobRefInvalid = WeakBlobRefInvalid Int
 
 {-# SPECIALISE withWeakBlobRef ::
      WeakBlobRef IO h
-  -> (RawBlobRef IO h -> IO a)
+  -> (StrongBlobRef IO h -> IO a)
   -> IO a #-}
 -- | 'WeakBlobRef's are weak references. They do not keep the blob file open.
 -- Dereference a 'WeakBlobRef' to a strong 'BlobRef' to allow I\/O using
@@ -89,82 +116,98 @@ newtype WeakBlobRefInvalid = WeakBlobRefInvalid Int
 withWeakBlobRef ::
      (MonadMask m, PrimMonad m)
   => WeakBlobRef m h
-  -> (RawBlobRef m h -> m a)
+  -> (StrongBlobRef m h -> m a)
   -> m a
 withWeakBlobRef wref = bracket (deRefWeakBlobRef wref) removeReference
 
 {-# SPECIALISE withWeakBlobRefs ::
      V.Vector (WeakBlobRef IO h)
-  -> (V.Vector (RawBlobRef IO h) -> IO a)
+  -> (V.Vector (StrongBlobRef IO h) -> IO a)
   -> IO a #-}
 -- | The same as 'withWeakBlobRef' but for many references in one go.
 --
 withWeakBlobRefs ::
      (MonadMask m, PrimMonad m)
   => V.Vector (WeakBlobRef m h)
-  -> (V.Vector (RawBlobRef m h) -> m a)
+  -> (V.Vector (StrongBlobRef m h) -> m a)
   -> m a
 withWeakBlobRefs wrefs = bracket (deRefWeakBlobRefs wrefs) removeReferences
 
 {-# SPECIALISE deRefWeakBlobRef ::
      WeakBlobRef IO h
-  -> IO (RawBlobRef IO h) #-}
+  -> IO (StrongBlobRef IO h) #-}
 deRefWeakBlobRef ::
      (MonadThrow m, PrimMonad m)
   => WeakBlobRef m h
-  -> m (RawBlobRef m h)
-deRefWeakBlobRef (WeakBlobRef ref) = do
-    ok <- RC.upgradeWeakReference (blobFileRefCounter (rawBlobRefFile ref))
+  -> m (StrongBlobRef m h)
+deRefWeakBlobRef WeakBlobRef{weakBlobRefFile, weakBlobRefSpan} = do
+    ok <- RC.upgradeWeakReference (blobFileRefCounter weakBlobRefFile)
     when (not ok) $ throwIO (WeakBlobRefInvalid 0)
-    pure ref
+    return StrongBlobRef{
+      strongBlobRefFile = weakBlobRefFile,
+      strongBlobRefSpan = weakBlobRefSpan
+    }
 
 {-# SPECIALISE deRefWeakBlobRefs ::
      V.Vector (WeakBlobRef IO h)
-  -> IO (V.Vector (RawBlobRef IO h)) #-}
+  -> IO (V.Vector (StrongBlobRef IO h)) #-}
 deRefWeakBlobRefs ::
     forall m h.
      (MonadMask m, PrimMonad m)
   => V.Vector (WeakBlobRef m h)
-  -> m (V.Vector (RawBlobRef m h))
+  -> m (V.Vector (StrongBlobRef m h))
 deRefWeakBlobRefs wrefs = do
-    let refs :: V.Vector (RawBlobRef m h)
-        refs = coerce wrefs -- safely coerce away the newtype wrappers
-    V.iforM_ wrefs $ \i (WeakBlobRef ref) -> do
-      ok <- RC.upgradeWeakReference (blobFileRefCounter (rawBlobRefFile ref))
-      when (not ok) $ do
-        -- drop refs on the previous ones taken successfully so far
-        V.mapM_ removeReference (V.take i refs)
-        throwIO (WeakBlobRefInvalid i)
-    pure refs
-
-{-# SPECIALISE removeReference :: RawBlobRef IO h -> IO () #-}
-removeReference :: (MonadMask m, PrimMonad m) => RawBlobRef m h -> m ()
-removeReference = BlobFile.removeReference . rawBlobRefFile
-
-{-# SPECIALISE removeReferences :: V.Vector (RawBlobRef IO h) -> IO () #-}
-removeReferences :: (MonadMask m, PrimMonad m) => V.Vector (RawBlobRef m h) -> m ()
+    refs <- VM.new (V.length wrefs)
+    V.iforM_ wrefs $ \i WeakBlobRef {weakBlobRefFile, weakBlobRefSpan} -> do
+      ok <- RC.upgradeWeakReference (blobFileRefCounter weakBlobRefFile)
+      if ok
+        then VM.write refs i StrongBlobRef {
+               strongBlobRefFile = weakBlobRefFile,
+               strongBlobRefSpan = weakBlobRefSpan
+             }
+        else do
+          -- drop refs on the previous ones taken successfully so far
+          VM.mapM_ removeReference (VM.take i refs)
+          throwIO (WeakBlobRefInvalid i)
+    V.unsafeFreeze refs
+
+{-# SPECIALISE removeReference :: StrongBlobRef IO h -> IO () #-}
+removeReference :: (MonadMask m, PrimMonad m) => StrongBlobRef m h -> m ()
+removeReference = BlobFile.removeReference . strongBlobRefFile
+
+{-# SPECIALISE removeReferences :: V.Vector (StrongBlobRef IO h) -> IO () #-}
+removeReferences :: (MonadMask m, PrimMonad m) => V.Vector (StrongBlobRef m h) -> m ()
 removeReferences = V.mapM_ removeReference
 
+{-# INLINE readRawBlobRef #-}
+readRawBlobRef ::
+     (MonadThrow m, PrimMonad m)
+  => HasFS m h
+  -> RawBlobRef m h
+  -> m SerialisedBlob
+readRawBlobRef fs RawBlobRef {rawBlobRefFile, rawBlobRefSpan} =
+    BlobFile.readBlobFile fs rawBlobRefFile rawBlobRefSpan
+
 {-# SPECIALISE readBlob ::
      HasFS IO h
-  -> RawBlobRef IO h
+  -> StrongBlobRef IO h
   -> IO SerialisedBlob #-}
 readBlob ::
      (MonadThrow m, PrimMonad m)
   => HasFS m h
-  -> RawBlobRef m h
+  -> StrongBlobRef m h
   -> m SerialisedBlob
-readBlob fs RawBlobRef {rawBlobRefFile, rawBlobRefSpan} =
-    BlobFile.readBlobFile fs rawBlobRefFile rawBlobRefSpan
+readBlob fs StrongBlobRef {strongBlobRefFile, strongBlobRefSpan} =
+    BlobFile.readBlobFile fs strongBlobRefFile strongBlobRefSpan
 
 readBlobIOOp ::
      P.MutableByteArray s -> Int
-  -> RawBlobRef m h
+  -> StrongBlobRef m h
   -> FS.IOOp s h
 readBlobIOOp buf bufoff
-             RawBlobRef {
-               rawBlobRefFile = BlobFile {blobFileHandle},
-               rawBlobRefSpan = BlobSpan {blobSpanOffset, blobSpanSize}
+             StrongBlobRef {
+               strongBlobRefFile = BlobFile {blobFileHandle},
+               strongBlobRefSpan = BlobSpan {blobSpanOffset, blobSpanSize}
              } =
     FS.IOOpRead
       blobFileHandle

src/Database/LSMTree/Internal/Cursor.hs

@@ -11,6 +11,7 @@ import           Control.Monad.Fix (MonadFix)
 import qualified Data.Vector as V
 import           Database.LSMTree.Internal.BlobRef (RawBlobRef,
                      WeakBlobRef (..))
+import qualified Database.LSMTree.Internal.BlobRef as BlobRef
 import           Database.LSMTree.Internal.Entry (Entry)
 import qualified Database.LSMTree.Internal.Entry as Entry
 import           Database.LSMTree.Internal.Lookup (ResolveSerialisedValue)
@@ -126,6 +127,6 @@ readEntriesWhile resolve keyIsWanted fromEntry readers n =
              -> Maybe res
     toResult key = \case
         Entry.Insert v -> Just $ fromEntry key v Nothing
-        Entry.InsertWithBlob v b -> Just $ fromEntry key v (Just (WeakBlobRef b))
+        Entry.InsertWithBlob v b -> Just $ fromEntry key v (Just (BlobRef.rawToWeakBlobRef b))
         Entry.Mupdate v -> Just $ fromEntry key v Nothing
         Entry.Delete -> Nothing

src/Database/LSMTree/Internal/Lookup.hs

@@ -47,7 +47,8 @@ import           Database.LSMTree.Internal.Page (PageSpan (..), getNumPages,
 import           Database.LSMTree.Internal.RawBytes (RawBytes (..))
 import qualified Database.LSMTree.Internal.RawBytes as RB
 import           Database.LSMTree.Internal.RawPage
-import           Database.LSMTree.Internal.Run (Run, mkBlobRefForRun)
+import           Database.LSMTree.Internal.Run (Run)
+import qualified Database.LSMTree.Internal.Run as Run
 import           Database.LSMTree.Internal.Serialise
 import qualified Database.LSMTree.Internal.Vector as V
 import qualified Database.LSMTree.Internal.WriteBuffer as WB
@@ -245,8 +246,7 @@ intraPageLookups !resolveV !wb !wbblobs !rs !ks !rkixs !ioops !ioress = do
     res <- VM.generateM (V.length ks) $ \ki ->
              case WB.lookup wb (V.unsafeIndex ks ki) of
                Nothing -> pure Nothing
-               Just e  -> pure $! Just $!
-                            fmap (WeakBlobRef . WBB.mkBlobRef wbblobs) e
+               Just e  -> pure $! Just $! fmap (WBB.mkWeakBlobRef wbblobs) e
                 -- TODO:  ^^ we should be able to avoid this allocation by
                 -- combining the conversion with other later conversions.
     loop res 0
@@ -275,8 +275,7 @@ intraPageLookups !resolveV !wb !wbblobs !rs !ks !rkixs !ioops !ioress = do
             -- Laziness ensures that we only compute the forcing of the value in
             -- the entry when the result is needed.
             LookupEntry e         -> do
-                let e' = bimap copySerialisedValue
-                               (WeakBlobRef . mkBlobRefForRun r) e
+                let e' = bimap copySerialisedValue (Run.mkWeakBlobRef r) e
                 -- TODO: ^^ we should be able to avoid this allocation by
                 -- combining the conversion with other later conversions.
                 V.unsafeInsertWithMStrict res (combine resolveV) kix e'
@@ -289,7 +288,7 @@ intraPageLookups !resolveV !wb !wbblobs !rs !ks !rkixs !ioops !ioress = do
                                   (unBufferOffset (ioopBufferOffset ioop) + 4096)
                                   (fromIntegral m)
                                   buf)
-                    e' = bimap v' (WeakBlobRef . mkBlobRefForRun r) e
+                    e' = bimap v' (Run.mkWeakBlobRef r) e
                 V.unsafeInsertWithMStrict res (combine resolveV) kix e'
           loop res (ioopix + 1)
 

src/Database/LSMTree/Internal/Run.hs

@@ -45,7 +45,8 @@ module Database.LSMTree.Internal.Run (
   , addReference
   , removeReference
   , removeReferenceN
-  , mkBlobRefForRun
+  , mkRawBlobRef
+  , mkWeakBlobRef
     -- ** Run creation
   , fromMutable
   , fromWriteBuffer
@@ -68,7 +69,7 @@ import           Data.Foldable (for_)
 import           Data.Word (Word64)
 import           Database.LSMTree.Internal.BlobFile hiding (removeReference)
 import qualified Database.LSMTree.Internal.BlobFile as BlobFile
-import           Database.LSMTree.Internal.BlobRef hiding (removeReference)
+import           Database.LSMTree.Internal.BlobRef  (RawBlobRef (..), WeakBlobRef (..))
 import           Database.LSMTree.Internal.BloomFilter (bloomFilterFromSBS)
 import qualified Database.LSMTree.Internal.CRC32C as CRC
 import           Database.LSMTree.Internal.Entry (NumEntries (..))
@@ -143,14 +144,22 @@ removeReference r = RC.removeReference (runRefCounter r)
 removeReferenceN :: (PrimMonad m, MonadMask m) => Run m h -> Word64 -> m ()
 removeReferenceN r = RC.removeReferenceN (runRefCounter r)
 
--- | Helper function to make a 'BlobRef' that points into a 'Run'.
-mkBlobRefForRun :: Run m h -> BlobSpan -> RawBlobRef m h
-mkBlobRefForRun Run{runBlobFile} blobspan =
+-- | Helper function to make a 'WeakBlobRef' that points into a 'Run'.
+mkRawBlobRef :: Run m h -> BlobSpan -> RawBlobRef m h
+mkRawBlobRef Run{runBlobFile} blobspan =
     RawBlobRef {
       rawBlobRefFile = runBlobFile,
       rawBlobRefSpan = blobspan
     }
 
+-- | Helper function to make a 'WeakBlobRef' that points into a 'Run'.
+mkWeakBlobRef :: Run m h -> BlobSpan -> WeakBlobRef m h
+mkWeakBlobRef Run{runBlobFile} blobspan =
+    WeakBlobRef {
+      weakBlobRefFile = runBlobFile,
+      weakBlobRefSpan = blobspan
+    }
+
 {-# SPECIALISE close ::
      Run IO h
   -> IO () #-}
@@ -250,7 +259,7 @@ fromWriteBuffer ::
 fromWriteBuffer fs hbio caching alloc fsPaths buffer blobs = do
     builder <- Builder.new fs hbio fsPaths (WB.numEntries buffer) alloc
     for_ (WB.toList buffer) $ \(k, e) ->
-      Builder.addKeyOp builder k (fmap (WBB.mkBlobRef blobs) e)
+      Builder.addKeyOp builder k (fmap (WBB.mkRawBlobRef blobs) e)
       --TODO: the fmap entry here reallocates even when there are no blobs
     fromMutable caching (RefCount 1) builder
 

src/Database/LSMTree/Internal/RunBuilder.hs

@@ -283,7 +283,7 @@ copyBlob ::
   -> RawBlobRef m h
   -> m BlobSpan
 copyBlob builder@RunBuilder {..} blobref = do
-    blob <- BlobRef.readBlob runBuilderHasFS blobref
+    blob <- BlobRef.readRawBlobRef runBuilderHasFS blobref
     writeBlob builder blob
 
 {-# SPECIALISE writeFilter ::

src/Database/LSMTree/Internal/RunReader.hs

@@ -246,14 +246,14 @@ next reader@RunReader {..} = do
                 go 0 p  -- try again on the new page
           IndexEntry key entry -> do
             modifyPrimVar readerCurrentEntryNo (+1)
-            let entry' = fmap (Run.mkBlobRefForRun readerRun) entry
+            let entry' = fmap (Run.mkRawBlobRef readerRun) entry
             let rawEntry = Entry entry'
             return (ReadEntry key rawEntry)
           IndexEntryOverflow key entry lenSuffix -> do
             -- TODO: we know that we need the next page, could already load?
             modifyPrimVar readerCurrentEntryNo (+1)
             let entry' :: E.Entry SerialisedValue (RawBlobRef m h)
-                entry' = fmap (Run.mkBlobRefForRun readerRun) entry
+                entry' = fmap (Run.mkRawBlobRef readerRun) entry
             overflowPages <- readOverflowPages readerHasFS readerKOpsHandle lenSuffix
             let rawEntry = mkEntryOverflow entry' page lenSuffix overflowPages
             return (ReadEntry key rawEntry)

src/Database/LSMTree/Internal/RunReaders.hs

@@ -132,7 +132,7 @@ new !offsetKey wbs runs = do
            -> m (Maybe (ReadCtx m h))
     fromWB wb wbblobs = do
         --TODO: this BlobSpan to BlobRef conversion involves quite a lot of allocation
-        kops <- newMutVar $ map (fmap (fmap (WB.mkBlobRef wbblobs))) $
+        kops <- newMutVar $ map (fmap (fmap (WB.mkRawBlobRef wbblobs))) $
                   Map.toList $ filterWB $ WB.toMap wb
         nextReadCtx (ReaderNumber 0) (ReadBuffer kops)
       where