{-# LANGUAGE CPP #-}

module Test.Sandwich.Util.Process (
  gracefullyStopProcess
  , gracefullyWaitForProcess

  , gracefullyStopProcess'
  , gracefullyWaitForProcess'

  , StopProcessResult(..)
  ) where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Logger
import Control.Retry
import Data.Maybe
import Data.String.Interpolate
import System.Process
import Test.Sandwich.Logging

#ifndef mingw32_HOST_OS
import System.Posix.Signals (signalProcess, sigKILL)
#endif


data StopProcessResult =
  StoppedByItself
  | StoppedAfterInterrupt
  | StoppedAfterTerminate
  | StoppedAfterKill
  | FailedToStop
  deriving (Int -> StopProcessResult -> ShowS
[StopProcessResult] -> ShowS
StopProcessResult -> String
(Int -> StopProcessResult -> ShowS)
-> (StopProcessResult -> String)
-> ([StopProcessResult] -> ShowS)
-> Show StopProcessResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> StopProcessResult -> ShowS
showsPrec :: Int -> StopProcessResult -> ShowS
$cshow :: StopProcessResult -> String
show :: StopProcessResult -> String
$cshowList :: [StopProcessResult] -> ShowS
showList :: [StopProcessResult] -> ShowS
Show, StopProcessResult -> StopProcessResult -> Bool
(StopProcessResult -> StopProcessResult -> Bool)
-> (StopProcessResult -> StopProcessResult -> Bool)
-> Eq StopProcessResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: StopProcessResult -> StopProcessResult -> Bool
== :: StopProcessResult -> StopProcessResult -> Bool
$c/= :: StopProcessResult -> StopProcessResult -> Bool
/= :: StopProcessResult -> StopProcessResult -> Bool
Eq, Eq StopProcessResult
Eq StopProcessResult =>
(StopProcessResult -> StopProcessResult -> Ordering)
-> (StopProcessResult -> StopProcessResult -> Bool)
-> (StopProcessResult -> StopProcessResult -> Bool)
-> (StopProcessResult -> StopProcessResult -> Bool)
-> (StopProcessResult -> StopProcessResult -> Bool)
-> (StopProcessResult -> StopProcessResult -> StopProcessResult)
-> (StopProcessResult -> StopProcessResult -> StopProcessResult)
-> Ord StopProcessResult
StopProcessResult -> StopProcessResult -> Bool
StopProcessResult -> StopProcessResult -> Ordering
StopProcessResult -> StopProcessResult -> StopProcessResult
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: StopProcessResult -> StopProcessResult -> Ordering
compare :: StopProcessResult -> StopProcessResult -> Ordering
$c< :: StopProcessResult -> StopProcessResult -> Bool
< :: StopProcessResult -> StopProcessResult -> Bool
$c<= :: StopProcessResult -> StopProcessResult -> Bool
<= :: StopProcessResult -> StopProcessResult -> Bool
$c> :: StopProcessResult -> StopProcessResult -> Bool
> :: StopProcessResult -> StopProcessResult -> Bool
$c>= :: StopProcessResult -> StopProcessResult -> Bool
>= :: StopProcessResult -> StopProcessResult -> Bool
$cmax :: StopProcessResult -> StopProcessResult -> StopProcessResult
max :: StopProcessResult -> StopProcessResult -> StopProcessResult
$cmin :: StopProcessResult -> StopProcessResult -> StopProcessResult
min :: StopProcessResult -> StopProcessResult -> StopProcessResult
Ord, Int -> StopProcessResult
StopProcessResult -> Int
StopProcessResult -> [StopProcessResult]
StopProcessResult -> StopProcessResult
StopProcessResult -> StopProcessResult -> [StopProcessResult]
StopProcessResult
-> StopProcessResult -> StopProcessResult -> [StopProcessResult]
(StopProcessResult -> StopProcessResult)
-> (StopProcessResult -> StopProcessResult)
-> (Int -> StopProcessResult)
-> (StopProcessResult -> Int)
-> (StopProcessResult -> [StopProcessResult])
-> (StopProcessResult -> StopProcessResult -> [StopProcessResult])
-> (StopProcessResult -> StopProcessResult -> [StopProcessResult])
-> (StopProcessResult
    -> StopProcessResult -> StopProcessResult -> [StopProcessResult])
-> Enum StopProcessResult
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: StopProcessResult -> StopProcessResult
succ :: StopProcessResult -> StopProcessResult
$cpred :: StopProcessResult -> StopProcessResult
pred :: StopProcessResult -> StopProcessResult
$ctoEnum :: Int -> StopProcessResult
toEnum :: Int -> StopProcessResult
$cfromEnum :: StopProcessResult -> Int
fromEnum :: StopProcessResult -> Int
$cenumFrom :: StopProcessResult -> [StopProcessResult]
enumFrom :: StopProcessResult -> [StopProcessResult]
$cenumFromThen :: StopProcessResult -> StopProcessResult -> [StopProcessResult]
enumFromThen :: StopProcessResult -> StopProcessResult -> [StopProcessResult]
$cenumFromTo :: StopProcessResult -> StopProcessResult -> [StopProcessResult]
enumFromTo :: StopProcessResult -> StopProcessResult -> [StopProcessResult]
$cenumFromThenTo :: StopProcessResult
-> StopProcessResult -> StopProcessResult -> [StopProcessResult]
enumFromThenTo :: StopProcessResult
-> StopProcessResult -> StopProcessResult -> [StopProcessResult]
Enum)

-- | Interrupt a process and wait for it to terminate, returning a 'StopProcessResult'.
gracefullyStopProcess :: (MonadIO m, MonadLogger m) => ProcessHandle -> Int -> m ()
gracefullyStopProcess :: forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m ()
gracefullyStopProcess ProcessHandle
p Int
gracePeriodUs = m StopProcessResult -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m StopProcessResult -> m ()) -> m StopProcessResult -> m ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> Int -> m StopProcessResult
forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m StopProcessResult
gracefullyStopProcess' ProcessHandle
p Int
gracePeriodUs

