Main.hs

module Main where
import Control.DeepSeq
import Control.Exception
import Control.Monad
import Criterion.Main hiding (defaultConfig)
import Data.String
import Database.Postgres.Temp.Internal
import Database.Postgres.Temp.Internal.Core
import Database.Postgres.Temp.Internal.Config
import qualified Database.PostgreSQL.Simple as PG
import           System.IO.Temp (createTempDirectory, withTempDirectory)
import qualified Database.PostgreSQL.Simple.Options as Options

data Once a = Once { unOnce :: a }

instance NFData (Once a) where
  rnf x = seq x ()

defaultConfigDefaultInitDb :: Config
defaultConfigDefaultInitDb = mempty
  { logger = pure mempty
  , postgresConfigFile = fastPostgresConfig
  , initDbConfig = pure mempty
  }

createFooDb :: PG.Connection -> Int -> IO ()
createFooDb conn index = void $ PG.execute_ conn $ fromString $ unlines
  [ "CREATE TABLE foo" <> show index
  , "( id int"
  , ");"
  ]

migrateDb :: DB -> IO ()
migrateDb db = do
  let theConnectionString = toConnectionString db

  bracket (PG.connectPostgreSQL theConnectionString) PG.close $
      \conn -> mapM_ (createFooDb conn) [0 .. 100]

testQuery :: DB -> IO ()
testQuery db = do
  let theConnectionString = toConnectionString db

  bracket (PG.connectPostgreSQL theConnectionString) PG.close $
    \conn -> void $ PG.execute_ conn "INSERT INTO foo1 (id) VALUES (1)"

setupCache :: Bool -> IO Cache
setupCache cow = do
  cacheInfo <- setupInitDbCache (defaultCacheConfig { cacheUseCopyOnWrite = cow})
  void (withConfig (defaultConfig <> cacheConfig cacheInfo) (const $ pure ()))
  pure cacheInfo

setupWithCache :: (Config -> Benchmark) -> Benchmark
setupWithCache f = envWithCleanup (setupCache True) cleanupInitDbCache $ f . (defaultConfig <>) . cacheConfig

setupWithCacheNoCow :: (Config -> Benchmark) -> Benchmark
setupWithCacheNoCow f = envWithCleanup (setupCache False) cleanupInitDbCache $ f . (defaultConfig <>) . cacheConfig

setupCacheAndSP :: IO (Cache, Snapshot, Once Config)
setupCacheAndSP = do
  cacheInfo <- setupCache True
  let theCacheConfig = defaultConfig <> cacheConfig cacheInfo
  sp <- either throwIO pure <=< withConfig theCacheConfig $ \db -> do
    migrateDb db
    either throwIO pure =<< takeSnapshot db

  let theConfig = defaultConfig <> snapshotConfig sp <> theCacheConfig

  pure (cacheInfo, sp, Once theConfig)

cleanupCacheAndSP :: (Cache, Snapshot, Once Config) -> IO ()
cleanupCacheAndSP (x, y, _) = cleanupSnapshot y >> cleanupInitDbCache x

setupWithCacheAndSP :: (Config -> Benchmark) -> Benchmark
setupWithCacheAndSP f = envWithCleanup setupCacheAndSP cleanupCacheAndSP $ \ ~(_, _, Once x) -> f x

setupWithCacheAndSP' :: (Snapshot -> Benchmark) -> Benchmark
setupWithCacheAndSP' f = envWithCleanup setupCacheAndSP cleanupCacheAndSP $ \ ~(_, x, _) -> f x

setupCacheAndAction :: IO (Cache, FilePath, Once Config)
setupCacheAndAction = do
  cacheInfo <- setupCache True
  snapshotDir <- createTempDirectory "/tmp" "tmp-postgres-bench-cache"
  let theCacheConfig = defaultConfig <> cacheConfig cacheInfo

  theConfig <- either throwIO pure =<< cacheAction snapshotDir migrateDb theCacheConfig

  pure (cacheInfo, snapshotDir, Once theConfig)

cleanupCacheAndAction :: (Cache, FilePath, Once Config) -> IO ()
cleanupCacheAndAction (c, f, _) = rmDirIgnoreErrors f >> cleanupInitDbCache c

setupWithCacheAndAction :: (FilePath -> Config -> Benchmark) -> Benchmark
setupWithCacheAndAction f = envWithCleanup setupCacheAndAction cleanupCacheAndAction $
  \ ~(_, filePath, Once x) -> f filePath x

main :: IO ()
main = defaultMain
  [ bench "with" $ whnfIO $ with $ const $ pure ()

  , bench "withConfig no --no-sync" $ whnfIO $
      withConfig defaultConfigDefaultInitDb $ const $ pure ()

  , bench "withConfig verbose" $ whnfIO $
      withConfig verboseConfig $ const $ pure ()

  , bench "withConfig db create" $ whnfIO $
      withConfig (optionsToDefaultConfig (mempty { Options.dbname = pure "test" } )) $
        const $ pure ()

  , setupWithCacheNoCow $ \theConfig -> bench "withConfig silent cache no cow" $ whnfIO $
      withConfig theConfig $ const $ pure ()

  , setupWithCache $ \theCacheConfig -> bench "withConfig silent cache" $ whnfIO $
      withConfig theCacheConfig $ const $ pure ()

 , setupWithCache $ \theCacheConfig -> bench "cache action and recache and cache" $ whnfIO $ withTempDirectory "/tmp" "tmp-postgres-bench-cache" $ \snapshotDir -> do
      newConfig <- either throwIO pure =<< cacheAction snapshotDir migrateDb theCacheConfig
      either throwIO pure =<< flip withConfig testQuery
        =<< either throwIO pure =<< cacheAction snapshotDir migrateDb newConfig

  , setupWithCacheAndAction $ \snapshotDir theCacheConfig -> bench "pre-cache action and recache" $ whnfIO $ do
      either throwIO pure =<< flip withConfig testQuery
        =<< either throwIO pure =<< cacheAction snapshotDir migrateDb theCacheConfig

  , setupWithCacheAndSP $ \theConfig -> bench "withConfig pre-setup with withSnapshot" $ whnfIO $
      void $ withConfig theConfig $ const $ pure ()

  , setupWithCacheAndSP' $ \sp -> bench "snapshotConfig" $ whnfIO $ void $ flip withConfig (const $ pure ())
    $ snapshotConfig sp

  , bench "migrateDb" $ perRunEnvWithCleanup (either throwIO (pure . Once) =<< startConfig defaultConfig) (stop . unOnce) $
      \ ~(Once db) -> migrateDb db

  , bench "withSnapshot" $ perRunEnvWithCleanup (either throwIO (pure . Once) =<< startConfig defaultConfig) (stop . unOnce) $
      \ ~(Once db) -> void $ withSnapshot db $ const $ pure ()

  , bench "stopGracefully" $ perRunEnvWithCleanup (either throwIO (pure . Once) =<< start) (stop . unOnce) $
    \ ~(Once db) -> do
      void $ stopPostgresGracefully db
      stop db

  , bench "stop" $ perRunEnvWithCleanup (either throwIO (pure . Once) =<< start) (stop . unOnce) $
      \ ~(Once db) -> stop db

  , bench "stop serial" $ perRunEnvWithCleanup (either throwIO (pure . Once) =<< start) (stop . unOnce) $
      \ ~(Once DB {..}) ->
          stopPlan dbPostgresProcess >> cleanupConfig dbResources
  ]