Limit the scope of forgotten reference checks

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

@@ -133,7 +133,7 @@ benchSalt :: Bloom.Salt
 benchSalt = 4
 
 benchmarks :: Run.RunDataCaching -> IO ()
-benchmarks !caching = withFS $ \hfs hbio -> do
+benchmarks !caching = withFS $ \hfs hbio refCtx -> do
 #ifdef NO_IGNORE_ASSERTS
     putStrLn "WARNING: Benchmarking in debug mode."
     putStrLn "         To benchmark in release mode, pass:"
@@ -163,7 +163,7 @@ benchmarks !caching = withFS $ \hfs hbio -> do
     -- instead of sequentially.
     let keyRng0 = mkStdGen 17
 
-    (!runs, !blooms, !indexes, !handles) <- lookupsEnv runSizes keyRng0 hfs hbio caching
+    (!runs, !blooms, !indexes, !handles) <- lookupsEnv runSizes keyRng0 hfs hbio refCtx caching
     putStrLn "<finished>"
 
     traceMarkerIO "Computing statistics for generated runs"
@@ -210,7 +210,7 @@ benchmarks !caching = withFS $ \hfs hbio -> do
                 "Calculate batches of keys, and perform disk lookups for each batch. This is roughly doing the same as benchPrepLookups, but also performing the disk I/O and resolving values. Net time/allocation is the result of subtracting the cost of benchGenKeyBatches."
                 (\n -> do
                     let wb_unused = WB.empty
-                    bracket (WBB.new hfs (FS.mkFsPath ["wbblobs_unused"])) releaseRef $ \wbblobs_unused ->
+                    bracket (WBB.new hfs refCtx (FS.mkFsPath ["wbblobs_unused"])) releaseRef $ \wbblobs_unused ->
                       benchLookupsIO hbio arenaManager benchmarkResolveSerialisedValue
                                      wb_unused wbblobs_unused runs blooms indexes handles
                                      keyRng0 n)
@@ -308,13 +308,14 @@ totalNumEntriesSanityCheck l1 runSizes =
     sum [ 2^l1 * sizeFactor | (_, sizeFactor) <- runSizes ]
 
 withFS ::
-     (FS.HasFS IO FS.HandleIO -> FS.HasBlockIO IO FS.HandleIO -> IO a)
+     (FS.HasFS IO FS.HandleIO -> FS.HasBlockIO IO FS.HandleIO -> RefCtx -> IO a)
   -> IO a
 withFS action =
+    withRefCtx $ \refCtx ->
     FS.withIOHasBlockIO (FS.MountPoint "_bench_lookups") FS.defaultIOCtxParams $ \hfs hbio -> do
       exists <- FS.doesDirectoryExist hfs (FS.mkFsPath [""])
       unless exists $ error ("_bench_lookups directory does not exist")
-      action hfs hbio
+      action hfs hbio refCtx
 
 -- | Input environment for benchmarking lookup functions.
 --
@@ -336,13 +337,14 @@ lookupsEnv ::
   -> StdGen -- ^ Key RNG
   -> FS.HasFS IO FS.HandleIO
   -> FS.HasBlockIO IO FS.HandleIO
+  -> RefCtx
   -> Run.RunDataCaching
   -> IO ( V.Vector (Ref (Run IO FS.HandleIO))
         , V.Vector (Bloom SerialisedKey)
         , V.Vector Index
         , V.Vector (FS.Handle FS.HandleIO)
         )
-lookupsEnv runSizes keyRng0 hfs hbio caching = do
+lookupsEnv runSizes keyRng0 hfs hbio refCtx caching = do
     -- create the vector of initial keys
     (mvec :: VUM.MVector RealWorld UTxOKey) <- VUM.unsafeNew (totalNumEntries runSizes)
     !keyRng1 <- vectorOfUniforms mvec keyRng0
@@ -381,7 +383,7 @@ lookupsEnv runSizes keyRng0 hfs hbio caching = do
     putStr "DONE"
 
     -- return runs
