    -- |
    -- Module      : System.Posix.Graceful
    -- Copyright   : 2013 Noriyuki OHKAWA
    -- License     : BSD3
    --
    -- Maintainer  : n.ohkawa@gmail.com
    -- Stability   : experimental
    -- Portability : unknown
    --
    -- Provides function to make process graceful.
    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

    -- | Server settings
    --
    -- Since 0.1.0.0
    --
    data GracefulSettings =
        GracefulSettings
        { gracefulSettingsListen :: IO Socket -- ^ Listen socket
        , gracefulSettingsWorkerCount :: Int -- ^ Prefork worker count
        , gracefulSettingsSockFile :: FilePath -- ^ Unix domain socket file
        , gracefulSettingsPidFile :: FilePath -- ^ The file to which the server records the process id
        , gracefulSettingsBinary :: FilePath -- ^ The binary file to upgrade
        }

    -- | Make server application enable shutdown/restart gracefully
    --
    -- Since 0.1.0.0
    --
    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