new improved benchmark command that can benchmark anything git-annex does

Command/Benchmark.hs

@@ -1,6 +1,6 @@
 {- git-annex benchmark
  -
- - Copyright 2016 Joey Hess <id@joeyh.name>
+ - Copyright 2016-2019 Joey Hess <id@joeyh.name>
  -
  - Licensed under the GNU GPL version 3 or higher.
  -}
@@ -10,114 +10,26 @@
 module Command.Benchmark where
 
 import Command
-import Database.Types
-import qualified Database.Keys.SQL as SQL
-import qualified Database.Queue as H
-import Utility.Tmp
-import Git.FilePath
+import Types.Benchmark
 
 import Criterion.Main
-import Criterion.Internal (runAndAnalyse)
-import Criterion.Monad
-import Control.Monad.IO.Class (liftIO)
-import Control.Monad
-import Control.DeepSeq
-import System.FilePath
-import System.Random
+import Criterion.Main.Options (parseWith, Mode)
 
-cmd :: Command
-cmd = noRepo (withParams benchmark) $
-	dontCheck repoExists $
-		command "benchmark" SectionTesting
-			"run benchmarks"
-			paramNothing
-			(withParams (liftIO . benchmark))
+cmd :: BenchmarkGenerator -> Command
+cmd generator = command "benchmark" SectionTesting
+	"benchmark git-annex commands"
+	paramNothing
+	(seek generator <$$> optParser)
 
-benchmark :: CmdParams -> IO ()
-benchmark _ = withTmpDirIn "." "benchmark" $ \tmpdir -> do
-	-- benchmark different sizes of databases
-	dbs <- mapM (benchDb tmpdir)
-		[ 1000
-		, 10000
-		-- , 100000
-		]
-	runCriterion $
-		bgroup "keys database" $ flip concatMap dbs $ \db ->
-			[ getAssociatedFilesHitBench db
-			, getAssociatedFilesMissBench db
-			, getAssociatedKeyHitBench db
-			, getAssociatedKeyMissBench db
-			, addAssociatedFileOldBench db
-			, addAssociatedFileNewBench db
-			]
+data BenchmarkOptions = BenchmarkOptions CmdParams Mode
 
-getAssociatedFilesHitBench :: BenchDb -> Benchmark
-getAssociatedFilesHitBench ( BenchDb h num) = bench ("getAssociatedFiles from " ++ show num ++ " (hit)") $ nfIO $ do
-	n <- getStdRandom (randomR (1,num))
-	SQL.getAssociatedFiles (keyN n) (SQL.ReadHandle h)
+optParser :: CmdParamsDesc -> Parser BenchmarkOptions
+optParser desc = BenchmarkOptions
+	<$> cmdParams desc
+	-- parse criterion's options
+	<*> parseWith defaultConfig
 
-getAssociatedFilesMissBench :: BenchDb -> Benchmark
-getAssociatedFilesMissBench ( BenchDb h num) = bench ("getAssociatedFiles from " ++ show num ++ " (miss)") $ nfIO $
-	SQL.getAssociatedFiles keyMiss (SQL.ReadHandle h)
-
-getAssociatedKeyHitBench :: BenchDb -> Benchmark
-getAssociatedKeyHitBench (BenchDb h num) = bench ("getAssociatedKey from " ++ show num ++ " (hit)") $ nfIO $ do
-	n <- getStdRandom (randomR (1,num))
-	SQL.getAssociatedKey (fileN n) (SQL.ReadHandle h)
-
-getAssociatedKeyMissBench :: BenchDb -> Benchmark
-getAssociatedKeyMissBench (BenchDb h num) = bench ("getAssociatedKey from " ++ show num ++ " (miss)") $ nfIO $
-	SQL.getAssociatedKey fileMiss (SQL.ReadHandle h)
-
-addAssociatedFileOldBench :: BenchDb -> Benchmark
-addAssociatedFileOldBench ( BenchDb h num) = bench ("addAssociatedFile to " ++ show num ++ " (old)") $ nfIO $ do
-	n <- getStdRandom (randomR (1,num))
-	SQL.addAssociatedFile (keyN n) (fileN n) (SQL.WriteHandle h)
-	H.flushDbQueue h
-
-addAssociatedFileNewBench :: BenchDb -> Benchmark
-addAssociatedFileNewBench ( BenchDb h num) = bench ("addAssociatedFile to " ++ show num ++ " (new)") $ nfIO $ do
-	n <- getStdRandom (randomR (1,num))
-	SQL.addAssociatedFile (keyN n) (fileN (n+1)) (SQL.WriteHandle h)
-	H.flushDbQueue h
-
-populateAssociatedFiles :: H.DbQueue -> Int -> IO ()
-populateAssociatedFiles h num = do
-	forM_ [1..num] $ \n ->
-		SQL.addAssociatedFile (keyN n) (fileN n) (SQL.WriteHandle h)
-	H.flushDbQueue h
-
-keyN :: Int -> IKey
-keyN n = IKey ("key" ++ show n)
-
-fileN :: Int -> TopFilePath
-fileN n = asTopFilePath ("file" ++ show n)
-
-keyMiss :: IKey
-keyMiss = keyN 0 -- 0 is never stored
-
-fileMiss :: TopFilePath
-fileMiss = fileN 0 -- 0 is never stored
-
-data BenchDb = BenchDb H.DbQueue Int
-
-benchDb :: FilePath -> Int -> IO BenchDb
-benchDb tmpdir num = do
-	putStrLn $ "setting up database with " ++ show num
-	H.initDb f SQL.createTables
-	h <- H.openDbQueue f SQL.containedTable
-	populateAssociatedFiles h num
-	return (BenchDb h num)
-  where
-	f = tmpdir </> "db" ++ show num
-
-instance NFData TopFilePath where
-	rnf = rnf . getTopFilePath
-
-instance NFData IKey where
-	rnf (IKey s) = rnf s
-	
--- can't use Criterion's defaultMain here because it looks at
--- command-line parameters
-runCriterion :: Benchmark -> IO ()
-runCriterion = withConfig defaultConfig . runAndAnalyse (const True)
+seek :: BenchmarkGenerator -> BenchmarkOptions -> CommandSeek
+seek generator (BenchmarkOptions ps mode) = do
+	runner <- generator ps
+	liftIO $ runMode mode [ bench (unwords ps) $ nfIO runner ]