-- | Interrupt a process and wait for it to terminate.
gracefullyStopProcess' :: (MonadIO m, MonadLogger m) => ProcessHandle -> Int -> m StopProcessResult
gracefullyStopProcess' :: forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m StopProcessResult
gracefullyStopProcess' ProcessHandle
p Int
gracePeriodUs = do
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO ()
interruptProcessGroupOf ProcessHandle
p
  ProcessHandle -> Int -> m StopProcessResult
forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m StopProcessResult
gracefullyWaitForProcess' ProcessHandle
p Int
gracePeriodUs

-- | Wait for a process to terminate. If it doesn't terminate within 'gracePeriodUs' microseconds,
-- send it an interrupt signal and wait for another 'gracePeriodUs' microseconds.
-- After this time elapses send a terminate signal and wait for the process to die.
gracefullyWaitForProcess :: (MonadIO m, MonadLogger m) => ProcessHandle -> Int -> m ()
gracefullyWaitForProcess :: forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m ()
gracefullyWaitForProcess ProcessHandle
p Int
gracePeriodUs = m StopProcessResult -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m StopProcessResult -> m ()) -> m StopProcessResult -> m ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> Int -> m StopProcessResult
forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m StopProcessResult
gracefullyWaitForProcess' ProcessHandle
p Int
gracePeriodUs

-- | Wait for a process to terminate. If it doesn't terminate within 'gracePeriodUs' microseconds,
-- send it an interrupt signal and wait for another 'gracePeriodUs' microseconds.
-- After this time elapses send a terminate signal and wait for the process to die.
gracefullyWaitForProcess' :: (MonadIO m, MonadLogger m) => ProcessHandle -> Int -> m StopProcessResult
gracefullyWaitForProcess' :: forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m StopProcessResult
gracefullyWaitForProcess' ProcessHandle
p Int
gracePeriodUs = do
  let waitForExit :: m (Maybe ExitCode)
waitForExit = do
        let policy :: RetryPolicyM m
