IntersectMBO · dcoutts · Jul 19, 2024
@@ -3,12 +3,10 @@
 module Main ( main ) where
 
 import           Control.DeepSeq
-import           Control.Exception (assert)
+import           Control.Exception (assert, evaluate)
 import           Control.Monad
-import           Control.Monad.Class.MonadST
 import           Control.Monad.Primitive
 import           Control.Monad.ST.Strict (ST, runST)
-import           Data.Arena (ArenaManager, newArenaManager, withArena)
 import           Data.Bits ((.&.))
 import           Data.BloomFilter (Bloom)
 import qualified Data.BloomFilter as Bloom
@@ -26,6 +24,7 @@ import           Database.LSMTree.Internal.Entry (Entry (Insert),
                      NumEntries (..))
 import           Database.LSMTree.Internal.IndexCompact (IndexCompact)
 import           Database.LSMTree.Internal.Lookup
+import           Database.LSMTree.Internal.PageAlloc
 import           Database.LSMTree.Internal.Paths (RunFsPaths (RunFsPaths))
 import           Database.LSMTree.Internal.Run (Run)
 import qualified Database.LSMTree.Internal.Run as Run
@@ -136,7 +135,7 @@ benchmarks !caching = withFS $ \hfs hbio -> do
 #ifdef NO_IGNORE_ASSERTS
     putStrLn "BENCHMARKING A BUILD WITH -fno-ignore-asserts"
 #endif
-    arenaManager <- newArenaManager
+    pagealloc <- newPageAlloc
     enabled <- getRTSStatsEnabled
     when (not enabled) $ fail "Need RTS +T statistics enabled"
     let runSizes = lsmStyleRuns benchmarkSizeBase
@@ -197,17 +196,17 @@ benchmarks !caching = withFS $ \hfs hbio -> do
     _bindexSearches <-
       benchmark "benchIndexSearches"
                 "Calculate batches of keys, perform bloom queries for each batch, and perform index searches for positively queried keys in each batch. Net time/allocation is the result of subtracting the cost of benchGenKeyBatches and benchBloomQueries."
-                (benchIndexSearches arenaManager blooms indexes handles keyRng0) benchmarkNumLookups
+                (benchIndexSearches pagealloc blooms indexes handles keyRng0) benchmarkNumLookups
                 (x1 + x2, y1 + y2)
     _bprepLookups <-
       benchmark "benchPrepLookups"
                 "Calculate batches of keys, and prepare lookups for each batch. This is roughly doing the same amount of work as benchIndexSearches. Net time/allocation is the result of subtracting the cost of benchGenKeyBatches."
-                (benchPrepLookups arenaManager blooms indexes handles keyRng0) benchmarkNumLookups
+                (benchPrepLookups pagealloc blooms indexes handles keyRng0) benchmarkNumLookups
                 bgenKeyBatches
     _blookupsIO <-
       benchmark "benchLookupsIO"
                 "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."
-                (benchLookupsIO hbio arenaManager benchmarkResolveSerialisedValue runs blooms indexes handles keyRng0) benchmarkNumLookups
+                (benchLookupsIO hbio pagealloc benchmarkResolveSerialisedValue runs blooms indexes handles keyRng0) benchmarkNumLookups
                 bgenKeyBatches
 
     traceMarkerIO "Cleaning up"
@@ -303,8 +302,8 @@ withFS ::
      (FS.HasFS IO FS.HandleIO -> FS.HasBlockIO IO FS.HandleIO -> IO a)
   -> IO a
 withFS action = do
-    let hfs = FS.ioHasFS (FS.MountPoint "")
-    exists <- FS.doesDirectoryExist hfs (FS.mkFsPath ["_bench_lookups"])
+    let hfs = FS.ioHasFS (FS.MountPoint "_bench_lookups")
+    exists <- FS.doesDirectoryExist hfs (FS.mkFsPath [""])
     unless exists $ error ("_bench_lookups directory does not exist")
     FS.withIOHasBlockIO hfs FS.defaultIOCtxParams $ \hbio ->
       action hfs hbio
@@ -427,43 +426,46 @@ benchBloomQueries !bs !keyRng !n
 -- | This gives us the combined cost of calculating batches of keys, performing
 -- bloom queries for each batch, and performing index searches for each batch.
 benchIndexSearches
-  :: ArenaManager RealWorld
+  :: PageAlloc RealWorld
   -> V.Vector (Bloom SerialisedKey)
   -> V.Vector IndexCompact
   -> V.Vector (FS.Handle h)
   -> StdGen
   -> Int
   -> IO ()
