gargoyle-0.1.2.0: src/Gargoyle.hs

-- | 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_exec :: FilePath
    -- ^ The path to the executable created with 'gargoyleMain' which will serve as the daemon
    -- monitor process.
  , _gargoyle_init :: FilePath -> IO ()
    -- ^ The action to run in order to populate the daemon's environment for the first run.
  , _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_stop :: pid -> IO ()
    -- ^ The action to run when the monitor process detects that no clients are connected anymore.
  , _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 = (</> "control")

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

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

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

checkThreadedRuntime :: IO ()
checkThreadedRuntime = when (not rtsSupportsBoundThreads) $ do
  hPutStrLn stderr "Gargoyle requires threaded run-time, aborting"
  assert rtsSupportsBoundThreads (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 g daemonDir b = do
  checkThreadedRuntime
  daemonExists <- doesDirectoryExist daemonDir
  if daemonExists
    then do
    let oldWrk = gOldWorkDir daemonDir
        wrk = gWorkDir 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 <- try $ connect s $ SockAddrUnix $ gControlDir daemonDir
        case connectResult of
          Right () -> return ()
          Left e
            | isDoesNotExistError e -> do
              let monProc = (proc (_gargoyle_exec g) [daemonDir])
                    { std_in = CreatePipe
                    , std_out = CreatePipe
                    , std_err = Inherit
                    , close_fds = True
                    , new_session = True
                    }
              (Just monIn, Just monOut, Nothing, monHnd) <- createProcess monProc
              void $ forkOS $ void $ waitForProcess monHnd
              hClose monIn
              r <- hGetLine monOut
              case r of
                "retry" -> do
                  threadDelay 500000 -- These are expensive ops so don't try too hard
                  acquire -- Try again
                "ready" -> acquire
                _ -> fail $ "Unexpected gargoyle message from monitor process: " <> r
            | otherwise -> throwIO 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 g = void $ forkProcess $ do
  checkThreadedRuntime
  [daemonDir] <- getArgs >>= \case
    x@[_] -> return x
    _ -> fail "Gargoyle monitor received unexpected number of arguments"
  let lockPath = gLockDir daemonDir
  -- Make sure the lock file is there
  catch (openFile lockPath WriteMode >>= hClose) $ \(e :: IOException) -> if
    | isAlreadyInUseError e -> return ()
    | isDoesNotExistError e -> throwIO e -- this means it's not a file but it exists
    | isPermissionError e -> throwIO e -- the daemon directory is in a bad state
    | otherwise -> throwIO 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 x -> return x
    Nothing -> do
      putStrLn "retry"
      hFlush stdout
      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 = gControlDir daemonDir
      createSocket = do
        result <- try $ bind controlSocket $ SockAddrUnix socketPath
        case result of
          Right () -> return ()
          Left e
            | isAlreadyInUseError 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.
              removePathForcibly socketPath
              putStrLn "retry"
              hFlush stdout
              exitFailure
            | otherwise -> throwIO e
  bracket createSocket (\_ -> removeFile 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
    listen controlSocket 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 <- newMVar (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 $ \n -> do
        return $ succ 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
          (\e -> if isEOFError e then Just () else e `traceShow` Nothing)
          (hGetChar h >> hPutStrLn stderr "Warning: client sent data over the control socket")
          return
        mask_ $ do
          n <- takeMVar numClientsVar
          case pred n of
            0 -> do
              -- If the client side beats us to this, we'll get an error here. This might be specific to Apple Silicon.
              catch (shutdown controlSocket ShutdownBoth) $ \(_ :: IOException) -> 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.
              killThread acceptThread
              putMVar shutdownVar ()
            n' -> putMVar numClientsVar n'
    bracket (_gargoyle_start g (gWorkDir daemonDir)) (_gargoyle_stop g) $ \_ -> do
      hSetBuffering stdout LineBuffering
      putStrLn "ready" -- Signal to the invoker that we're ready
      takeMVar shutdownVar
  unlockFile lock