Merge pull request #417 from IntersectMBO/dcoutts/refcount-vs-reference

Convert reference counted resources to a new Ref API

Author: dcoutts

bench/macro/lsm-tree-bench-lookups.hs

@@ -7,7 +7,7 @@ import           Control.Monad
 import           Control.Monad.Class.MonadST
 import           Control.Monad.Primitive
 import           Control.Monad.ST.Strict (ST, runST)
-import           Control.RefCount (RefCount (..))
+import           Control.RefCount
 import           Data.Arena (ArenaManager, newArenaManager, withArena)
 import           Data.Bits ((.&.))
 import           Data.BloomFilter (Bloom)
@@ -162,9 +162,9 @@ benchmarks !caching = withFS $ \hfs hbio -> do
     putStrLn "<finished>"
 
     traceMarkerIO "Computing statistics for generated runs"
-    let numEntries = V.map Run.runNumEntries runs
-        numPages = V.map Run.sizeInPages runs
-        nhashes = V.map Bloom.hashesN blooms
+    let numEntries = V.map Run.size runs
+        numPages   = V.map Run.sizeInPages runs
+        nhashes    = V.map Bloom.hashesN blooms
         bitsPerEntry = V.zipWith
                          (\b (NumEntries n) -> fromIntegral (Bloom.length b) / fromIntegral n :: Double)
                          blooms
@@ -213,7 +213,7 @@ benchmarks !caching = withFS $ \hfs hbio -> do
 
     traceMarkerIO "Cleaning up"
     putStrLn "Cleaning up"
-    V.mapM_ Run.removeReference runs
+    V.mapM_ releaseRef runs
 
     traceMarkerIO "Computing statistics for prepLookups results"
     putStr "<Computing statistics for prepLookups>"
@@ -331,7 +331,7 @@ lookupsEnv ::
   -> FS.HasFS IO FS.HandleIO
   -> FS.HasBlockIO IO FS.HandleIO
   -> Run.RunDataCaching
