Skip to content

Commit

Permalink
Switch to checked exit code functions
Browse files Browse the repository at this point in the history
  • Loading branch information
snoyberg committed Oct 19, 2016
1 parent 7ef0877 commit 22bad35
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 69 deletions.
179 changes: 113 additions & 66 deletions src/System/Process/Typed.hs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ module System.Process.Typed
, setChildGroup
, setChildUser
#endif
, setCheckExitCode

-- * Stream specs
, mkStreamSpec
Expand All @@ -54,7 +53,9 @@ module System.Process.Typed
, startProcess
, stopProcess
, withProcess
, withProcess_
, readProcess
, readProcess_
, runProcess
, runProcess_

Expand All @@ -63,6 +64,8 @@ module System.Process.Typed
-- ** Process exit code
, waitExitCode
, waitExitCodeSTM
, getExitCode
, getExitCodeSTM
, checkExitCode
, checkExitCodeSTM

Expand Down Expand Up @@ -149,9 +152,30 @@ data ProcessConfig stdin stdout stderr = ProcessConfig
, pcChildGroup :: !(Maybe GroupID)
, pcChildUser :: !(Maybe UserID)
#endif

, pcCheckExitCode :: !Bool
}
instance Show (ProcessConfig stdin stdout stderr) where
show pc = concat
[ case pcCmdSpec pc of
P.ShellCommand s -> "Shell command: " ++ s
P.RawCommand x xs -> "Raw command: " ++ unwords (map escape (x:xs))
, "\n"
, case pcWorkingDir pc of
Nothing -> ""
Just wd -> concat
[ "Run from: "
, wd
, "\n"
]
, case pcEnv pc of
Nothing -> ""
Just e -> unlines
$ "Modified environment:"
: map (\(k, v) -> concat [k, "=", v]) e
]
where
escape x
| any (`elem` " \\\"'") x = show x
| otherwise = x
instance (stdin ~ (), stdout ~ (), stderr ~ ())
=> IsString (ProcessConfig stdin stdout stderr) where
fromString s
Expand All @@ -172,7 +196,7 @@ data StreamType = STInput | STOutput
-- @since 0.1.0.0
data StreamSpec (streamType :: StreamType) a = StreamSpec
{ ssStream :: !P.StdStream
, ssCreate :: !(Maybe Handle -> Cleanup a)
, ssCreate :: !(ProcessConfig () () () -> Maybe Handle -> Cleanup a)
}
deriving Functor

Expand Down Expand Up @@ -202,13 +226,16 @@ instance Applicative Cleanup where
--
-- @since 0.1.0.0
data Process stdin stdout stderr = Process
{ pCleanup :: !(IO ())
{ pConfig :: !(ProcessConfig () () ())
, pCleanup :: !(IO ())
, pStdin :: !stdin
, pStdout :: !stdout
, pStderr :: !stderr
, pHandle :: !P.ProcessHandle
, pExitCode :: !(TMVar ExitCode)
}
instance Show (Process stdin stdout stderr) where
show p = "Running process: " ++ show (pConfig p)

-- | Internal helper
defaultProcessConfig :: ProcessConfig () () ()
Expand All @@ -233,8 +260,6 @@ defaultProcessConfig = ProcessConfig
, pcChildGroup = Nothing
, pcChildUser = Nothing
#endif

, pcCheckExitCode = False
}

-- | Create a 'ProcessConfig' from the given command and arguments.
Expand Down Expand Up @@ -409,27 +434,6 @@ setChildUser
setChildUser x pc = pc { pcChildUser = Just x }
#endif

-- | Should we throw an exception when the process exits with a
-- non-success code?
--
-- If set to 'True', then when 'stopProcess' is called - either
-- directly or via 'withProcess' or other wrappers - the processes
-- exit code will be checked. Any exit code besides 'ExitSuccess' will
-- result in an 'ExitCodeException' being thrown.
--
-- Default: 'False'
--
-- @since 0.1.0.0
setCheckExitCode :: Bool
-> ProcessConfig stdin stdout stderr
-> ProcessConfig stdin stdout stderr
setCheckExitCode x p = p { pcCheckExitCode = x }

-- TODO: Instead of having this setting, we could consider just having
-- alternatives to readProcess, runProcess, etc, that check the exit
-- code. This could actually be a really nice convention: readProcess
-- does not check, readProcess_ or readProcessCheck does.