-    runs <- V.fromList <$> mapM Run.fromBuilder rbs
+    runs <- V.fromList <$> mapM (Run.fromBuilder refCtx) 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

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

@@ -90,7 +90,7 @@ benchSalt = 4
 
 benchLookups :: Config -> Benchmark
 benchLookups conf@Config{name} =
-    withEnv $ \ ~(_dir, arenaManager, _hasFS, hasBlockIO, wbblobs, rs, ks) ->
+    withEnv $ \ ~(_dir, arenaManager, _hasFS, hasBlockIO, _refCtx, wbblobs, rs, ks) ->
       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
@@ -182,6 +182,7 @@ lookupsInBatchesEnv ::
         , ArenaManager RealWorld
         , FS.HasFS IO FS.HandleIO
         , FS.HasBlockIO IO FS.HandleIO
+        , RefCtx
         , Ref (WBB.WriteBufferBlobs IO FS.HandleIO)
         , V.Vector (Ref (Run IO FS.HandleIO))
         , V.Vector SerialisedKey
@@ -192,10 +193,11 @@ lookupsInBatchesEnv Config {..} = do
     benchTmpDir <- createTempDirectory sysTmpDir "lookupsInBatchesEnv"
     (storedKeys, lookupKeys) <- lookupsEnv (mkStdGen 17) nentries npos nneg
     (hasFS, hasBlockIO) <- FS.ioHasBlockIO (FS.MountPoint benchTmpDir) (fromMaybe FS.defaultIOCtxParams ioctxps)
-    wbblobs <- WBB.new hasFS (FS.mkFsPath ["0.wbblobs"])
+    refCtx <- newRefCtx
+    wbblobs <- WBB.new hasFS refCtx (FS.mkFsPath ["0.wbblobs"])
     wb <- WB.fromMap <$> traverse (traverse (WBB.addBlob hasFS wbblobs)) storedKeys
     let fsps = RunFsPaths (FS.mkFsPath []) (RunNumber 0)
-    r <- Run.fromWriteBuffer hasFS hasBlockIO benchSalt runParams fsps wb wbblobs
+    r <- Run.fromWriteBuffer hasFS hasBlockIO refCtx benchSalt runParams fsps wb wbblobs
     let NumEntries nentriesReal = Run.size r
     assertEqual nentriesReal nentries $ pure ()
     -- 42 to 43 entries per page
@@ -204,6 +206,7 @@ lookupsInBatchesEnv Config {..} = do
          , arenaManager
          , hasFS
          , hasBlockIO
+         , refCtx
          , wbblobs
          , V.singleton r
          , lookupKeys
@@ -222,16 +225,18 @@ lookupsInBatchesCleanup ::
      , ArenaManager RealWorld
      , FS.HasFS IO FS.HandleIO
      , FS.HasBlockIO IO FS.HandleIO
+     , RefCtx
      , Ref (WBB.WriteBufferBlobs IO FS.HandleIO)
      , V.Vector (Ref (Run IO FS.HandleIO))
      , V.Vector SerialisedKey
      )
   -> IO ()
-lookupsInBatchesCleanup (tmpDir, _arenaManager, _hasFS, hasBlockIO, wbblobs, rs, _) = do
+lookupsInBatchesCleanup (tmpDir, _arenaManager, _hasFS, hasBlockIO, refCtx, wbblobs, rs, _) = do
     FS.close hasBlockIO
     forM_ rs releaseRef
     releaseRef wbblobs
     removeDirectoryRecursive tmpDir
+    closeRefCtx refCtx
 
 -- | Generate keys to store and keys to lookup
 lookupsEnv ::

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

@@ -234,7 +234,7 @@ runParams =
 
 benchMerge :: Config -> Benchmark
 benchMerge conf@Config{name} =
