tmp-postgres-1.28.0.0: benchmark/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.Config
import qualified Database.PostgreSQL.Simple as PG

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"
  , ");"
  ]

snapshotDir :: DirectoryType
snapshotDir = Permanent "~/.tmp-postgres/bench-mark-1"

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 :: IO Cache
setupCache = do
  cacheInfo <- setupInitDbCache defaultCacheConfig
  void (withConfig (defaultConfig <> cacheConfig cacheInfo) (const $ pure ()))
  pure cacheInfo

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

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

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

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

  , bench "withConfig silent" $ whnfIO $
      withConfig defaultConfig $ const $ pure ()

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

  , bench "with migrate 10x" $ whnfIO $ replicateM 10 $ withConfig defaultConfig $ \db ->
      migrateDb db >> testQuery db

  , setupWithCache $ \theCacheConfig -> do
      bench "with migrate 10x and cache" $ whnfIO $ withConfig theCacheConfig $ \_ -> do
        replicateM_ 10 $ withConfig theCacheConfig $ \db ->
          migrateDb db >> testQuery db

  , setupWithCache $ \theCacheConfig -> bench "withSnapshot migrate 10x and cache" $ whnfIO $ withConfig theCacheConfig $ \db -> do
      migrateDb db
      void $ withSnapshot Temporary db $ \theSnapshotDir -> do
        let theSnapshotConfig = defaultConfig <> snapshotConfig theSnapshotDir
        replicateM_ 10 $ withConfig theSnapshotConfig testQuery

  , setupWithCacheAndSP $ \theCacheConfig -> bench "preexisting snapshot withSnapshot migrate 10x and cache" $ whnfIO $ withConfig theCacheConfig $ \db -> do
      void $ withSnapshot snapshotDir db $ \theSnapshotDir -> do
        let theSnapshotConfig = defaultConfig <> snapshotConfig theSnapshotDir
        replicateM_ 10 $ withConfig theSnapshotConfig testQuery

  , 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 Temporary db $ const $ pure ()

  ]