-- | Create a new 'StreamSpec' from the given 'P.StdStream' and a
-- helper function. This function:
--
Expand All @@ -442,25 +446,25 @@ setCheckExitCode x p = p { pcCheckExitCode = x }
--
-- @since 0.1.0.0
mkStreamSpec :: P.StdStream
-> (Maybe Handle -> IO (a, IO ()))
-> (ProcessConfig () () () -> Maybe Handle -> IO (a, IO ()))
-> StreamSpec streamType a
mkStreamSpec ss f = StreamSpec ss (Cleanup . f)
mkStreamSpec ss f = StreamSpec ss (\pc mh -> Cleanup (f pc mh))

-- | A stream spec which simply inherits the stream of the parent
-- process.
--
-- @since 0.1.0.0
inherit :: StreamSpec anyStreamType ()
inherit = mkStreamSpec P.Inherit (\Nothing -> pure ((), return ()))
inherit = mkStreamSpec P.Inherit (\_ Nothing -> pure ((), return ()))

-- | A stream spec which will close the stream for the child process.
--
-- @since 0.1.0.0
closed :: StreamSpec anyStreamType ()
#if MIN_VERSION_process(1, 4, 0)
closed = mkStreamSpec P.NoStream (\Nothing -> pure ((), return ()))
closed = mkStreamSpec P.NoStream (\_ Nothing -> pure ((), return ()))
#else
closed = mkStreamSpec P.CreatePipe (\(Just h) -> (((), return ()) <$ hClose h))
closed = mkStreamSpec P.CreatePipe (\_ (Just h) -> (((), return ()) <$ hClose h))
#endif

-- | An input stream spec which sets the input to the given
Expand All @@ -469,7 +473,7 @@ closed = mkStreamSpec P.CreatePipe (\(Just h) -> (((), return ()) <$ hClose h))
--
-- @since 0.1.0.0
byteStringInput :: L.ByteString -> StreamSpec 'STInput ()
byteStringInput lbs = StreamSpec P.CreatePipe $ \(Just h) -> Cleanup $ do
byteStringInput lbs = mkStreamSpec P.CreatePipe $ \_ (Just h) -> do
void $ async $ do
L.hPut h lbs
hClose h
Expand All @@ -489,7 +493,7 @@ byteStringInput lbs = StreamSpec P.CreatePipe $ \(Just h) -> Cleanup $ do
--
-- @since 0.1.0.0
byteStringOutput :: StreamSpec 'STOutput (STM (Either ByteStringOutputException L.ByteString))
byteStringOutput = StreamSpec P.CreatePipe $ \(Just h) -> Cleanup $ do
byteStringOutput = mkStreamSpec P.CreatePipe $ \pc (Just h) -> do
mvar <- newEmptyTMVarIO

void $ async $ do
Expand All @@ -499,7 +503,7 @@ byteStringOutput = StreamSpec P.CreatePipe $ \(Just h) -> Cleanup $ do
then atomically $ putTMVar mvar $ Right $ L.fromChunks $ front []
else loop $ front . (bs:)
loop id `catch` \e -> do
atomically $ void $ tryPutTMVar mvar $ Left $ ByteStringOutputException e
atomically $ void $ tryPutTMVar mvar $ Left $ ByteStringOutputException e pc
throwIO e

return (readTMVar mvar, hClose h)
Expand All @@ -509,7 +513,7 @@ byteStringOutput = StreamSpec P.CreatePipe $ \(Just h) -> Cleanup $ do
--
-- @since 0.1.0.0
createPipe :: StreamSpec anyStreamType Handle
createPipe = StreamSpec P.CreatePipe $ \(Just h) -> Cleanup $ return (h, hClose h)
createPipe = mkStreamSpec P.CreatePipe $ \_ (Just h) -> return (h, hClose h)

