Add removeReferenceN

In some cases, we have to remove more than one reference at a time. For example,
when a merge is completed, then the reference count of each input run should be
reduced by value of the reference count of the output run. You can view this as
the references being "transferred" from the input runs to the output run.

Author: jorisdral

src-control/Control/RefCount.hs

@@ -10,6 +10,7 @@ module Control.RefCount (
   , mkRefCounter1
   , addReference
   , removeReference
+  , removeReferenceN
   , upgradeWeakReference
   , readRefCount
   ) where
@@ -21,7 +22,9 @@ import           Control.Monad.Class.MonadThrow
 import           Control.Monad.Primitive
 import           Data.Maybe
 import           Data.Primitive.PrimVar
+import           Data.Word
 import           GHC.Stack
+import           Text.Printf
 
 -- | A reference counter with an optional finaliser action. Once the reference
 -- count reaches @0@, the finaliser will be run.
@@ -89,6 +92,43 @@ removeReference RefCounter{countVar, finaliser} = mask_ $ do
     assertWithCallStack (prevCount > 0) $ pure ()
     when (prevCount == 1) $ sequence_ finaliser
 
+-- TODO: remove uses of this API. Eventually all references should be singular,
+-- and not use patterns where if A contains B then N references on A becomes N
+-- references on B. Instead this should be a single reference from A to B,
+-- irrespective of the number of references to A.
+{-# SPECIALISE removeReferenceN :: HasCallStack => RefCounter IO -> Word64 -> IO () #-}
+-- | Decrease the reference counter by @n@. @n@ must be a positive number.
+--
+-- The count must be known (from context) to be non-zero and at least as large
+-- as @n@. Typically this will be because the caller has @n@ references already
+-- (that they took out themselves or were given).
+removeReferenceN :: (HasCallStack, PrimMonad m, MonadMask m) => RefCounter m -> Word64 -> m ()
+removeReferenceN RefCounter{countVar, finaliser} n = mask_ $ do
+    -- n should be positive
+    assert (n > 0) $ pure ()
+    let !n' = fromIntegralChecked n
+    prevCount <- fetchSubInt countVar n'
+    -- the reference count must not already be 0, because then the finaliser
+    -- will have run already
+    assertWithCallStack (prevCount > 0) $ pure ()
+    -- the reference count can not go below zero
+    assertWithCallStack (prevCount >= n') $ pure ()
+    when (prevCount <= n') $ sequence_ finaliser
+
+-- TODO: remove when removeReferenceN is removed
+{-# INLINABLE fromIntegralChecked #-}
+-- | Like 'fromIntegral', but throws an error when @(x :: a) /= fromIntegral
+-- (fromIntegral x :: b)@.
+fromIntegralChecked :: (HasCallStack, Integral a, Integral b, Show a) => a -> b
+fromIntegralChecked x
+  | x'' == x
+  = x'
+  | otherwise
+  = error $ printf "fromIntegralChecked: conversion failed, %s /= %s" (show x) (show x'')
+  where
+    x' = fromIntegral x
+    x'' = fromIntegral x'
+
 -- | Try to turn a \"weak\" reference on something into a proper reference.
 -- This is by analogy with @deRefWeak :: Weak v -> IO (Maybe v)@, but for
 -- reference counts.
@@ -120,6 +160,7 @@ upgradeWeakReference RefCounter{countVar} = do
             then return True
             else casLoop prevCount'
 
+-- TODO: remove when removeRefenceN is removed
 {-# SPECIALISE readRefCount :: RefCounter IO -> IO RefCount #-}
 -- | Warning: reading the current reference count is inherently racy as there is
 -- no way to reliably act on the information. It can be useful for debugging.

src/Database/LSMTree/Internal/Merge.hs

@@ -9,6 +9,8 @@ module Database.LSMTree.Internal.Merge (
   , new
   , addReference
   , removeReference
+  , removeReferenceN
+  , readRefCount
   , complete
   , stepsToCompletion
   , stepsToCompletionCounted
@@ -24,12 +26,13 @@ import           Control.Monad.Class.MonadThrow (MonadCatch, MonadMask (..),
                      MonadThrow (..))
 import           Control.Monad.Fix (MonadFix)
 import           Control.Monad.Primitive (PrimMonad, PrimState, RealWorld)
-import           Control.RefCount (RefCounter)
+import           Control.RefCount (RefCount (..), RefCounter)
 import qualified Control.RefCount as RC
 import           Data.Coerce (coerce)
 import           Data.Primitive.MutVar
 import           Data.Traversable (for)
 import qualified Data.Vector as V
+import           Data.Word
 import           Database.LSMTree.Internal.BlobRef (BlobRef)
 import           Database.LSMTree.Internal.Entry
 import           Database.LSMTree.Internal.Run (Run, RunDataCaching)
@@ -130,6 +133,14 @@ addReference Merge{..} = RC.addReference mergeRefCounter
 removeReference :: (HasCallStack, PrimMonad m, MonadMask m) => Merge m h -> m ()
 removeReference Merge{..} = RC.removeReference mergeRefCounter
 
+{-# SPECIALISE removeReferenceN :: Merge IO h -> Word64 -> IO () #-}
+removeReferenceN :: (HasCallStack, PrimMonad m, MonadMask m) => Merge m h -> Word64 -> m ()
+removeReferenceN r = RC.removeReferenceN (mergeRefCounter r)
+
+{-# SPECIALISE readRefCount :: Merge IO h -> IO RefCount #-}
+readRefCount :: PrimMonad m => Merge m h -> m RefCount
+readRefCount Merge{..} = RC.readRefCount mergeRefCounter
+
 {-# SPECIALISE finaliser ::
      MutVar RealWorld MergeState
   -> RunBuilder IO h

src/Database/LSMTree/Internal/Run.hs

@@ -44,6 +44,7 @@ module Database.LSMTree.Internal.Run (
   , sizeInPages
   , addReference
   , removeReference
+  , removeReferenceN
   , mkBlobRefForRun
     -- ** Run creation
   , fromMutable
@@ -64,6 +65,7 @@ import qualified Control.RefCount as RC
 import           Data.BloomFilter (Bloom)
 import qualified Data.ByteString.Short as SBS
 import           Data.Foldable (for_)
+import           Data.Word (Word64)
 import           Database.LSMTree.Internal.BlobRef (BlobRef (..), BlobSpan (..))
 import           Database.LSMTree.Internal.BloomFilter (bloomFilterFromSBS)
 import qualified Database.LSMTree.Internal.CRC32C as CRC
@@ -134,6 +136,10 @@ addReference r = RC.addReference (runRefCounter r)
 removeReference :: (PrimMonad m, MonadMask m) => Run m h -> m ()
 removeReference r = RC.removeReference (runRefCounter r)
 
+{-# SPECIALISE removeReferenceN :: Run IO h -> Word64 -> IO () #-}
+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 -> BlobRef m (FS.Handle h)
 mkBlobRefForRun Run{runBlobFile, runRefCounter} blobRefSpan =

test-control/Test/Control/RefCount.hs

@@ -3,16 +3,17 @@
 module Test.Control.RefCount (tests) where
 
 import           Control.Concurrent.Class.MonadMVar
-import           Control.Exception (AssertionFailed (..))
+import           Control.Exception (AssertionFailed (..), try)
 import           Control.Monad
-import           Control.Monad.Class.MonadThrow
 import           Control.RefCount
+import           Data.Either (isRight)
 import           Test.Tasty (TestTree, testGroup)
 import           Test.Tasty.QuickCheck
 
 tests :: TestTree
 tests = testGroup "Control.RefCount" [
       testProperty "prop_refCount" prop_refCount
+    , testProperty "prop_removeReferenceN" prop_removeReferenceN
     ]
 
 prop_refCount :: Property
@@ -80,3 +81,25 @@ prop_refCount = once $ ioProperty $ do
 #else
     check = \case Left (AssertionFailed _) -> False; Right () -> True
 #endif
+
+prop_removeReferenceN :: Positive Int -> NonNegative Int -> Property
+prop_removeReferenceN (Positive n) (NonNegative m) = ioProperty $ do
+    obj <- newMVar False
+    ref <- unsafeMkRefCounterN (RefCount n) $ Just (void $ modifyMVar_ obj (\x -> pure (not x)) )
+
+    e1 <- try @AssertionFailed $ removeReferenceN ref (fromIntegral m)
+    n1 <- readRefCount ref -- 0
+    b1 <- readMVar obj -- True
+
+    pure $
+        counterexample "e1" (if n < m || m == 0
+                             then check e1
+                             else isRight e1) .&&.
+        counterexample "n1" (n1 == RefCount (n - m)) .&&.
+        counterexample "b1" (b1 == (isRight e1 && m >= n))
+  where
+#ifdef NO_IGNORE_ASSERTS
+    check = \case Left (AssertionFailed _) -> True; Right () -> False
+#else
+    check = \case Left (AssertionFailed _) -> False; Right () -> True
+#endif