-benchIndexSearches !arenaManager !bs !ics !hs !keyRng !n
+benchIndexSearches !pagealloc !bs !ics !hs !keyRng !n
   | n <= 0 = pure ()
   | otherwise = do
     let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize
         !rkixs = bloomQueriesDefault bs ks
-    !_ioops <- withArena arenaManager $ \arena -> stToIO $ indexSearches arena ics hs ks rkixs
-    benchIndexSearches arenaManager bs ics hs keyRng' (n-benchmarkGenBatchSize)
+    withPages pagealloc (VU.length rkixs) $ \pages -> do
+       _ <- evaluate (indexSearches ics hs ks pages rkixs)
+       return ()
+    benchIndexSearches pagealloc bs ics hs keyRng' (n-benchmarkGenBatchSize)
 
 -- | This gives us the combined cost of calculating batches of keys, and
 -- preparing lookups for each batch.
 benchPrepLookups
-  :: ArenaManager RealWorld
+  :: PageAlloc RealWorld
   -> V.Vector (Bloom SerialisedKey)
   -> V.Vector IndexCompact
   -> V.Vector (FS.Handle h)
   -> StdGen
   -> Int
   -> IO ()
-benchPrepLookups !arenaManager !bs !ics !hs !keyRng !n
+benchPrepLookups !pagealloc !bs !ics !hs !keyRng !n
   | n <= 0 = pure ()
   | otherwise = do
       let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize
-      (!_rkixs, !_ioops) <- withArena arenaManager $ \arena -> stToIO $ prepLookups arena bs ics hs ks
-      benchPrepLookups arenaManager bs ics hs keyRng' (n-benchmarkGenBatchSize)
+      withPreparedLookups pagealloc bs ics hs ks $ \rkixs ioops ->
+        void $ evaluate rkixs >> evaluate ioops
+      benchPrepLookups pagealloc bs ics hs keyRng' (n-benchmarkGenBatchSize)
 
 -- | This gives us the combined cost of calculating batches of keys, and
 -- performing disk lookups for each batch.
 benchLookupsIO ::
      FS.HasBlockIO IO h
-  -> ArenaManager RealWorld
+  -> PageAlloc RealWorld
   -> ResolveSerialisedValue
   -> V.Vector (Run (FS.Handle h))
   -> V.Vector (Bloom SerialisedKey)
@@ -472,12 +474,12 @@ benchLookupsIO ::
   -> StdGen
   -> Int
   -> IO ()
-benchLookupsIO !hbio !arenaManager !resolve !rs !bs !ics !hs !keyRng !n
+benchLookupsIO !hbio !pagealloc !resolve !rs !bs !ics !hs !keyRng !n
   | n <= 0 = pure ()
   | otherwise = do
       let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize
-      !_ <- lookupsIO hbio arenaManager resolve rs bs ics hs ks
-      benchLookupsIO hbio arenaManager resolve rs bs ics hs keyRng' (n-benchmarkGenBatchSize)
+      !_ <- lookupsIO hbio pagealloc resolve rs bs ics hs ks
+      benchLookupsIO hbio pagealloc resolve rs bs ics hs keyRng' (n-benchmarkGenBatchSize)
 
 {-------------------------------------------------------------------------------
   Utilities

@@ -2,13 +2,10 @@
 
 module Bench.Database.LSMTree.Internal.Lookup (benchmarks) where
 
-import           Control.Exception (assert)
+import           Control.Exception (assert, evaluate)
 import           Control.Monad
-import           Control.Monad.ST.Strict (stToIO)
 import           Criterion.Main (Benchmark, bench, bgroup, env, envWithCleanup,
-                     perRunEnv, perRunEnvWithCleanup, whnf, whnfAppIO)
-import           Data.Arena (ArenaManager, closeArena, newArena,
-                     newArenaManager, withArena)
+                     perRunEnvWithCleanup, whnf, whnfAppIO)
 import           Data.Bifunctor (Bifunctor (..))
 import qualified Data.List as List
 import           Data.Map.Strict (Map)
@@ -21,7 +18,10 @@ import           Database.LSMTree.Extras.Random (frequency,
 import           Database.LSMTree.Extras.UTxO
 import           Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))
 import           Database.LSMTree.Internal.Lookup (bloomQueriesDefault,
-                     indexSearches, intraPageLookups, lookupsIO, prepLookups)
+                     cleanupPreparedLookups, indexSearches, intraPageLookups,
+                     lookupsIO, unmanagedAllocatePagesForIndexSearches,
+                     unmanagedPreparedLookups, withPreparedLookups)
+import           Database.LSMTree.Internal.PageAlloc
 import           Database.LSMTree.Internal.Paths (RunFsPaths (..))
 import           Database.LSMTree.Internal.Run (Run)
 import qualified Database.LSMTree.Internal.Run as Run