-- | Use the provided 'Handle' for the child process, and when the
-- process exits, do /not/ close it. This is useful if, for example,
Expand All @@ -518,15 +522,15 @@ createPipe = StreamSpec P.CreatePipe $ \(Just h) -> Cleanup $ return (h, hClose
--
-- @since 0.1.0.0
useHandleOpen :: Handle -> StreamSpec anyStreamType ()
useHandleOpen h = StreamSpec (P.UseHandle h) $ \Nothing -> Cleanup $ return ((), return ())
useHandleOpen h = mkStreamSpec (P.UseHandle h) $ \_ Nothing -> return ((), return ())

-- | Use the provided 'Handle' for the child process, and when the
-- process exits, close it. If you have no reason to keep the 'Handle'
-- open, you should use this over 'useHandleOpen'.
--
-- @since 0.1.0.0
useHandleClose :: Handle -> StreamSpec anyStreamType ()
useHandleClose h = StreamSpec (P.UseHandle h) $ \Nothing -> Cleanup $ return ((), hClose h)
useHandleClose h = mkStreamSpec (P.UseHandle h) $ \_ Nothing -> return ((), hClose h)

-- | Provide input to a process by writing to a conduit.
--
Expand All @@ -553,7 +557,7 @@ source =
startProcess :: MonadIO m
=> ProcessConfig stdin stdout stderr
-> m (Process stdin stdout stderr)
startProcess ProcessConfig {..} = liftIO $ do
startProcess pConfig'@ProcessConfig {..} = liftIO $ do
let cp0 =
case pcCmdSpec of
P.ShellCommand cmd -> P.shell cmd
Expand Down Expand Up @@ -584,16 +588,16 @@ startProcess ProcessConfig {..} = liftIO $ do
(minH, moutH, merrH, pHandle) <- P.createProcess_ "startProcess" cp

((pStdin, pStdout, pStderr), pCleanup1) <- runCleanup $ (,,)
<$> ssCreate pcStdin minH
<*> ssCreate pcStdout moutH
<*> ssCreate pcStderr merrH
<$> ssCreate pcStdin pConfig minH
<*> ssCreate pcStdout pConfig moutH
<*> ssCreate pcStderr pConfig merrH

pExitCode <- newEmptyTMVarIO
void $ async $ do
ec <- P.waitForProcess pHandle
atomically $ putTMVar pExitCode ec

let pCleanup2 = pCleanup1 `finally` do
let pCleanup = pCleanup1 `finally` do
mec <- atomically $ tryReadTMVar pExitCode
case mec of
Nothing -> do
Expand All @@ -602,24 +606,16 @@ startProcess ProcessConfig {..} = liftIO $ do
-- a SIGKILL on Unix?
void $ atomically $ readTMVar pExitCode
Just _ -> return ()
pCleanup
| pcCheckExitCode = do
eres <- try pCleanup2
ec <- atomically $ readTMVar pExitCode
case (ec, eres) of
(ExitSuccess, Right ()) -> return ()
(ExitSuccess, Left e) -> throwIO e
_ -> throwIO $ ExitCodeException ec $ either Just (const Nothing) eres
| otherwise = pCleanup2

return Process {..}
where
pConfig = clearStreams pConfig'

-- | Close a process and release any resources acquired. This will
-- ensure 'P.terminateProcess' is called, wait for the process to
-- actually exit, and then close out resources allocated for the
-- streams. In the event of any cleanup exceptions being thrown, or if
-- a non-success exit code was received and 'setCheckExitCode' was
-- used, this will throw an exception.
-- streams. In the event of any cleanup exceptions being thrown this
-- will throw an exception.
--
-- @since 0.1.0.0
stopProcess :: MonadIO m
Expand All @@ -637,6 +633,17 @@ withProcess :: (MonadIO m, C.MonadMask m)
-> m a
withProcess config = C.bracket (startProcess config) stopProcess

-- | Same as 'withProcess', but also calls 'checkExitCode'
--
-- @since 0.1.0.0
withProcess_ :: (MonadIO m, C.MonadMask m)
=> ProcessConfig stdin stdout stderr
-> (Process stdin stdout stderr -> m a)
-> m a
withProcess_ config = C.bracket
(startProcess config)
(\p -> stopProcess p `finally` checkExitCode p)

-- | Run a process, capture its standard output and error as a
-- 'L.ByteString', wait for it to complete, and then return its exit
-- code, output, and error.
Expand All @@ -657,6 +664,22 @@ readProcess pc =
pc' = setStdout byteStringOutput
$ setStderr byteStringOutput pc

-- | Same as 'readProcess', but instead of returning the 'ExitCode',
-- checks it with 'checkExitCode'.
--
-- @since 0.1.0.0
readProcess_ :: MonadIO m
=> ProcessConfig stdin stdoutIgnored stderrIgnored
-> m (L.ByteString, L.ByteString)
readProcess_ pc =
liftIO $ withProcess pc' $ \p -> atomically $ (,)
<$> (checkExitCodeSTM p
*> (getStdout p >>= either throwSTM return))
<*> (getStderr p >>= either throwSTM return)
where
pc' = setStdout byteStringOutput
$ setStderr byteStringOutput pc

-- | Run the given process, wait for it to exit, and returns its
-- 'ExitCode'.
--
Expand All @@ -672,7 +695,7 @@ runProcess pc = liftIO $ withProcess pc waitExitCode
runProcess_ :: MonadIO m
=> ProcessConfig stdin stdout stderr
-> m ()
runProcess_ = liftIO . void . runProcess
runProcess_ pc = liftIO $ withProcess pc checkExitCode

-- | Wait for the process to exit and then return its 'ExitCode'.
--
Expand All @@ -689,14 +712,39 @@ waitExitCodeSTM = readTMVar . pExitCode
-- | Check if a process has exited and, if so, return its 'ExitCode'.
--
-- @since 0.1.0.0
checkExitCode :: MonadIO m => Process stdin stdout stderr -> m (Maybe ExitCode)
getExitCode :: MonadIO m => Process stdin stdout stderr -> m (Maybe ExitCode)
getExitCode = liftIO . atomically . getExitCodeSTM

-- | Same as 'getExitCode', but in 'STM'.
--
-- @since 0.1.0.0
getExitCodeSTM :: Process stdin stdout stderr -> STM (Maybe ExitCode)
getExitCodeSTM = tryReadTMVar . pExitCode

-- | Wait for a process to exit, and ensure that it exited
-- successfully. If not, throws an 'ExitCodeException'.
--
-- @since 0.1.0.0
checkExitCode :: MonadIO m => Process stdin stdout stderr -> m ()
checkExitCode = liftIO . atomically . checkExitCodeSTM

-- | Same as 'checkExitCode', but in 'STM'.
--
-- @since 0.1.0.0
checkExitCodeSTM :: Process stdin stdout stderr -> STM (Maybe ExitCode)
checkExitCodeSTM = tryReadTMVar . pExitCode
checkExitCodeSTM :: Process stdin stdout stderr -> STM ()
checkExitCodeSTM p = do
ec <- readTMVar (pExitCode p)
case ec of
ExitSuccess -> return ()
_ -> throwSTM (ExitCodeException ec (clearStreams (pConfig p)))

-- | Internal
clearStreams :: ProcessConfig stdin stdout stderr -> ProcessConfig () () ()
clearStreams pc = pc
{ pcStdin = inherit
, pcStdout = inherit
, pcStderr = inherit
}

-- | Get the child's standard input stream value.
--
Expand All @@ -716,20 +764,19 @@ getStdout = pStdout
getStderr :: Process stdin stdout stderr -> stderr
getStderr = pStderr

-- | Exit code generated by 'stopProcess' when 'setCheckExitCode' is
-- 'True' and a process exits with a non-success code. Contains the
-- non-success code, and if any other exceptions occur during cleanup,
-- that exception.
-- | Exception thrown by 'checkExitCode' in the event of a non-success
-- exit code. Note that 'checkExitCode' is called by other functions
-- as well, like 'runProcess_' or 'readProcess_'.
--
-- @since 0.1.0.0
data ExitCodeException = ExitCodeException ExitCode (Maybe SomeException)
data ExitCodeException = ExitCodeException ExitCode (ProcessConfig () () ())
deriving (Show, Typeable)
instance Exception ExitCodeException

-- | Wrapper for when an exception is thrown when reading from a child
-- process, used by 'byteStringOutput'.
--
-- @since 0.1.0.0
newtype ByteStringOutputException = ByteStringOutputException SomeException
data ByteStringOutputException = ByteStringOutputException SomeException (ProcessConfig () () ())
deriving (Show, Typeable)
instance Exception ByteStringOutputException
5 changes: 2 additions & 3 deletions test/System/Process/TypedSpec.hs
Original file line number Diff line number Diff line change
Expand Up @@ -72,9 +72,8 @@ spec = do
res <- runProcess "false"
res `shouldBe` ExitFailure 1

it "checked exit code" $ do
runProcess_ (setCheckExitCode True "false")
`shouldThrow` \ExitCodeException{} -> True
it "checked exit code" $
runProcess_ "false" `shouldThrow` \ExitCodeException{} -> True

it "async" $ withSystemTempFile "httpbin" $ \fp h -> do
bss <- withProcess (setStdin sink $ setStdout source "base64") $ \p ->
Expand Down

0 comments on commit 22bad35

Please sign in to comment.