Fix cleanup when stdin/stdout handles are inherited and are Nothing

Author: pranaysashank

src/Streamly/Internal/System/Process.hs

@@ -77,6 +77,8 @@ module Streamly.Internal.System.Process
     , inheritStdin
     , inheritStdout
     , pipeStdErr
+    , pipeStdin
+    , pipeStdout
 
     -- * Exceptions
     , ProcessFailure (..)
@@ -214,6 +216,12 @@ mkConfig _ _ = Config False
 pipeStdErr :: Config -> Config
 pipeStdErr (Config _) = Config True
 
+pipeStdin :: Config -> Config
+pipeStdin (Config _) = Config True
+
+pipeStdout :: Config -> Config
+pipeStdout (Config _) = Config True
+
 inheritStdin :: Config -> Config
 inheritStdin (Config _) = Config True
 
@@ -453,6 +461,12 @@ waitForChildTree = waitForDescendants
 pipeStdErr :: Config -> Config
 pipeStdErr (Config cfg) = Config $ cfg { std_err = CreatePipe }
 
+pipeStdin :: Config -> Config
+pipeStdin (Config cleanupCfg cfg) = Config cleanupCfg $ cfg { std_in = CreatePipe }
+
+pipeStdout :: Config -> Config
+pipeStdout (Config cleanupCfg cfg) = Config cleanupCfg $ cfg { std_out = CreatePipe }
+
 inheritStdin :: Config -> Config
 inheritStdin (Config cfg) = Config $ cfg { std_in = Inherit }
 
@@ -518,7 +532,7 @@ cleanupNormal (_, _, _, procHandle) = do
 -- still hanging around.
 cleanupException ::
     (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle) -> IO ()
-cleanupException (Just stdinH, Just stdoutH, stderrMaybe, ph) = do
+cleanupException (stdinMaybe, stdoutMaybe, stderrMaybe, ph) = do
     -- Send a SIGTERM to the process
 #ifdef USE_NATIVE
     terminate ph
@@ -529,8 +543,8 @@ cleanupException (Just stdinH, Just stdoutH, stderrMaybe, ph) = do
     -- Ideally we should be closing the handle without flushing the buffers so
     -- that we cannot get a SIGPIPE. But there seems to be no way to do that as
     -- of now so we just ignore the SIGPIPE.
-    hClose stdinH `catch` eatSIGPIPE
-    hClose stdoutH
+    whenJust (\stdinH -> hClose stdinH `catch` eatSIGPIPE) stdinMaybe
+    whenJust hClose stdoutMaybe
     whenJust hClose stderrMaybe
 
     -- Non-blocking wait for the process to go away
@@ -553,7 +567,6 @@ cleanupException (Just stdinH, Just stdoutH, stderrMaybe, ph) = do
             _ -> False
 
     eatSIGPIPE e = unless (isSIGPIPE e) $ throwIO e
-cleanupException _ = error "cleanupProcess: Not reachable"
 
 -- | Creates a system process from an executable path and arguments. For the
 -- default attributes used to create the process see 'mkConfig'.
@@ -618,6 +631,11 @@ pipeChunksWithAction run modCfg path args =
 
     alloc = createProc' modCfg path args
 