policy = Int -> RetryPolicyM m -> RetryPolicyM m
forall (m :: * -> *).
Monad m =>
Int -> RetryPolicyM m -> RetryPolicyM m
limitRetriesByCumulativeDelay Int
gracePeriodUs (RetryPolicyM m -> RetryPolicyM m)
-> RetryPolicyM m -> RetryPolicyM m
forall a b. (a -> b) -> a -> b
$ Int -> RetryPolicyM m -> RetryPolicyM m
forall (m :: * -> *).
Monad m =>
Int -> RetryPolicyM m -> RetryPolicyM m
capDelay Int
200_000 (RetryPolicyM m -> RetryPolicyM m)
-> RetryPolicyM m -> RetryPolicyM m
forall a b. (a -> b) -> a -> b
$ Int -> RetryPolicyM m
forall (m :: * -> *). Monad m => Int -> RetryPolicyM m
exponentialBackoff Int
1_000
        RetryPolicyM m
-> (RetryStatus -> Maybe ExitCode -> m Bool)
-> (RetryStatus -> m (Maybe ExitCode))
-> m (Maybe ExitCode)
forall (m :: * -> *) b.
MonadIO m =>
RetryPolicyM m
-> (RetryStatus -> b -> m Bool) -> (RetryStatus -> m b) -> m b
retrying RetryPolicyM m
policy (\RetryStatus
_ Maybe ExitCode
x -> Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Maybe ExitCode -> Bool
forall a. Maybe a -> Bool
isNothing Maybe ExitCode
x) ((RetryStatus -> m (Maybe ExitCode)) -> m (Maybe ExitCode))
-> (RetryStatus -> m (Maybe ExitCode)) -> m (Maybe ExitCode)
forall a b. (a -> b) -> a -> b
$ \RetryStatus
_ -> do
          IO (Maybe ExitCode) -> m (Maybe ExitCode)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ExitCode) -> m (Maybe ExitCode))
-> IO (Maybe ExitCode) -> m (Maybe ExitCode)
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO (Maybe ExitCode)
getProcessExitCode ProcessHandle
p

  m (Maybe ExitCode)
waitForExit m (Maybe ExitCode)
-> (Maybe ExitCode -> m StopProcessResult) -> m StopProcessResult
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just ExitCode
_ -> StopProcessResult -> m StopProcessResult
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return StopProcessResult
StoppedByItself
    Maybe ExitCode
Nothing -> do
      IO (Maybe Pid) -> m (Maybe Pid)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ProcessHandle -> IO (Maybe Pid)
getPid ProcessHandle
p) m (Maybe Pid)
-> (Maybe Pid -> m StopProcessResult) -> m StopProcessResult
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe Pid
Nothing -> StopProcessResult -> m StopProcessResult
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return StopProcessResult
StoppedByItself
        Just Pid
pid -> do
          Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
warn [i|(#{pid}) Process didn't stop after #{gracePeriodUs}us; trying to interrupt|]

          IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO ()
interruptProcessGroupOf ProcessHandle
p
          m (Maybe ExitCode)
waitForExit m (Maybe ExitCode)
-> (Maybe ExitCode -> m StopProcessResult) -> m StopProcessResult
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Just ExitCode
_ -> StopProcessResult -> m StopProcessResult
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return StopProcessResult
StoppedAfterInterrupt
            Maybe ExitCode
Nothing -> do
              Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
warn [i|(#{pid}) Process didn't stop after another sigINT and a further #{gracePeriodUs}us; going to terminate|]
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO ()
terminateProcess ProcessHandle
p

#ifdef mingw32_HOST_OS
              waitForExit >>= \case
                Just _ -> return StoppedAfterKill
                Nothing -> return FailedToStop
#else
              m (Maybe ExitCode)
waitForExit m (Maybe ExitCode)
-> (Maybe ExitCode -> m StopProcessResult) -> m StopProcessResult
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                Just ExitCode
_ -> StopProcessResult -> m StopProcessResult
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return StopProcessResult
StoppedAfterTerminate
                Maybe ExitCode
Nothing -> do
                  Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
warn [i|(#{pid}) Process didn't stop after sigTERM and a further #{gracePeriodUs}us; going to kill|]
                  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Signal -> Pid -> IO ()
signalProcess Signal
sigKILL Pid
pid
                  m (Maybe ExitCode)
waitForExit m (Maybe ExitCode)
-> (Maybe ExitCode -> m StopProcessResult) -> m StopProcessResult
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                    Just ExitCode
_ -> StopProcessResult -> m StopProcessResult
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return StopProcessResult
StoppedAfterKill
                    Maybe ExitCode
Nothing -> StopProcessResult -> m StopProcessResult
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return StopProcessResult
FailedToStop
#endif