-  -> IO ( V.Vector (Run IO FS.HandleIO)
+  -> IO ( V.Vector (Ref (Run IO FS.HandleIO))
         , V.Vector (Bloom SerialisedKey)
         , V.Vector IndexCompact
         , V.Vector (FS.Handle FS.HandleIO)
@@ -371,11 +371,10 @@ lookupsEnv runSizes keyRng0 hfs hbio caching = do
     putStr "DONE"
 
     -- return runs
-    runs <- V.fromList <$>
-              mapM (Run.fromMutable caching (RefCount 1)) rbs
-    let blooms = V.map Run.runFilter runs
-        indexes = V.map Run.runIndex runs
-        handles = V.map Run.runKOpsFile runs
+    runs <- V.fromList <$> mapM (Run.fromMutable caching) rbs
+    let blooms  = V.map (\(DeRef r) -> Run.runFilter   r) runs
+        indexes = V.map (\(DeRef r) -> Run.runIndex    r) runs
+        handles = V.map (\(DeRef r) -> Run.runKOpsFile r) runs
     pure $!! (runs, blooms, indexes, handles)
 
 genLookupBatch :: StdGen -> Int -> (V.Vector SerialisedKey, StdGen)
@@ -466,8 +465,8 @@ benchLookupsIO ::
   -> ArenaManager RealWorld
   -> ResolveSerialisedValue
   -> WB.WriteBuffer
-  -> WBB.WriteBufferBlobs IO h
-  -> V.Vector (Run IO h)
+  -> Ref (WBB.WriteBufferBlobs IO h)
+  -> V.Vector (Ref (Run IO h))
   -> V.Vector (Bloom SerialisedKey)
   -> V.Vector IndexCompact
   -> V.Vector (FS.Handle h)

bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs

@@ -5,6 +5,7 @@ module Bench.Database.LSMTree.Internal.Lookup (benchmarks) where
 import           Control.Exception (assert)
 import           Control.Monad
 import           Control.Monad.ST.Strict (stToIO)
+import           Control.RefCount
 import           Criterion.Main (Benchmark, bench, bgroup, env, envWithCleanup,
                      perRunEnv, perRunEnvWithCleanup, whnf, whnfAppIO)
 import           Data.Arena (ArenaManager, closeArena, newArena,
@@ -20,7 +21,7 @@ import           Database.LSMTree.Extras.Random (frequency,
                      sampleUniformWithReplacement, uniformWithoutReplacement)
 import           Database.LSMTree.Extras.UTxO
 import           Database.LSMTree.Internal.BlobRef (BlobSpan (..))
-import           Database.LSMTree.Internal.Entry (Entry (..), unNumEntries)
+import           Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))
 import           Database.LSMTree.Internal.Lookup (bloomQueries, indexSearches,
                      intraPageLookups, lookupsIO, prepLookups)
 import           Database.LSMTree.Internal.Page (getNumPages)
@@ -83,9 +84,9 @@ benchmarks = bgroup "Bench.Database.LSMTree.Internal.Lookup" [
 benchLookups :: Config -> Benchmark
 benchLookups conf@Config{name} =
     withEnv $ \ ~(_dir, arenaManager, hasFS, hasBlockIO, rs, ks) ->
-      env ( pure ( V.map Run.runFilter rs
-                 , V.map Run.runIndex rs
-                 , V.map Run.runKOpsFile rs
+      env ( pure ( V.map (\(DeRef r) -> Run.runFilter   r) rs
+                 , V.map (\(DeRef r) -> Run.runIndex    r) rs
+                 , V.map (\(DeRef r) -> Run.runKOpsFile r) rs
                  )
           ) $ \ ~(blooms, indexes, kopsFiles) ->
         bgroup name [
@@ -130,7 +131,7 @@ benchLookups conf@Config{name} =
                 )
                 (\(_, _, _, arena, wbblobs) -> do
                     closeArena arenaManager arena
-                    WBB.removeReference wbblobs)
+                    releaseRef wbblobs)
                 (\ ~(rkixs, ioops, ioress, _, wbblobs_unused) -> do
                   !_ <- intraPageLookups resolveV WB.empty wbblobs_unused
                                          rs ks rkixs ioops ioress
@@ -180,7 +181,7 @@ lookupsInBatchesEnv ::
         , ArenaManager RealWorld
         , FS.HasFS IO FS.HandleIO
         , FS.HasBlockIO IO FS.HandleIO
-        , V.Vector (Run IO FS.HandleIO)
+        , V.Vector (Ref (Run IO FS.HandleIO))
         , V.Vector SerialisedKey
         )
 lookupsInBatchesEnv Config {..} = do
@@ -194,7 +195,7 @@ lookupsInBatchesEnv Config {..} = do
         fsps = RunFsPaths (FS.mkFsPath []) (RunNumber 0)
     wbblobs <- WBB.new hasFS (FS.mkFsPath [])
     r <- Run.fromWriteBuffer hasFS hasBlockIO caching (RunAllocFixed 10) fsps wb wbblobs
-    let nentriesReal = unNumEntries $ Run.runNumEntries r
+    let NumEntries nentriesReal = Run.size r
     assert (nentriesReal == nentries) $ pure ()
     let npagesReal = Run.sizeInPages r
     assert (getNumPages npagesReal * 42 <= nentriesReal) $ pure ()
@@ -212,13 +213,13 @@ lookupsInBatchesCleanup ::
      , ArenaManager RealWorld
      , FS.HasFS IO FS.HandleIO
      , FS.HasBlockIO IO FS.HandleIO
-     , V.Vector (Run IO FS.HandleIO)
+     , V.Vector (Ref (Run IO FS.HandleIO))
      , V.Vector SerialisedKey
      )
   -> IO ()
 lookupsInBatchesCleanup (tmpDir, _arenaManager, _hasFS, hasBlockIO, rs, _) = do
     FS.close hasBlockIO
-    forM_ rs Run.removeReference
+    forM_ rs releaseRef
     removeDirectoryRecursive tmpDir
 
 -- | Generate keys to store and keys to lookup

bench/micro/Bench/Database/LSMTree/Internal/Merge.hs

@@ -1,6 +1,7 @@
 module Bench.Database.LSMTree.Internal.Merge (benchmarks) where
 
 import           Control.Monad (when, zipWithM)
+import           Control.RefCount
 import           Criterion.Main (Benchmark, bench, bgroup)
 import qualified Criterion.Main as Cr
 import           Data.Bifunctor (first)
@@ -206,7 +207,7 @@ benchMerge conf@Config{name} =
                 -- Make sure to immediately close resulting runs so we don't run
                 -- out of file handles. Ideally this would not be measured, but at
                 -- least it's pretty cheap.
-                Run.removeReference run
+                releaseRef run
         ]
   where
     withEnv =
@@ -228,7 +229,7 @@ merge ::
   -> Config
   -> Run.RunFsPaths
   -> InputRuns
-  -> IO (Run IO FS.HandleIO)
+  -> IO (Ref (Run IO FS.HandleIO))
 merge fs hbio Config {..} targetPaths runs = do
     let f = fromMaybe const mergeMappend
     m <- fromMaybe (error "empty inputs, no merge created") <$>
@@ -241,7 +242,7 @@ outputRunPaths = RunFsPaths (FS.mkFsPath []) (RunNumber 0)
 inputRunPaths :: [Run.RunFsPaths]
 inputRunPaths = RunFsPaths (FS.mkFsPath []) . RunNumber <$> [1..]
 
-type InputRuns = V.Vector (Run IO FS.HandleIO)
+type InputRuns = V.Vector (Ref (Run IO FS.HandleIO))
 
 type Mappend = SerialisedValue -> SerialisedValue -> SerialisedValue
 
@@ -331,7 +332,7 @@ mergeEnvCleanup ::
      )
   -> IO ()
 mergeEnvCleanup (tmpDir, _hasFS, hasBlockIO, runs) = do
-    traverse_ Run.removeReference runs
+    traverse_ releaseRef runs
     removeDirectoryRecursive tmpDir
     FS.close hasBlockIO
 

lsm-tree.cabal

@@ -476,6 +476,7 @@ benchmark lsm-tree-micro-bench
     , lsm-tree
     , lsm-tree:blockio-api
     , lsm-tree:bloomfilter
+    , lsm-tree:control
     , lsm-tree:extras
     , QuickCheck
     , random
@@ -864,6 +865,7 @@ test-suite control-test
     , io-classes
     , io-sim
     , lsm-tree:control
+    , primitive
     , QuickCheck
     , tasty
     , tasty-quickcheck