{-# LANGUAGE RecordWildCards, LambdaCase, ScopedTypeVariables, DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric, OverloadedStrings #-}
module Database.Postgres.Temp.Internal where
import System.IO.Temp
import System.Process
import System.Process.Internals
import Control.Concurrent
import System.IO
import System.Exit
import System.Directory
import qualified Network.Socket as N
import Control.Exception
import Data.Typeable
import GHC.Generics
import System.Posix.Signals
import qualified Database.PostgreSQL.Simple as PG
import qualified Data.ByteString.Char8 as BSC
import Control.Monad (void, forever)
import Network.Socket.Free (openFreePort)
import Data.Foldable
import Control.Concurrent.Async(race_)
import Data.IORef
getFreePort :: IO Int
getFreePort = do
(port, socket) <- openFreePort
N.close socket
pure port
waitForDB :: String -> IO ()
waitForDB connStr = do
eresult <- try $ bracket (PG.connectPostgreSQL (BSC.pack connStr)) PG.close $ \_ -> return ()
case eresult of
Left (_ :: IOError) -> threadDelay 10000 >> waitForDB connStr
Right _ -> return ()
withLock :: MVar a -> IO b -> IO b
withLock m f = withMVar m (const f)
data DB = DB
{ mainDir :: FilePath
, connectionString :: String
, extraOptions :: [(String, String)]
, stdErr :: Handle
, stdOut :: Handle
, pidLock :: MVar ()
, port :: Int
, socketClass :: SocketClass
, pid :: IORef (Maybe ProcessHandle)
}
data SocketClass = Localhost | Unix
deriving (Show, Eq, Read, Ord, Enum, Bounded, Generic, Typeable)
start :: [(String, String)]
-> IO (Either StartError DB)
start options = startWithHandles Unix options stdout stderr
startLocalhost :: [(String, String)]
-> IO (Either StartError DB)
startLocalhost options = startWithHandles Localhost options stdout stderr
fourth :: (a, b, c, d) -> d
fourth (_, _, _, x) = x
procWith :: Handle -> Handle -> String -> [String] -> CreateProcess
procWith stdOut stdErr cmd args =
(proc cmd args)
{ std_err = UseHandle stdErr
, std_out = UseHandle stdOut
}
config :: Maybe FilePath -> String
config mMainDir = unlines $
[ "shared_buffers = 12MB"
, "fsync = off"
, "synchronous_commit = off"
, "full_page_writes = off"
, "log_min_duration_statement = 0"
, "log_connections = on"
, "log_disconnections = on"
, "client_min_messages = ERROR"
] ++ maybe ["listen_addresses = '127.0.0.1'"] (\x -> ["unix_socket_directories = '" ++ x ++ "'", "listen_addresses = ''"]) mMainDir
data StartError
= InitDBFailed ExitCode
| CreateDBFailed [String] ExitCode
| StartPostgresFailed [String] ExitCode
| StartPostgresDisappeared [String]
deriving (Show, Eq, Typeable)
instance Exception StartError
throwIfError :: (ExitCode -> StartError) -> ExitCode -> IO ()
throwIfError f e = case e of
ExitSuccess -> return ()
_ -> throwIO $ f e
pidString :: ProcessHandle -> IO String
pidString phandle = withProcessHandle phandle (\case
OpenHandle p -> return $ show p
OpenExtHandle _ _ _ -> return ""
ClosedHandle _ -> return ""
)
runProcessWith :: Handle -> Handle -> String -> String -> [String] -> IO ExitCode
runProcessWith stdOut stdErr name cmd args
= createProcess_ name (procWith stdOut stdErr cmd args)
>>= waitForProcess . fourth
startWithHandles :: SocketClass
-> [(String, String)]
-> Handle
-> Handle
-> IO (Either StartError DB)
startWithHandles socketClass options stdOut stdErr = do
mainDir <- createTempDirectory "/tmp" "tmp-postgres"
startWithHandlesAndDir socketClass options mainDir stdOut stdErr
startWithHandlesAndDir :: SocketClass
-> [(String, String)]
-> FilePath
-> Handle
-> Handle
-> IO (Either StartError DB)
startWithHandlesAndDir = startWithLogger $ \_ -> return ()
data AnotherPostgresProcessActive = AnotherPostgresProcessActive
deriving (Show, Eq, Typeable)
instance Exception AnotherPostgresProcessActive
waitOnPostgres :: DB -> IO ()
waitOnPostgres DB {..} = do
let postgresOptions = makePostgresOptions extraOptions (mainDir ++ "/data") port
checkForCrash = readIORef pid >>= \case
Nothing -> throwIO $ StartPostgresDisappeared postgresOptions
Just thePid -> do
mExitCode <- getProcessExitCode thePid
for_ mExitCode (throwIO . StartPostgresFailed postgresOptions)
host = case socketClass of
Localhost -> "127.0.0.1"
Unix -> mainDir
makeConnectionString dbName = "postgresql:///"
++ dbName ++ "?host=" ++ host ++ "&port=" ++ show port
waitForDB (makeConnectionString "template1") `race_`
forever (checkForCrash >> threadDelay 100000)
reloadConfig :: DB -> IO ()
reloadConfig DB {..} = do
mHandle <- readIORef pid
for_ mHandle $ \theHandle -> do
mPid <- getPid theHandle
for_ mPid $ signalProcess sigHUP
startPostgres :: DB -> IO ()
startPostgres db@DB {..} = withLock pidLock $ do
readIORef pid >>= \case
Just _ -> throwIO AnotherPostgresProcessActive
Nothing -> do
let postgresOptions = makePostgresOptions extraOptions (mainDir ++ "/data") port
bracketOnError
(runPostgres stdErr stdOut postgresOptions)
(const $ stopPostgres db)
$ \thePid -> do
writeIORef pid $ Just thePid
waitOnPostgres db
stopPostgres :: DB -> IO (Maybe ExitCode)
stopPostgres db@DB {..} = withLock pidLock $ readIORef pid >>= \case
Nothing -> pure Nothing
Just pHandle -> do
withProcessHandle pHandle (\case
OpenHandle p -> do
terminateConnections db
signalProcess sigINT p
OpenExtHandle _ _ _ -> pure ()
ClosedHandle _ -> return ()
)
exitCode <- waitForProcess pHandle
writeIORef pid Nothing
pure $ Just exitCode
makePostgresOptions :: [(String, String)]
-> FilePath
-> Int
-> [String]
makePostgresOptions options dataDir port =
let extraOptions = map (\(key, value) -> "--" ++ key ++ "=" ++ value) options
in ["-D", dataDir, "-p", show port] ++ extraOptions
runPostgres :: Handle
-> Handle
-> [String]
-> IO ProcessHandle
runPostgres theStdOut theStdErr postgresOptions = do
fmap fourth $ createProcess_ "postgres" $
procWith theStdOut theStdErr "postgres" postgresOptions
data Event
= InitDB
| WriteConfig
| FreePort
| StartPostgres
| WaitForDB
| CreateDB
| Finished
deriving (Show, Eq, Enum, Bounded, Ord)
rmDirIgnoreErrors :: FilePath -> IO ()
rmDirIgnoreErrors mainDir =
removeDirectoryRecursive mainDir `catch` (\(_ :: IOException) -> return ())
startWithLogger :: (Event -> IO ())
-> SocketClass
-> [(String, String)]
-> FilePath
-> Handle
-> Handle
-> IO (Either StartError DB)
startWithLogger logger socketClass options mainDir stdOut stdErr = try $ flip onException (rmDirIgnoreErrors mainDir) $ do
let dataDir = mainDir ++ "/data"
logger InitDB
initDBExitCode <- runProcessWith stdOut stdErr "initdb"
"initdb" ["-E", "UNICODE", "-A", "trust", "--nosync", "-D", dataDir]
throwIfError InitDBFailed initDBExitCode
logger WriteConfig
writeFile (dataDir ++ "/postgresql.conf") $ config $ if socketClass == Unix then Just mainDir else Nothing
logger FreePort
port <- getFreePort
let host = case socketClass of
Localhost -> "127.0.0.1"
Unix -> mainDir
let makeConnectionString dbName = "postgresql:///"
++ dbName ++ "?host=" ++ host ++ "&port=" ++ show port
connectionString = makeConnectionString "test"
logger StartPostgres
pidLock <- newMVar ()
let postgresOptions = makePostgresOptions options dataDir port
createDBResult = do
thePid <- runPostgres stdOut stdErr postgresOptions
pid <- newIORef $ Just thePid
pure $ DB mainDir connectionString options stdErr stdOut pidLock port socketClass pid
bracketOnError createDBResult stop $ \result -> do
let checkForCrash = readIORef (pid result) >>= \case
Nothing -> throwIO $ StartPostgresDisappeared postgresOptions
Just thePid -> do
mExitCode <- getProcessExitCode thePid
for_ mExitCode (throwIO . StartPostgresFailed postgresOptions)
logger WaitForDB
waitForDB (makeConnectionString "template1") `race_`
forever (checkForCrash >> threadDelay 100000)
logger CreateDB
let createDBHostArgs = case socketClass of
Unix -> ["-h", mainDir]
Localhost -> ["-h", "127.0.0.1"]
let createDBArgs = createDBHostArgs ++ ["-p", show port, "test"]
throwIfError (CreateDBFailed createDBArgs) =<<
runProcessWith stdOut stdErr "createDB" "createdb" createDBArgs
logger Finished
return result
startAndLogToTmp :: [(String, String)]
-> IO (Either StartError DB)
startAndLogToTmp options = do
mainDir <- createTempDirectory "/tmp" "tmp-postgres"
stdOutFile <- openFile (mainDir ++ "/" ++ "output.txt") WriteMode
stdErrFile <- openFile (mainDir ++ "/" ++ "error.txt") WriteMode
startWithHandlesAndDir Unix options mainDir stdOutFile stdErrFile
terminateConnections :: DB -> IO ()
terminateConnections DB {..} = do
e <- try $ bracket (PG.connectPostgreSQL $ BSC.pack connectionString)
PG.close
$ \conn -> do
void $ PG.execute_ conn "select pg_terminate_backend(pid) from pg_stat_activity where datname='test';"
case e of
Left (_ :: IOError) -> pure ()
Right _ -> pure ()
stop :: DB -> IO (Maybe ExitCode)
stop db@DB {..} = do
result <- stopPostgres db
removeDirectoryRecursive mainDir
return result