module System.Posix.Graceful
( GracefulSettings(..)
, GracefulWorker(..)
, graceful
) where
import Control.Concurrent ( newEmptyMVar, putMVar, takeMVar )
import Control.Concurrent.STM ( newTVarIO )
import Control.Exception ( IOException, bracket, bracket_, try, throwIO )
import Control.Monad ( replicateM, void, when )
import Network.Socket.Wrapper ( Socket(..), socket, mkSocket
, connect, close, accept, bindSocket, listen
, send, recv, sendFd, recvFd, fdSocket, SocketStatus(..)
, Family(..), SocketType(..), SockAddr(..) )
import System.Directory ( doesFileExist, removeFile, renameFile )
import System.Environment ( getArgs )
import System.Posix.IO ( dup )
import System.Posix.Process ( getProcessID, forkProcess, executeFile, getProcessStatus )
import System.Posix.Signals ( blockSignals, unblockSignals, fullSignalSet )
import System.Posix.Types ( ProcessID, Fd(Fd) )
import System.Posix.Graceful.Handler
import System.Posix.Graceful.Worker
data GracefulSettings =
GracefulSettings
{ gracefulSettingsListen :: IO Socket
, gracefulSettingsWorkerCount :: Int
, gracefulSettingsSockFile :: FilePath
, gracefulSettingsPidFile :: FilePath
, gracefulSettingsBinary :: FilePath
}
graceful :: GracefulSettings -> GracefulWorker -> IO ()
graceful settings worker = do
quit <- newEmptyMVar
result <- tryIO $ bracket_ (blockSignals fullSignalSet) (unblockSignals fullSignalSet) $ do
esock <- tryRecvSocket settings
sock <- either (const $ gracefulSettingsListen settings) return esock
let launch = launchWorkers (gracefulSettingsWorkerCount settings) $ do
unblockSignals fullSignalSet
defaultHandlers
workerProcess worker sock
pids <- launch >>= newTVarIO
resetHandlers HandlerSettings { handlerSettingsProcessIDs = pids
, handlerSettingsQuitProcess = putMVar quit True
, handlerSettingsLaunchWorkers = launch
, handlerSettingsSpawnProcess = spawnProcess settings sock
}
writeProcessId settings
either throwIO (const $ void $ takeMVar quit) result
tryRecvSocket :: GracefulSettings -> IO (Either IOException Socket)
tryRecvSocket settings =
tryIO $ bracket (socket AF_UNIX Stream 0) close $ \uds -> do
connect uds $ SockAddrUnix $ gracefulSettingsSockFile settings
recvSock uds
writeProcessId :: GracefulSettings -> IO ()
writeProcessId settings =
getProcessID >>= writeFile (gracefulSettingsPidFile settings) . show
clearUnixDomainSocket :: FilePath -> IO ()
clearUnixDomainSocket sockFile = do
exist <- doesFileExist sockFile
when exist $ removeFile sockFile
spawnProcess :: GracefulSettings -> Socket -> IO ()
spawnProcess GracefulSettings { gracefulSettingsSockFile = sockFile
, gracefulSettingsBinary = binary
, gracefulSettingsPidFile = pidFile
} sock = do
exist <- doesFileExist pidFile
when exist $ do
clearUnixDomainSocket sockFile
bracket (socket AF_UNIX Stream 0) close $ \uds -> do
bindSocket uds $ SockAddrUnix sockFile
listen uds 1
args <- getArgs
pid <- forkProcess $ executeFile binary False args Nothing
bracket (accept uds) (close . fst) $ \(s, _) -> sendSock s sock
renameFile pidFile (pidFile ++ ".old")
void $ getProcessStatus True False pid
tryIO :: IO a -> IO (Either IOException a)
tryIO = try
sendSock :: Socket -> Socket -> IO ()
sendSock uds sock = do
Fd fd <- dup $ Fd $ fdSocket sock
sendFd uds fd
let MkSocket _ family socktype protocol _ = sock
void $ send uds $ show (family, socktype, protocol)
recvSock :: Socket -> IO Socket
recvSock uds = do
fd <- recvFd uds
(family, socktype, protocol) <- fmap read $ recv uds 2048
mkSocket fd family socktype protocol Listening
launchWorkers :: Int -> IO () -> IO [ProcessID]
launchWorkers n = replicateM n . forkProcess