-- | Utilities for running a daemon in a local directory
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Gargoyle
  ( Gargoyle (..)
  , withGargoyle
  , gargoyleMain
  ) where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Network.Socket
import System.Directory
import System.Exit
import System.FilePath
import System.IO
import System.IO.Error
import System.Environment
import System.FileLock
import System.Posix.Process
import System.Process

import Debug.Trace

data Gargoyle pid a = Gargoyle
  { Gargoyle pid a -> FilePath
_gargoyle_exec :: FilePath
    -- ^ The path to the executable created with 'gargoyleMain' which will serve as the daemon
    -- monitor process.
  , Gargoyle pid a -> FilePath -> IO ()
_gargoyle_init :: FilePath -> IO ()
    -- ^ The action to run in order to populate the daemon's environment for the first run.
  , Gargoyle pid a -> FilePath -> IO pid
_gargoyle_start :: FilePath -> IO pid
    -- ^ The action to run in order to spin up the daemon on every run. This happens after
    -- '_gargoyle_init' if it also runs.
  , Gargoyle pid a -> pid -> IO ()
_gargoyle_stop :: pid -> IO ()
    -- ^ The action to run when the monitor process detects that no clients are connected anymore.
  , Gargoyle pid a -> FilePath -> IO a
_gargoyle_getInfo :: FilePath -> IO a
    -- ^ Run a command which knows about the working directory of the daemon to collect runtime
    -- information to pass to client code in 'withGargoyle'.
  }

gControlDir :: FilePath -> FilePath
gControlDir :: FilePath -> FilePath
gControlDir = (FilePath -> FilePath -> FilePath
</> FilePath
"control")

gWorkDir :: FilePath -> FilePath
gWorkDir :: FilePath -> FilePath
gWorkDir = (FilePath -> FilePath -> FilePath
</> FilePath
"work")

gOldWorkDir :: FilePath -> FilePath
gOldWorkDir :: FilePath -> FilePath
gOldWorkDir = (FilePath -> FilePath -> FilePath
</> FilePath
"db")

gLockDir :: FilePath -> FilePath
gLockDir :: FilePath -> FilePath
gLockDir = (FilePath -> FilePath -> FilePath
</> FilePath
"lock")

checkThreadedRuntime :: IO ()
checkThreadedRuntime :: IO ()
checkThreadedRuntime = Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
rtsSupportsBoundThreads) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  Handle -> FilePath -> IO ()
hPutStrLn Handle
stderr FilePath
"Gargoyle requires threaded run-time, aborting"
  Bool -> IO () -> IO ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert Bool
rtsSupportsBoundThreads (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) -- throws an AssertionFailed exception

-- | Run an IO action while maintaining a connection to a daemon. The daemon will automatically be
-- stopped when no clients remain. If the daemon has not yet been initialized, it will be.
-- The counterpart of this function is 'gargoyleMain' which should be used to produce an executable
-- that will monitor the daemon's status.
withGargoyle :: Gargoyle pid a -- ^ Description of how to manage the daemon.
             -> FilePath -- ^ The directory where the daemon should be initialized.
             -> (a -> IO b)
                -- ^ Client action which has access to runtime information provided by
                -- the 'Gargoyle'.
             -> IO b
                -- ^ By the time this function returns, the monitor process is aware that the
                -- the client is no longer interested in the daemon.
withGargoyle :: Gargoyle pid a -> FilePath -> (a -> IO b) -> IO b
withGargoyle Gargoyle pid a
g FilePath
daemonDir a -> IO b
b = do
  IO ()
checkThreadedRuntime
  Bool
daemonExists <- FilePath -> IO Bool
doesDirectoryExist FilePath
daemonDir
  if Bool
daemonExists
    then do
    let oldWrk :: FilePath
oldWrk = FilePath -> FilePath
gOldWorkDir FilePath
daemonDir
        wrk :: FilePath
wrk = FilePath -> FilePath
gWorkDir FilePath
daemonDir
    Bool
oldWorkDirExists <- FilePath -> IO Bool
doesDirectoryExist FilePath
oldWrk
    Bool