@@ -77,7 +77,7 @@ benchmarks = bgroup "Bench.Database.LSMTree.Internal.Lookup" [
 
 benchLookups :: Config -> Benchmark
 benchLookups conf@Config{name} =
-    withEnv $ \ ~(_dir, arenaManager, _hasFS, hasBlockIO, rs, ks) ->
+    withEnv $ \ ~(_dir, pagealloc, _hasFS, hasBlockIO, rs, ks) ->
       env ( pure ( V.map Run.runFilter rs
                  , V.map Run.runIndex rs
                  , V.map Run.runKOpsFile rs
@@ -90,24 +90,33 @@ benchLookups conf@Config{name} =
               whnf (\ks' -> bloomQueriesDefault blooms ks') ks
             -- The compact index is only searched for (true and false) positive
             -- lookup keys. We use whnf here because the result is
-          , env (pure $ bloomQueriesDefault blooms ks) $ \rkixs ->
-              bench "Compact index search" $
-                whnfAppIO (\ks' -> withArena arenaManager $ \arena -> stToIO $ indexSearches arena indexes kopsFiles ks' rkixs) ks
+          , env (do let !rkixs = bloomQueriesDefault blooms ks
+                    pages <- unmanagedAllocatePagesForIndexSearches pagealloc rkixs
+                    return (rkixs, pages))
+                (\ ~(rkixs, pages) ->
+                  bench "Compact index search" $
+                    whnf (\ks' -> indexSearches indexes kopsFiles
+                                                ks' pages rkixs) ks)
             -- prepLookups combines bloom filter querying and index searching.
             -- The implementation forces the results to WHNF, so we use
             -- whnfAppIO here instead of nfAppIO.
           , bench "Lookup preparation in memory" $
-              whnfAppIO (\ks' -> withArena arenaManager $ \arena -> stToIO $ prepLookups arena blooms indexes kopsFiles ks') ks
+              whnfAppIO (\ks' -> withPreparedLookups
+                                     pagealloc blooms indexes
+                                     kopsFiles ks' $ \rkixs ioops ->
+                                   void $ evaluate rkixs >> evaluate ioops
+                        ) ks
             -- Submit the IOOps we get from prepLookups to HasBlockIO. We use
             -- perRunEnv because IOOps contain mutable buffers, so we want fresh
             -- ones for each run of the benchmark. We manually evaluate the
             -- result to WHNF since it is unboxed vector.
           , bench "Submit IOOps" $
               -- TODO: here arena is destroyed too soon
               -- but it should be fine for non-debug code
-              perRunEnv (withArena arenaManager $ \arena -> stToIO $ prepLookups arena blooms indexes kopsFiles ks) $ \ ~(_rkixs, ioops) -> do
-                !_ioress <- FS.submitIO hasBlockIO ioops
-                pure ()
+              perRunEnvWithCleanup
+                (unmanagedPreparedLookups pagealloc blooms indexes kopsFiles ks)
+                (cleanupPreparedLookups pagealloc)
+                (\ ~(_, ioops, _) -> void $ evaluate =<< FS.submitIO hasBlockIO ioops)
             -- When IO result have been collected, intra-page lookups searches
             -- through the raw bytes (representing a disk page) for the lookup
             -- key. Again, we use perRunEnv here because IOOps contain mutable
@@ -117,20 +126,21 @@ benchLookups conf@Config{name} =
             -- only compute WHNF.
           , bench "Perform intra-page lookups" $
               perRunEnvWithCleanup
-                ( newArena arenaManager >>= \arena ->
-                  stToIO (prepLookups arena blooms indexes kopsFiles ks) >>= \(rkixs, ioops) ->
-                  FS.submitIO hasBlockIO ioops >>= \ioress ->
-                  pure (rkixs, ioops, ioress, arena)
-                )
-                (\(_, _, _, arena) -> closeArena arenaManager arena) $ \ ~(rkixs, ioops, ioress, _) -> do
-                  !_ <- intraPageLookups resolveV rs ks rkixs ioops ioress
-                  pure ()
+                (do (rkixs, ioops, pages) <-
+                      unmanagedPreparedLookups pagealloc blooms
+                                               indexes kopsFiles ks
+                    ioress <- FS.submitIO hasBlockIO ioops
+                    return ((rkixs, ioops, pages), ioress))
+                (cleanupPreparedLookups pagealloc . fst)
+                (\ ~((rkixs, ioops, _pages), ioress) ->
+                    void $ evaluate =<< intraPageLookups resolveV rs ks
+                                                         rkixs ioops ioress)
             -- The whole shebang: lookup preparation, doing the IO, and then
             -- performing intra-page-lookups. Again, we evaluate the result to
             -- WHNF because it is the same result that intraPageLookups produces
             -- (see above).
           , bench "Lookups in IO" $
-              whnfAppIO (\ks' -> lookupsIO hasBlockIO arenaManager resolveV rs blooms indexes kopsFiles ks') ks
+              whnfAppIO (\ks' -> lookupsIO hasBlockIO pagealloc resolveV rs blooms indexes kopsFiles ks') ks
           ]
   where
     withEnv = envWithCleanup
@@ -162,14 +172,14 @@ data Config = Config {
 lookupsInBatchesEnv ::
      Config
   -> IO ( FilePath -- ^ Temporary directory
-        , ArenaManager RealWorld
+        , PageAlloc RealWorld
         , FS.HasFS IO FS.HandleIO
         , FS.HasBlockIO IO FS.HandleIO
         , V.Vector (Run (FS.Handle FS.HandleIO))
         , V.Vector SerialisedKey
         )
 lookupsInBatchesEnv Config {..} = do
-    arenaManager <- newArenaManager
+    pagealloc <- newPageAlloc
     sysTmpDir <- getCanonicalTemporaryDirectory
     benchTmpDir <- createTempDirectory sysTmpDir "lookupsInBatchesEnv"
     (storedKeys, lookupKeys) <- lookupsEnv (mkStdGen 17) nentries npos nneg
@@ -184,7 +194,7 @@ lookupsInBatchesEnv Config {..} = do
     assert (npagesReal * 42 <= nentriesReal) $ pure ()
     assert (npagesReal * 43 >= nentriesReal) $ pure ()
     pure ( benchTmpDir
-         , arenaManager
+         , pagealloc
          , hasFS
          , hasBlockIO
          , V.singleton r
@@ -193,14 +203,14 @@ lookupsInBatchesEnv Config {..} = do
 
 lookupsInBatchesCleanup ::
      ( FilePath -- ^ Temporary directory
-     , ArenaManager RealWorld
+     , PageAlloc RealWorld
      , FS.HasFS IO FS.HandleIO
      , FS.HasBlockIO IO FS.HandleIO
      , V.Vector (Run (FS.Handle FS.HandleIO))
      , V.Vector SerialisedKey
      )
   -> IO ()
-lookupsInBatchesCleanup (tmpDir, _arenaManager, hasFS, hasBlockIO, rs, _) = do
+lookupsInBatchesCleanup (tmpDir, _pagealloc, hasFS, hasBlockIO, rs, _) = do
     FS.close hasBlockIO
     forM_ rs $ Run.removeReference hasFS
     removeDirectoryRecursive tmpDir

@@ -30,7 +30,7 @@ benchmarks: True
 constraints: bloomfilter <0
 
 -- comment me if you are benchmarking
-import: cabal.project.debug
+--import: cabal.project.debug
 
 -- comment me if you don't have liburing installed
 --

@@ -51,7 +51,6 @@ library
   hs-source-dirs:  src
   exposed-modules:
     Control.Concurrent.Class.MonadSTM.RWVar
-    Data.Arena
     Data.Map.Range
     Database.LSMTree.Common
     Database.LSMTree.Internal
@@ -73,6 +72,7 @@ library
     Database.LSMTree.Internal.Normal
     Database.LSMTree.Internal.PageAcc
     Database.LSMTree.Internal.PageAcc1
+    Database.LSMTree.Internal.PageAlloc
     Database.LSMTree.Internal.Paths
     Database.LSMTree.Internal.Primitive
     Database.LSMTree.Internal.Range
@@ -233,7 +233,6 @@ test-suite lsm-tree-test
     Database.LSMTree.ModelIO.Normal
     Database.LSMTree.ModelIO.Session
     Test.Control.Concurrent.Class.MonadSTM.RWVar
-    Test.Data.Arena
     Test.Database.LSMTree.Class.Monoidal
     Test.Database.LSMTree.Class.Normal
     Test.Database.LSMTree.Generators
@@ -247,6 +246,7 @@ test-suite lsm-tree-test
     Test.Database.LSMTree.Internal.Monkey
     Test.Database.LSMTree.Internal.PageAcc
     Test.Database.LSMTree.Internal.PageAcc1
+    Test.Database.LSMTree.Internal.PageAlloc
     Test.Database.LSMTree.Internal.RawOverflowPage
     Test.Database.LSMTree.Internal.RawPage
     Test.Database.LSMTree.Internal.Run
@@ -394,7 +394,6 @@ benchmark lsm-tree-bench-lookups
     , bytestring
     , deepseq
     , fs-api
-    , io-classes
     , lsm-tree
     , lsm-tree:blockio-api
     , lsm-tree:bloomfilter