[lmdb facebookincubator#4] Add support for LMDB in benchmarks

Author: simonmar

glean.cabal.in

@@ -1884,6 +1884,7 @@ library bench-lib
         BenchDB
     build-depends:
         glean:backend-api,
+        glean:bench-util,
         glean:client-hs,
         glean:core,
         glean:db,

glean/bench/BenchDB.hs

@@ -16,7 +16,7 @@ import qualified Data.Text as Text
 
 import Glean
 import Glean.Backend.Types (loadPredicates)
-import Glean.Database.Test (withEmptyTestDB, completeTestDB)
+import Glean.Database.Test
 import Glean.Database.Open
 import Glean.Database.Write.Batch
 import qualified Glean.Schema.Sys.Types as Sys
@@ -26,9 +26,20 @@ import qualified Glean.Schema.Sys as Sys
 import qualified Glean.Schema.Cxx1 as Cxx
 import qualified Glean.Schema.GleanTest as Glean.Test
 import Glean.Typed
+import Glean.Util.Benchmark
 
-withBenchDB :: Int -> (forall b . Backend b => b -> Repo -> IO a) -> IO a
-withBenchDB num act = withEmptyTestDB [] $ \env repo -> do
+withBenchDB
+  :: GleanBenchConfig
+  -> Int
+  -> (forall b . Backend b => b -> Repo -> IO a)
+  -> IO a
+withBenchDB conf num act = do
+  let
+    settings
+      | useLMDB conf = [setLMDBStorage]
+      | otherwise = []
+
+  withEmptyTestDB settings $ \env repo -> do
   withOpenDatabase env repo $ \odb ->
     void $ return odb
 

glean/bench/MakeFactBench.hs

@@ -21,7 +21,7 @@ import Glean.Util.Benchmark
 import Glean.Typed
 
 main :: IO ()
-main = benchmarkMain $ \run -> do
+main = benchmarkMain $ \_conf run -> do
   withEmptyTestDB [] $ \env repo -> do
     predicates <- Backend.loadPredicates env repo
       [ Cxx.allPredicates

glean/bench/QueryBench.hs

@@ -29,7 +29,8 @@ import Glean.Util.Benchmark
 import BenchDB
 
 main :: IO ()
-main = benchmarkMain $ \run -> withBenchDB 10000 $ \env repo -> do
+main = benchmarkMain $ \conf run ->
+  withBenchDB conf 10000 $ \env repo -> do
   let
     nestedAngle :: Query Cxx.FunctionName
     nestedAngle = angle "cxx1.FunctionName { name = \"x1\" }"

glean/bench/RenameBench.hs

@@ -22,7 +22,7 @@ import Glean.Util.Benchmark
 import TestBatch
 
 main :: IO ()
-main = benchmarkMain $ \run ->
+main = benchmarkMain $ \_conf run ->
   withEmptyTestDB [] $ \env repo -> do
   schema <- loadDbSchema env repo
   batch <- testBatch 500000 env repo

glean/bench/lib/Glean/Util/Benchmark.hs

@@ -9,15 +9,31 @@
 module Glean.Util.Benchmark
   ( BenchmarkRunner
   , benchmarkMain
+  , GleanBenchConfig(..)
   ) where
 
 import Criterion.Main
 import Criterion.Main.Options
+import Data.Default
+import Options.Applicative
 
 import Glean.Init
 
 type BenchmarkRunner = [Benchmark] -> IO ()
 
-benchmarkMain :: (BenchmarkRunner -> IO ()) -> IO ()
-benchmarkMain exec = withOptions (describe defaultConfig) $
-  exec . runMode
+benchmarkMain :: (GleanBenchConfig -> BenchmarkRunner -> IO ()) -> IO ()
+benchmarkMain exec =
+  withOptions
+    (describeWith ((,) <$> opts <*> parseWith defaultConfig)) $ \(conf, crit) ->
+    exec conf (runMode crit)
+
+data GleanBenchConfig = GleanBenchConfig {
+    useLMDB :: Bool
+  }
+
+instance Default GleanBenchConfig where
+  def = GleanBenchConfig { useLMDB = False }
+
+opts :: Parser GleanBenchConfig
+opts = GleanBenchConfig
+  <$> switch (long "lmdb" <> help "use LMDB (otherwise use RocksDB)")

glean/test/tests/TimeoutTest.hs

@@ -9,6 +9,7 @@
 {-# LANGUAGE QuasiQuotes #-}
 module TimeoutTest (main) where
 
+import Data.Default
 import Test.HUnit
 
 import TestRunner
@@ -21,7 +22,7 @@ import qualified Glean.Schema.GleanTest.Types as Glean.Test
 import BenchDB
 
 timeoutTest :: Test
-timeoutTest = TestCase $ withBenchDB 10000 $ \env repo -> do
+timeoutTest = TestCase $ withBenchDB def 10000 $ \env repo -> do
   -- The test spends most of its time building the DB (withBenchDB above),
   -- so to keep things reasonable in debug mode we limit the number of
   -- facts to 10K and set the time limit to 1ms which is as low as we