workDirExists <- FilePath -> IO Bool
doesDirectoryExist FilePath
wrk
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool
oldWorkDirExists Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
workDirExists) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> FilePath -> IO ()
renameDirectory FilePath
oldWrk FilePath
wrk
    else do
    FilePath -> IO ()
createDirectory FilePath
daemonDir
    Gargoyle pid a -> FilePath -> IO ()
forall pid a. Gargoyle pid a -> FilePath -> IO ()
_gargoyle_init Gargoyle pid a
g (FilePath -> FilePath
gWorkDir FilePath
daemonDir)
  Socket
s <- Family -> SocketType -> ProtocolNumber -> IO Socket
socket Family
AF_UNIX SocketType
Stream ProtocolNumber
defaultProtocol
  let acquire :: IO ()
acquire = do
        Either IOError ()
connectResult <- IO () -> IO (Either IOError ())
forall e a. Exception e => IO a -> IO (Either e a)
try (IO () -> IO (Either IOError ()))
-> IO () -> IO (Either IOError ())
forall a b. (a -> b) -> a -> b
$ Socket -> SockAddr -> IO ()
connect Socket
s (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> SockAddr
SockAddrUnix (FilePath -> SockAddr) -> FilePath -> SockAddr
forall a b. (a -> b) -> a -> b
$ FilePath -> FilePath
gControlDir FilePath
daemonDir
        case Either IOError ()
connectResult of
          Right () -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
          Left IOError
e
            | IOError -> Bool
isDoesNotExistError IOError
e -> do
              let monProc :: CreateProcess
monProc = (FilePath -> [FilePath] -> CreateProcess
proc (Gargoyle pid a -> FilePath
forall pid a. Gargoyle pid a -> FilePath
_gargoyle_exec Gargoyle pid a
g) [FilePath
daemonDir])
                    { std_in :: StdStream
std_in = StdStream
CreatePipe
                    , std_out :: StdStream
std_out = StdStream
CreatePipe
                    , std_err :: StdStream
std_err = StdStream
Inherit
                    , close_fds :: Bool
close_fds = Bool
True
                    , new_session :: Bool
new_session = Bool
True
                    }
              (Just Handle
monIn, Just Handle
monOut, Maybe Handle
Nothing, ProcessHandle
monHnd) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess CreateProcess
monProc
              IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkOS (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ IO ExitCode -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ExitCode -> IO ()) -> IO ExitCode -> IO ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
monHnd
              Handle -> IO ()
hClose Handle
monIn
              FilePath
r <- Handle -> IO FilePath
hGetLine Handle
monOut
              case FilePath
r of
                FilePath
"retry" -> do
                  Int -> IO ()
threadDelay Int
500000 -- These are expensive ops so don't try too hard
                  IO ()
acquire -- Try again
                FilePath
"ready" -> IO ()
acquire
                FilePath
_ -> FilePath -> IO ()
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Unexpected gargoyle message from monitor process: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
r
            | Bool
otherwise -> IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e
  IO () -> IO () -> IO b -> IO b
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ IO ()
acquire (Socket -> ShutdownCmd -> IO ()
shutdown Socket
s ShutdownCmd
ShutdownBoth IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Socket -> IO ()
close Socket
s) (IO b -> IO b) -> IO b -> IO b
forall a b. (a -> b) -> a -> b
$
    a -> IO b
b (a -> IO b) -> IO a -> IO b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Gargoyle pid a -> FilePath -> IO a
forall pid a. Gargoyle pid a -> FilePath -> IO a
_gargoyle_getInfo Gargoyle pid a
g (FilePath -> FilePath
gWorkDir FilePath
daemonDir)

-- | Run a local daemon over a domain socket; the daemon will be automatically stopped when
-- no clients remain. This function assumes that the daemon has already been initialized
-- in the specified location. This function should be used as the main function of an executable
-- which will then be invoked by calling 'withGargoyle' in the client code to monitor
-- the daemon's status.
gargoyleMain :: Gargoyle pid a
             -- ^ Description of how to initialize, spin up, and spin down a daemon.
             -> IO () -- ^ Returns only when all clients have disconnected.
gargoyleMain :: Gargoyle pid a -> IO ()
gargoyleMain Gargoyle pid a
g = IO ProcessID -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ProcessID -> IO ()) -> IO ProcessID -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ProcessID
forkProcess (IO () -> IO ProcessID) -> IO () -> IO ProcessID
forall a b. (a -> b) -> a -> b
$ do
  IO ()
