module Test.Sandwich.Util.Process (
  gracefullyStopProcess
  , gracefullyWaitForProcess
  ) 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


-- | Interrupt a process and wait for it to terminate.
gracefullyStopProcess :: (MonadIO m, MonadLogger m) => ProcessHandle -> Int -> m ()
gracefullyStopProcess :: forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m ()
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 ()
forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m ()
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 = 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 ()) -> m ()
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
_ -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Maybe ExitCode
Nothing -> do
      Maybe Pid
pid <- IO (Maybe Pid) -> m (Maybe Pid)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe Pid) -> m (Maybe Pid))
-> IO (Maybe Pid) -> m (Maybe Pid)
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO (Maybe Pid)
getPid ProcessHandle
p
      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 ()) -> m ()
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
_ -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Maybe ExitCode
Nothing -> m ExitCode -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m ExitCode -> m ()) -> m ExitCode -> m ()
forall a b. (a -> b) -> a -> b
$ do
          Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
warn [i|(#{pid}) Process didn't stop after 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
$ ProcessHandle -> IO ()
terminateProcess ProcessHandle
p
          IO ExitCode -> m ExitCode
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ExitCode -> m ExitCode) -> IO ExitCode -> m ExitCode
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
p