content: {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE ForeignFunctionInterface #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  module Keter.Conduit.Process.Unix
      ( -- * Process tracking
        -- $processTracker
  
        -- ** Types
        ProcessTracker
        -- ** Functions
      , initProcessTracker
  
        -- * Monitored process
      , MonitoredProcess
      , monitorProcess
      , terminateMonitoredProcess
      , printStatus
      ) where
  
  import Control.Applicative (pure, (<$>), (<*>))
  import Control.Arrow ((***))
  import Control.Concurrent (forkIO, threadDelay)
  import Control.Concurrent.MVar
         ( MVar
         , modifyMVar
         , modifyMVar_
         , newEmptyMVar
         , newMVar
         , putMVar
         , readMVar
         , swapMVar
         , takeMVar
         , tryReadMVar
         )
  import Control.Exception
         ( Exception
         , SomeException
         , bracketOnError
         , finally
         , handle
         , mask_
         , throwIO
         , try
         )
  import Control.Monad (void, when)
  import Control.Monad.IO.Unlift
  import Control.Monad.Logger
  import Data.ByteString (ByteString)
  import Data.ByteString.Char8 qualified as S8
  import Data.Conduit (ConduitM, runConduit, (.|))
  import Data.Conduit.Binary (sinkHandle, sourceHandle)
  import Data.Conduit.List qualified as CL
  import Data.IORef (IORef, newIORef, readIORef, writeIORef)
  import Data.Monoid ((<>))
  import Data.Text (Text, pack)
  import Data.Text.Encoding (decodeUtf8)
  import Data.Time (diffUTCTime, getCurrentTime)
  import Data.Typeable (Typeable)
  import Foreign.C.Types
  import Prelude
         ( Bool(..)
         , Either(..)
         , IO
         , Maybe(..)
         , Monad(..)
         , Show
         , String
         , const
         , error
         , map
         , maybe
         , show
         , ($!)
         , ($)
         , (*)
         , (<)
         , (==)
         )
  import System.Exit (ExitCode)
  import System.IO (hClose)
  import System.Posix.IO.ByteString (closeFd, createPipe, fdToHandle)
  import System.Posix.Signals (sigKILL, signalProcess)
  import System.Posix.Types (CPid(..))
  import System.Process
         ( CmdSpec(..)
         , CreateProcess(..)
         , StdStream(..)
         , createProcess
         , getPid
         , terminateProcess
         , waitForProcess
         )
  import System.Process.Internals (ProcessHandle(..), ProcessHandle__(..))
  
  processHandleMVar :: ProcessHandle -> MVar ProcessHandle__
  processHandleMVar (ProcessHandle m _ _) = m
  
  withProcessHandle_ ::
       ProcessHandle
    -> (ProcessHandle__ -> IO ProcessHandle__)
    -> IO ()
  withProcessHandle_ ph = modifyMVar_ (processHandleMVar ph)
  
  -- | Kill a process by sending it the KILL (9) signal.
  --
  -- Since 0.1.0
  killProcess :: ProcessHandle -> IO ()
  killProcess ph = withProcessHandle_ ph $ \p_ -> case p_ of
    ClosedHandle _ -> return p_
    OpenHandle h -> do
      signalProcess sigKILL h
      return p_
    _ -> error "Not implemented"
  
  ignoreExceptions :: IO () -> IO ()
  ignoreExceptions = handle (\(_ :: SomeException) -> return ())
  
  -- $processTracker
  --
  -- Ensure that child processes are killed, regardless of how the parent process exits.
  --
  -- The technique used here is:
  --
  -- * Create a pipe.
  --
  -- * Fork a new child process that listens on the pipe.
  --
  -- * In the current process, send updates about processes that should be auto-killed.
  --
  -- * When the parent process dies, listening on the pipe in the child process will get an EOF.
  --
  -- * When the child process receives that EOF, it kills all processes it was told to auto-kill.
  --
  -- This code was originally written for Keter, but was moved to unix-process
  -- conduit in the 0.2.1 release.
  
  foreign import ccall unsafe "launch_process_tracker"
      c_launch_process_tracker :: IO CInt
  
  foreign import ccall unsafe "track_process"
      c_track_process :: ProcessTracker -> CPid -> CInt -> IO ()
  
  -- | Represents the child process which handles process cleanup.
  --
  -- Since 0.2.1
  newtype ProcessTracker = ProcessTracker CInt
  
  -- | Represents a child process which is currently being tracked by the cleanup
  -- child process.
  --
  -- Since 0.2.1
  data TrackedProcess = TrackedProcess !ProcessTracker !(IORef MaybePid) !(IO ExitCode)
  
  data MaybePid = NoPid | Pid !CPid
  
  -- | Fork off the child cleanup process.
  --
  -- This will ideally only be run once for your entire application.
  --
  -- Since 0.2.1
  initProcessTracker :: IO ProcessTracker
  initProcessTracker = do
      i <- c_launch_process_tracker
      if i == -1
          then throwIO CannotLaunchProcessTracker
          else return $! ProcessTracker i
  
  -- | Since 0.2.1
  data ProcessTrackerException = CannotLaunchProcessTracker
      deriving (Show, Typeable)
  instance Exception ProcessTrackerException
  
  -- | Begin tracking the given process. If the 'ProcessHandle' refers to a
  -- closed process, no tracking will occur. If the process is closed, then it
  -- will be untracked automatically.
  --
  -- Note that you /must/ compile your program with @-threaded@; see
  -- 'waitForProcess'.
  --
  -- Since 0.2.1
  trackProcess :: ProcessTracker -> ProcessHandle -> IO TrackedProcess
  trackProcess pt ph = mask_ $ do
      mpid <- readMVar $ processHandleMVar ph
      mpid' <- case mpid of
          ClosedHandle{} -> return NoPid
          OpenHandle pid -> do
              c_track_process pt pid 1
              return $ Pid pid
          _ -> error "Not implemented"
      ipid <- newIORef mpid'
      baton <- newEmptyMVar
      let tp = TrackedProcess pt ipid (takeMVar baton)
      case mpid' of
          NoPid -> return ()
          Pid _ -> void $ forkIO $ do
              waitForProcess ph >>= putMVar baton
              untrackProcess tp
      return $! tp
  
  -- | Explicitly remove the given process from the tracked process list in the
  -- cleanup process.
  --
  -- Since 0.2.1
  untrackProcess :: TrackedProcess -> IO ()
  untrackProcess (TrackedProcess pt ipid _) = mask_ $ do
      mpid <- readIORef ipid
      case mpid of
          NoPid -> return ()
          Pid pid -> do
              c_track_process pt pid 0
              writeIORef ipid NoPid
  
  -- | Fork and execute a subprocess, sending stdout and stderr to the specified
  -- rotating log.
  --
  -- Since 0.2.1
  forkExecuteLog :: ByteString -- ^ command
                 -> [ByteString] -- ^ args
                 -> Maybe [(ByteString, ByteString)] -- ^ environment
                 -> Maybe ByteString -- ^ working directory
                 -> Maybe (ConduitM () ByteString IO ()) -- ^ stdin
                 -> (ByteString -> IO ()) -- ^ both stdout and stderr will be sent to this location
                 -> IO ProcessHandle
  forkExecuteLog cmd args menv mwdir mstdin log = bracketOnError
      setupPipe
      cleanupPipes
      usePipes
    where
      setupPipe = bracketOnError
          createPipe
          (\(x, y) -> closeFd x `finally` closeFd y)
          (\(x, y) -> (,) <$> fdToHandle x <*> fdToHandle y)
      cleanupPipes (x, y) = hClose x `finally` hClose y
  
      usePipes pipes@(readerH, writerH) = do
          (min, _, _, ph) <- createProcess CreateProcess
              { cmdspec = RawCommand (S8.unpack cmd) (map S8.unpack args)
              , cwd = S8.unpack <$> mwdir
              , env = map (S8.unpack *** S8.unpack) <$> menv
              , std_in = maybe Inherit (const CreatePipe) mstdin
              , std_out = UseHandle writerH
              , std_err = UseHandle writerH
              , close_fds = True
              , create_group = True
              , use_process_jobs = False
              , delegate_ctlc = False
              , detach_console = True
              , create_new_console = False
              , new_session = True
              , child_group = Nothing
              , child_user = Nothing
              }
          ignoreExceptions $ addAttachMessage pipes ph
          void $ forkIO $ ignoreExceptions $
              runConduit (sourceHandle readerH .| CL.mapM_ log) `finally` hClose readerH
          case (min, mstdin) of
              (Just h, Just source) -> void $ forkIO $ ignoreExceptions $
                  runConduit (source .| sinkHandle h) `finally` hClose h
              (Nothing, Nothing) -> return ()
              _ -> error "Invariant violated: Data.Conduit.Process.Unix.forkExecuteLog"
          return ph
  
      addAttachMessage pipes ph = withProcessHandle_ ph $ \p_ -> do
          now <- getCurrentTime
          case p_ of
              ClosedHandle ec -> do
                  log $ S8.concat
                      [ "\n\n"
                      , S8.pack $ show now
                      , ": Process immediately died with exit code "
                      , S8.pack $ show ec
                      , "\n\n"
                      ]
                  cleanupPipes pipes
              OpenHandle h -> do
                  log $ S8.concat
                      [ "\n\n"
                      , S8.pack $ show now
                      , ": Attached new process "
                      , S8.pack $ show h
                      , "\n\n"
                      ]
              _ -> error "Not implemented"
          return p_
  
  data Status = NeedsRestart | NoRestart | Running ProcessHandle
  
  -- | Run the given command, restarting if the process dies.
  monitorProcess
      :: (MonadUnliftIO m, MonadLogger m)
      => ProcessTracker
      -> Maybe S8.ByteString -- ^ setuid
      -> S8.ByteString -- ^ executable
      -> S8.ByteString -- ^ working directory
      -> [S8.ByteString] -- ^ command line parameter
      -> [(S8.ByteString, S8.ByteString)] -- ^ environment
      -> (ByteString -> IO ())
      -> (ExitCode -> IO Bool) -- ^ should we restart?
      -> m MonitoredProcess
  monitorProcess processTracker msetuid exec dir args env' rlog shouldRestart =
      withRunInIO $ \rio -> do
          mstatus <- newMVar NeedsRestart
          let loop mlast = do
                next <- modifyMVar mstatus $ \case
                  NoRestart -> return (NoRestart, return ())
                  _ -> do
                      now <- getCurrentTime
                      case mlast of
                          Just last | diffUTCTime now last < 5 -> do
                              rio $ $logWarn $ "Process restarting too quickly, waiting before trying again: " <> decodeUtf8 exec
                              threadDelay $ 5 * 1000 * 1000
                          _ -> return ()
                      let (cmd, args') =
                              case msetuid of
                                  Nothing -> (exec, args)
                                  Just setuid -> ("sudo", "-E" : "-u" : setuid : "--" : exec : args)
                      res <- try $ forkExecuteLog
                          cmd
                          args'
                          (Just env')
                          (Just dir)
                          (Just $ return ())
                          rlog
                      case res of
                          Left e -> do
                              rio $ $logError $ "Data.Conduit.Process.Unix.monitorProcess: " <> pack (show (e :: SomeException))
                              return (NeedsRestart, return ())
                          Right pid -> do
                              rio $ $logInfo $ "Process created: " <> decodeUtf8 exec
                              return (Running pid, do
                                  TrackedProcess _ _ wait <- trackProcess processTracker pid
                                  ec <- wait
                                  shouldRestart' <- shouldRestart ec
                                  when shouldRestart' $ loop (Just now))
                next
          _ <- forkIO $ loop Nothing
          return $ MonitoredProcess mstatus
  {-# ANN monitorProcess ("HLint: ignore Use join" :: String) #-}
  
  -- | Abstract type containing information on a process which will be restarted.
  newtype MonitoredProcess = MonitoredProcess (MVar Status)
  
  printStatus :: MonitoredProcess -> IO Text
  printStatus (MonitoredProcess mstatus) = do
    mStatus <- tryReadMVar mstatus
    case mStatus of
      Nothing -> pure "no status set process"
      Just NeedsRestart -> pure "needs-restart process"
      Just NoRestart -> pure "no-restart process"
      Just (Running running) -> do
        x <- getPid running
        case x of
          Just y -> pure ("running process '" <> pack (show y) <> "'")
          Nothing -> pure "just closed process"
  
  -- | Terminate the process and prevent it from being restarted.
  terminateMonitoredProcess :: MonitoredProcess -> IO ()
  terminateMonitoredProcess (MonitoredProcess mstatus) = do
      status <- swapMVar mstatus NoRestart
      case status of
          Running pid -> do
              terminateProcess pid
              threadDelay 1000000
              killProcess pid
          _ -> return ()
  