checkThreadedRuntime
  [FilePath
daemonDir] <- IO [FilePath]
getArgs IO [FilePath] -> ([FilePath] -> IO [FilePath]) -> IO [FilePath]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    x :: [FilePath]
x@[FilePath
_] -> [FilePath] -> IO [FilePath]
forall (m :: * -> *) a. Monad m => a -> m a
return [FilePath]
x
    [FilePath]
_ -> FilePath -> IO [FilePath]
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail FilePath
"Gargoyle monitor received unexpected number of arguments"
  let lockPath :: FilePath
lockPath = FilePath -> FilePath
gLockDir FilePath
daemonDir
  -- Make sure the lock file is there
  IO () -> (IOError -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (FilePath -> IOMode -> IO Handle
openFile FilePath
lockPath IOMode
WriteMode IO Handle -> (Handle -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Handle -> IO ()
hClose) ((IOError -> IO ()) -> IO ()) -> (IOError -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(IOError
e :: IOException) -> if
    | IOError -> Bool
isAlreadyInUseError IOError
e -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    | IOError -> Bool
isDoesNotExistError IOError
e -> IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e -- this means it's not a file but it exists
    | IOError -> Bool
isPermissionError IOError
e -> IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e -- the daemon directory is in a bad state
    | Bool
otherwise -> IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e
  -- The daemon tries to hold on to the lock file for its lifetime, signaling that it is
  -- accepting connections.
  FileLock
lock <- FilePath -> SharedExclusive -> IO (Maybe FileLock)
tryLockFile FilePath
lockPath SharedExclusive
Exclusive IO (Maybe FileLock)
-> (Maybe FileLock -> IO FileLock) -> IO FileLock
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just FileLock
x -> FileLock -> IO FileLock
forall (m :: * -> *) a. Monad m => a -> m a
return FileLock
x
    Maybe FileLock
Nothing -> do
      FilePath -> IO ()
putStrLn FilePath
"retry"
      Handle -> IO ()
hFlush Handle
stdout
      IO FileLock
forall a. IO a
exitFailure
  -- Clients must maintain a connection to controlSocket to ensure that
  -- the daemon doesn't get shut down
  Socket
controlSocket <- Family -> SocketType -> ProtocolNumber -> IO Socket
socket Family
AF_UNIX SocketType
Stream ProtocolNumber
defaultProtocol
  let socketPath :: FilePath
socketPath = FilePath -> FilePath
gControlDir FilePath
daemonDir
      createSocket :: IO ()
createSocket = do
        Either IOError ()
result <- IO () -> IO (Either IOError ())
forall e a. Exception e => IO a -> IO (Either e a)
try (IO () -> IO (Either IOError ()))
-> IO () -> IO (Either IOError ())
forall a b. (a -> b) -> a -> b
$ Socket -> SockAddr -> IO ()
bind Socket
controlSocket (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> SockAddr
SockAddrUnix FilePath
socketPath
        case Either IOError ()
result of
          Right () -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
          Left IOError
e
            | IOError -> Bool
isAlreadyInUseError IOError
e
            -> do
              -- This is safe because all gargoyle monitors try to take a lock before trying
              -- to mess with the control socket. This avoids race conditions provided that
              -- only gargoyle monitors touch the daemon directory.
              FilePath -> IO ()
removePathForcibly FilePath
socketPath
              FilePath -> IO ()
putStrLn FilePath
"retry"
              Handle -> IO ()
hFlush Handle
stdout
              IO ()
forall a. IO a
exitFailure
            | Bool
otherwise -> IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO IOError
e
  IO () -> (() -> IO ()) -> (() -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO ()
createSocket (\()
_ -> FilePath -> IO ()
removeFile FilePath
socketPath) ((() -> IO ()) -> IO ()) -> (() -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \()
_ -> do
    -- Between bind and listen, the socket will be in a non-accepting state;
    -- this should last a very brief time, so the client should just briefly wait and then retry
    Socket -> Int -> IO ()
listen Socket
controlSocket Int
128
    -- TODO: There is a failure mode here: if an interloper connects and disconnects before
    -- the initial caller connects, the initial caller will fail to connect; instead, we should
    -- start up with an existing connection (possibly a pipe passed in from the parent process) and
    -- with this var set to 1
    MVar Int
numClientsVar <- Int -> IO (MVar Int)
forall a. a -> IO (MVar a)
newMVar (Int
0 :: Int)
    -- When this var is filled, the server will shut down
    MVar ()
shutdownVar <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
    IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkOS (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ IO ThreadId -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      (Socket
s, SockAddr
_) <- Socket -> IO (Socket, SockAddr)
accept Socket
controlSocket
      ThreadId
acceptThread <- IO ThreadId
myThreadId
      --TODO: What happens if we decide we're shutting down here?
      MVar Int -> (Int -> IO Int) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar Int
numClientsVar ((Int -> IO Int) -> IO ()) -> (Int -> IO Int) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
n -> do
        Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a. Enum a => a -> a
succ Int
n
      IO () -> IO ThreadId
forkOS (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
        Handle
h <- Socket -> IOMode -> IO Handle
socketToHandle Socket
s IOMode
ReadMode
        -- Block until we hit EOF; if we successfully read a character that means the client is
        -- in violation of the protocol, so we shut them down too
        (IOError -> Maybe ()) -> IO () -> (() -> IO ()) -> IO ()
forall e b a.
Exception e =>
(e -> Maybe b) -> IO a -> (b -> IO a) -> IO a
catchJust
          (\IOError
e -> if IOError -> Bool
isEOFError IOError
e then () -> Maybe ()
forall a. a -> Maybe a
Just () else IOError
e IOError -> Maybe () -> Maybe ()
forall a b. Show a => a -> b -> b
`traceShow` Maybe ()
forall a. Maybe a
Nothing)
          (Handle -> IO Char
hGetChar Handle
h IO Char -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> FilePath -> IO ()
hPutStrLn Handle
stderr FilePath
"Warning: client sent data over the control socket")
          () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return
        IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          Int
n <- MVar Int -> IO Int
forall a. MVar a -> IO a
takeMVar MVar Int
numClientsVar
          case Int -> Int
forall a. Enum a => a -> a
pred Int
n of
            Int
0 -> do
              Socket -> ShutdownCmd -> IO ()
shutdown Socket
controlSocket ShutdownCmd
ShutdownBoth
              -- We have to explicitly kill the accept thread, because otherwise (sometimes?) 'accept' will begin returning EAGAIN, and ghc will continuously retry it.  This busywait consumes 100% of CPU, prevents the monitor from actually exiting, and leaves the control socket and lock file in a state where another monitor can't be started.
              ThreadId -> IO ()
killThread ThreadId
acceptThread
              MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
shutdownVar ()
            Int
n' -> MVar Int -> Int -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Int
numClientsVar Int
n'
    IO pid -> (pid -> IO ()) -> (pid -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (Gargoyle pid a -> FilePath -> IO pid
forall pid a. Gargoyle pid a -> FilePath -> IO pid
_gargoyle_start Gargoyle pid a
g (FilePath -> FilePath
gWorkDir FilePath
daemonDir)) (Gargoyle pid a -> pid -> IO ()
forall pid a. Gargoyle pid a -> pid -> IO ()
_gargoyle_stop Gargoyle pid a
g) ((pid -> IO ()) -> IO ()) -> (pid -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \pid
_ -> do
      Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
LineBuffering
      FilePath -> IO ()
putStrLn FilePath
"ready" -- Signal to the invoker that we're ready
      MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
shutdownVar
  FileLock -> IO ()
unlockFile FileLock
lock