timeout.hs 5.66 KB
Newer Older
1
{-# OPTIONS -cpp #-}
dterei's avatar
dterei committed
2
module Main where
3 4 5

import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.MVar (putMVar, takeMVar, newEmptyMVar)
dterei's avatar
dterei committed
6
import Control.Monad
7
import Control.Exception
8 9
import Data.Maybe (isNothing)
import System.Environment (getArgs)
10
import System.Exit
11
import System.IO (hPutStrLn, stderr)
12

13
#if !defined(mingw32_HOST_OS)
14 15
import System.Posix hiding (killProcess)
import System.IO.Error hiding (try,catch)
16 17
#endif

18
#if defined(mingw32_HOST_OS)
19
import System.Process
20 21 22 23 24
import WinCBindings
import Foreign
import System.Win32.DebugApi
import System.Win32.Types
#endif
25

26
main :: IO ()
27 28 29
main = do
  args <- getArgs
  case args of
Ian Lynagh's avatar
Ian Lynagh committed
30 31 32 33 34 35
      [secs,cmd] ->
          case reads secs of
          [(secs', "")] -> run secs' cmd
          _ -> die ("Can't parse " ++ show secs ++ " as a number of seconds")
      _ -> die ("Bad arguments " ++ show args)

36 37 38
run :: Int -> String -> IO ()
#if !defined(mingw32_HOST_OS)
run secs cmd = do
39 40 41
        m <- newEmptyMVar
        mp <- newEmptyMVar
        installHandler sigINT (Catch (putMVar m Nothing)) Nothing
42 43 44 45 46 47
        forkIO $ do threadDelay (secs * 1000000)
                    putMVar m Nothing
        forkIO $ do ei <- try $ do pid <- systemSession cmd
                                   return pid
                    putMVar mp ei
                    case ei of
48
                       Left _ -> return ()
49 50 51
                       Right pid -> do
                           r <- getProcessStatus True False pid
                           putMVar m r
52 53 54 55 56
        ei_pid_ph <- takeMVar mp
        case ei_pid_ph of
            Left e -> do hPutStrLn stderr
                                   ("Timeout:\n" ++ show (e :: IOException))
                         exitWith (ExitFailure 98)
57
            Right pid -> do
58 59 60
                r <- takeMVar m
                case r of
                  Nothing -> do
61
                        killProcess pid
62
                        exitWith (ExitFailure 99)
63 64 65
                  Just (Exited r) -> exitWith r
                  Just (Terminated s) -> raiseSignal s
                  Just _ -> exitWith (ExitFailure 1)
66 67 68 69 70 71 72 73 74 75 76

systemSession cmd =
 forkProcess $ do
   createSession
   executeFile "/bin/sh" False ["-c", cmd] Nothing
   -- need to use exec() directly here, rather than something like
   -- System.Process.system, because we are in a forked child and some
   -- pthread libraries get all upset if you start doing certain
   -- things in a forked child of a pthread process, such as forking
   -- more threads.

77
killProcess pid = do
78
  ignoreIOExceptions (signalProcessGroup sigTERM pid)
79 80 81 82 83
  checkReallyDead 10
  where
    checkReallyDead 0 = hPutStrLn stderr "checkReallyDead: Giving up"
    checkReallyDead (n+1) =
      do threadDelay (3*100000) -- 3/10 sec
84 85 86 87 88 89 90 91
         m <- tryJust (guard . isDoesNotExistError) $
                 getProcessStatus False False pid
         case m of
            Right Nothing -> return ()
            Left _ -> return ()
            _ -> do
              ignoreIOExceptions (signalProcessGroup sigKILL pid)
              checkReallyDead n
92

93 94 95
ignoreIOExceptions :: IO () -> IO ()
ignoreIOExceptions io = io `catch` ((\_ -> return ()) :: IOException -> IO ())

96
#else
97
run secs cmd =
98 99 100 101
    let escape '\\' = "\\\\"
        escape '"'  = "\\\""
        escape c    = [c]
        cmd' = "sh -c \"" ++ concatMap escape cmd ++ "\"" in
102 103
    alloca $ \p_startupinfo ->
    alloca $ \p_pi ->
104
    withTString cmd' $ \cmd'' ->
105
    do job <- createJobObjectW nullPtr nullPtr
106 107 108 109 110 111
       b_info <- setJobParameters job
       unless b_info $ errorWin "setJobParameters"

       ioPort <- createCompletionPort job
       when (ioPort == nullPtr) $ errorWin "createCompletionPort, cannot continue."

112
       -- We're explicitly turning off handle inheritance to prevent misc handles
Joachim Breitner's avatar
Joachim Breitner committed
113 114
       -- from being inherited by the child. Notable we don't want the I/O Completion
       -- Ports and Job handles to be inherited. So we mark them as non-inheritable.
115 116
       setHandleInformation job    cHANDLE_FLAG_INHERIT 0
       setHandleInformation ioPort cHANDLE_FLAG_INHERIT 0
117 118 119

       -- Now create the process suspended so we can add it to the job and then resume.
       -- This is so we don't miss any events on the receiving end of the I/O port.
120
       let creationflags = cCREATE_SUSPENDED
121
       b <- createProcessW nullPtr cmd'' nullPtr nullPtr True
122 123 124
                           creationflags
                           nullPtr nullPtr p_startupinfo p_pi
       unless b $ errorWin "createProcessW"
125

126
       pi <- peek p_pi
127 128 129
       b_assign <- assignProcessToJobObject job (piProcess pi)
       unless b_assign $ errorWin "assignProcessToJobObject, cannot continue."

130 131
       let handleInterrupt action =
               action `onException` terminateJobObject job 99
132

133 134 135 136 137
       handleInterrupt $ do
          resumeThread (piThread pi)
          -- The program is now running
          let handle = piProcess pi
          let millisecs = secs * 1000
138 139 140 141
          rc <- waitForJobCompletion job ioPort (fromIntegral millisecs)
          closeHandle ioPort

          if not rc
142
              then do terminateJobObject job 99
143
                      closeHandle job
144 145
                      exitWith (ExitFailure 99)
              else alloca $ \p_exitCode ->
146 147 148
                    do terminateJobObject job 0 -- Ensure it's all really dead.
                       closeHandle job
                       r <- getExitCodeProcess handle p_exitCode
149 150 151 152 153 154
                       if r then do ec <- peek p_exitCode
                                    let ec' = if ec == 0
                                              then ExitSuccess
                                              else ExitFailure $ fromIntegral ec
                                    exitWith ec'
                            else errorWin "getExitCodeProcess"
155
#endif
156