-    withEnv $ \ ~(_dir, hasFS, hasBlockIO, runs) ->
+    withEnv $ \ ~(_dir, hasFS, hasBlockIO, refCtx, runs) ->
       bgroup name [
           bench "merge" $
             -- We'd like to do: `whnfAppIO (runs' -> ...) runs`.
@@ -252,7 +252,7 @@ benchMerge conf@Config{name} =
             Cr.perRunEnvWithCleanup
               ((runs,) <$> newIORef Nothing)
               (releaseRun . snd) $ \(runs', ref) -> do
-                !run <- merge hasFS hasBlockIO conf outputRunPaths runs'
+                !run <- merge hasFS hasBlockIO refCtx conf outputRunPaths runs'
                 writeIORef ref $ Just $ releaseRef run
         ]
   where
@@ -270,15 +270,16 @@ benchMerge conf@Config{name} =
 merge ::
      FS.HasFS IO FS.HandleIO
   -> FS.HasBlockIO IO FS.HandleIO
+  -> RefCtx
   -> Config
   -> Run.RunFsPaths
   -> InputRuns
   -> IO (Ref (Run IO FS.HandleIO))
-merge fs hbio Config {..} targetPaths runs = do
+merge fs hbio refCtx Config {..} targetPaths runs = do
     let f = fromMaybe const mergeResolve
     m <- fromMaybe (error "empty inputs, no merge created") <$>
       Merge.new fs hbio benchSalt runParams mergeType f targetPaths runs
-    Merge.stepsToCompletion m stepSize
+    Merge.stepsToCompletion refCtx m stepSize
 
 fsPath :: FS.FsPath
 fsPath = FS.mkFsPath []
@@ -368,39 +369,44 @@ mergeEnv ::
   -> IO ( FilePath -- ^ Temporary directory
         , FS.HasFS IO FS.HandleIO
         , FS.HasBlockIO IO FS.HandleIO
+        , RefCtx
         , InputRuns
         )
 mergeEnv config = do
     sysTmpDir <- getCanonicalTemporaryDirectory
     benchTmpDir <- createTempDirectory sysTmpDir "mergeEnv"
     (hasFS, hasBlockIO) <- FS.ioHasBlockIO (FS.MountPoint benchTmpDir) FS.defaultIOCtxParams
-    runs <- randomRuns hasFS hasBlockIO config (mkStdGen 17)
-    pure (benchTmpDir, hasFS, hasBlockIO, runs)
+    refCtx <- newRefCtx
+    runs <- randomRuns hasFS hasBlockIO refCtx config (mkStdGen 17)
+    pure (benchTmpDir, hasFS, hasBlockIO, refCtx, runs)
 
 mergeEnvCleanup ::
      ( FilePath -- ^ Temporary directory
      , FS.HasFS IO FS.HandleIO
      , FS.HasBlockIO IO FS.HandleIO
+     , RefCtx
      , InputRuns
      )
   -> IO ()
-mergeEnvCleanup (tmpDir, _hasFS, hasBlockIO, runs) = do
+mergeEnvCleanup (tmpDir, _hasFS, hasBlockIO, refCtx, runs) = do
     traverse_ releaseRef runs
     removeDirectoryRecursive tmpDir
     FS.close hasBlockIO
+    checkForgottenRefs refCtx
 
 -- | Generate keys and entries to insert into the write buffer.
 -- They are already serialised to exclude the cost from the benchmark.
 randomRuns ::
      FS.HasFS IO FS.HandleIO
   -> FS.HasBlockIO IO FS.HandleIO
+  -> RefCtx
   -> Config
   -> StdGen
   -> IO InputRuns
-randomRuns hasFS hasBlockIO config@Config {..} rng0 = do
+randomRuns hasFS hasBlockIO refCtx config@Config {..} rng0 = do
     counter <- inputRunPathsCounter
     fmap V.fromList $
-      mapM (unsafeCreateRun hasFS hasBlockIO benchSalt runParams fsPath counter) $
+      mapM (unsafeCreateRun hasFS hasBlockIO refCtx benchSalt runParams fsPath counter) $
         zipWith
           (randomRunData config)
           nentries