+-- Note: It is allowed to inheritStdout or inheritStderr but not both as that
+-- would not generate a stream for further processing.
+-- inheritStdin has no affect, we always pipe the input stream to the process's
+-- stdin
+
 -- | Like 'pipeChunksEither' but use the specified configuration to run the
 -- process.
 {-# INLINE pipeChunksEitherWith #-}
@@ -629,10 +647,12 @@ pipeChunksEitherWith ::
     -> Stream m (Array Word8)    -- ^ Input stream
     -> Stream m (Either (Array Word8) (Array Word8))     -- ^ Output stream
 pipeChunksEitherWith modifier path args input =
-    pipeChunksWithAction run (modifier . pipeStdErr) path args
+    pipeChunksWithAction run (pipeStdin . modifier . pipeStdErr) path args
 
     where
 
+    run (_, Nothing, Nothing, _) =
+        error "pipeChunksEitherWith: only one of stdout or stderr can be inherited"
     run (Just stdinH, Just stdoutH, Just stderrH, _) =
         putChunksClose stdinH input
             `parallel` fmap Left (toChunksClose stderrH)
@@ -692,6 +712,10 @@ pipeBytesEither path args input =
         rightRdr = fmap Right Array.reader
      in Stream.unfoldMany (Unfold.either leftRdr rightRdr) output
 
+-- Note: inheritStdin, inheritStdout have no affect, we always pipe
+-- the input stream to the process's stdin and pipe the stdout to the
+-- resulting stream
+
 -- | Like 'pipeChunks' but use the specified configuration to run the process.
 {-# INLINE pipeChunksWith #-}
 pipeChunksWith ::
@@ -702,13 +726,14 @@ pipeChunksWith ::
     -> Stream m (Array Word8)    -- ^ Input stream
     -> Stream m (Array Word8)    -- ^ Output stream
 pipeChunksWith modifier path args input =
-    pipeChunksWithAction run modifier path args
+    pipeChunksWithAction run (pipeStdout . pipeStdin . modifier) path args
 
     where
 
     run (Just stdinH, Just stdoutH, _, _) =
-        putChunksClose stdinH input `parallel` toChunksClose stdoutH
-    run _ = error "pipeChunksWith: Not reachable"
+        putChunksClose stdinH input
+            `parallel` toChunksClose stdoutH
+    run _ = error "pipeChunksWith: unreachable"
 
 -- | @pipeChunks file args input@ runs the executable @file@ specified by
 -- its name or path using @args@ as arguments and @input@ stream as its
@@ -841,6 +866,10 @@ pipeChars path args input =
 -- Generation
 -------------------------------------------------------------------------------
 
+-- Note: It is allowed to inheritStdout or inheritStderr but not both as that
+-- would not generate a stream for further processing and would result in unintuitive
+-- behaviour
+
 -- | Like 'toChunksEither' but use the specified configuration to run the
 -- process.
 {-# INLINE toChunksEitherWith #-}
@@ -851,14 +880,18 @@ toChunksEitherWith ::
     -> [String]             -- ^ Arguments
     -> Stream m (Either (Array Word8) (Array Word8))     -- ^ Output stream
 toChunksEitherWith modifier path args =
-    pipeChunksWithAction run (modifier . inheritStdin . pipeStdErr) path args
+    pipeChunksWithAction run (modifier . inheritStdin . pipeStdErr . pipeStdout) path args
 
     where
 
-    run (_, Just stdoutH, Just stderrH, _) =
-        fmap Left (toChunksClose stderrH)
-            `parallel` fmap Right (toChunksClose stdoutH)
-    run _ = error "toChunksEitherWith: Not reachable"
+    run (_, Nothing, Nothing, _) =
+        error "toChunksEitherWith: only one of stdout or stderr can be inherited"
+    run (_, stdoutMaybe, stderrMaybe, _) =
+        fmap Left (whenJustS toChunksClose stderrMaybe)
+            `parallel` fmap Right (whenJustS toChunksClose stdoutMaybe)
+
+-- Note: inheritStdout has no affect, we always pipe stdout to the resulting
+-- stream
 
 -- | Like 'toChunks' but use the specified configuration to run the process.
 {-# INLINE toChunksWith #-}
@@ -869,12 +902,12 @@ toChunksWith ::
     -> [String]             -- ^ Arguments
     -> Stream m (Array Word8)    -- ^ Output stream
 toChunksWith modifier path args =
-    pipeChunksWithAction run (modifier . inheritStdin) path args
+    pipeChunksWithAction run (pipeStdout . modifier . inheritStdin) path args
 
     where
 
     run (_, Just stdoutH, _, _) = toChunksClose stdoutH
-    run _ = error "toChunksWith: Not reachable"
+    run _                       = error "toChunksWith: Not reachable"
 
 -- | @toBytesEither path args@ runs the executable at @path@ using @args@ as
 -- arguments and returns the output of the process as a stream of 'Either'
@@ -1181,3 +1214,7 @@ daemon modCfg path args =
                 (setSession NewSession . modCfg)
                 path args
      in fmap (either undefined id) r
+
+
+whenJustS :: Applicative m => (a -> Stream m b) -> Maybe a -> Stream m b
+whenJustS action mb = maybe Stream.nil action mb