-- | 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
  { forall pid a. Gargoyle pid a -> FilePath
_gargoyle_exec :: FilePath
    -- ^ The path to the executable created with 'gargoyleMain' which will serve as the daemon
    -- monitor process.
  , forall pid a. 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.
  , forall pid a. 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.
  , forall pid a. Gargoyle pid a -> pid -> IO ()
_gargoyle_stop :: pid -> IO ()
    -- ^ The action to run when the monitor process detects that no clients are connected anymore.
  , forall pid a. 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 a. a -> IO a
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 :: forall pid a b. Gargoyle pid a -> FilePath -> (a -> IO b) -> IO b
withGargoyle Gargoyle pid a
g FilePath
daemonDir a -> IO b
b = do
  IO ()
checkThreadedRuntime
  daemonExists <- FilePath -> IO Bool
doesDirectoryExist FilePath
daemonDir
  if daemonExists
    then do
    let oldWrk = FilePath -> FilePath
gOldWorkDir FilePath
daemonDir
        wrk = FilePath -> FilePath
gWorkDir FilePath
daemonDir
    oldWorkDirExists <- doesDirectoryExist oldWrk
    workDirExists <- doesDirectoryExist wrk
    when (oldWorkDirExists && not workDirExists) $ renameDirectory oldWrk wrk
    else do
    createDirectory daemonDir
    _gargoyle_init g (gWorkDir daemonDir)
  s <- socket AF_UNIX Stream defaultProtocol
  let acquire = do
        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 connectResult of
          Right () -> () -> IO ()
forall a. a -> IO a
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 = CreatePipe
                    , std_out = CreatePipe
                    , std_err = Inherit
                    , close_fds = True
                    , new_session = True
                    }
              (Just monIn, Just monOut, Nothing, monHnd) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess CreateProcess
monProc
              void $ forkOS $ void $ waitForProcess monHnd
              hClose monIn
              r <- hGetLine monOut
              case 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 a. FilePath -> IO a
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. (?callStack::CallStack, Exception e) => e -> IO a
throwIO IOError
e
  bracket_ acquire (shutdown s ShutdownBoth >> close s) $
    b =<< _gargoyle_getInfo g (gWorkDir 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 :: forall pid a. 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
  [daemonDir] <- IO [FilePath]
getArgs IO [FilePath] -> ([FilePath] -> IO [FilePath]) -> IO [FilePath]
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    x :: [FilePath]
x@[FilePath
_] -> [FilePath] -> IO [FilePath]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return [FilePath]
x
    [FilePath]
_ -> FilePath -> IO [FilePath]
forall a. FilePath -> IO a
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail FilePath
"Gargoyle monitor received unexpected number of arguments"
  let lockPath = FilePath -> FilePath
gLockDir FilePath
daemonDir
  -- Make sure the lock file is there
  catch (openFile lockPath WriteMode >>= hClose) $ \(IOError
e :: IOException) -> if
    | IOError -> Bool
isAlreadyInUseError IOError
e -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    | IOError -> Bool
isDoesNotExistError IOError
e -> IOError -> IO ()
forall e a. (?callStack::CallStack, 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. (?callStack::CallStack, Exception e) => e -> IO a
throwIO IOError
e -- the daemon directory is in a bad state
    | Bool
otherwise -> IOError -> IO ()
forall e a. (?callStack::CallStack, 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.
  lock <- tryLockFile lockPath Exclusive >>= \case
    Just FileLock
x -> FileLock -> IO FileLock
forall a. a -> IO a
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
  controlSocket <- socket AF_UNIX Stream defaultProtocol
  let socketPath = FilePath -> FilePath
gControlDir FilePath
daemonDir
      createSocket = do
        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 result of
          Right () -> () -> IO ()
forall a. a -> IO a
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. (?callStack::CallStack, Exception e) => e -> IO a
throwIO IOError
e
  bracket createSocket (\()
_ -> FilePath -> IO ()
removeFile FilePath
socketPath) $ \()
_ -> 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
    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
    shutdownVar <- newEmptyMVar
    void $ forkOS $ forever $ do
      (s, _) <- accept controlSocket
      acceptThread <- myThreadId
      --TODO: What happens if we decide we're shutting down here?
      modifyMVar_ numClientsVar $ \Int
n -> do
        Int -> IO Int
forall a. a -> IO a
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
      forkOS $ do
        h <- socketToHandle s 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
        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)
          (hGetChar h >> hPutStrLn stderr "Warning: client sent data over the control socket")
          return
        mask_ $ do
          n <- takeMVar numClientsVar
          case pred n of
            Int
0 -> do
              -- If the client side beats us to this, we'll get an error here. This might be specific to Apple Silicon.
              IO () -> (IOError -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (Socket -> ShutdownCmd -> IO ()
shutdown Socket
controlSocket ShutdownCmd
ShutdownBoth) ((IOError -> IO ()) -> IO ()) -> (IOError -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(IOError
_ :: IOException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              -- 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'
    bracket (_gargoyle_start g (gWorkDir daemonDir)) (_gargoyle_stop g) $ \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
  unlockFile lock