async exception safety for coprocesses

Tested the forcerestart code path and it works.

The hairy part is, what if an async exception is caught when it's in
restart?

If it's in the part that stops the old process, the old process
is left in the handle. The next attempt to use the CoProcessHandle
will then throw an IO exception, which will result in restart getting
run again. So I think this will work, but have not actually tested it.

The use of withMVarMasked lets it start the new process and fill the
mvar with it, even if there's an async exception at that point.

Note that exceptions are masked while running forcerestart, so
do not need to worry about an async exception being thrown while it's
recovering from an async exception.

Utility/CoProcess.hs

@@ -1,7 +1,7 @@
 {- Interface for running a shell command as a coprocess,
  - sending it queries and getting back results.
  -
- - Copyright 2012-2013 Joey Hess <id@joeyh.name>
+ - Copyright 2012-2020 Joey Hess <id@joeyh.name>
  -
  - License: BSD-2-clause
  -}
@@ -62,28 +62,46 @@ stop ch = do
 	let p = proc (coProcessCmd $ coProcessSpec s) (coProcessParams $ coProcessSpec s)
 	forceSuccessProcess p (coProcessPid s)
 
-{- To handle a restartable process, any IO exception thrown by the send and
+{- Note that concurrent queries are not safe to perform; caller should
+ - serialize calls to query.
+ -
+ - To handle a restartable process, any IO exception thrown by the send and
  - receive actions are assumed to mean communication with the process
- - failed, and the failed action is re-run with a new process. -}
+ - failed, and the query is re-run with a new process.
+ -
+ - If an async exception is received during a query, the state of
+ - communication with the process is unknown, so it is killed, and a new
+ - one started so the CoProcessHandle can continue to be used by other
+ - threads.
+ -}
 query :: CoProcessHandle -> (Handle -> IO a) -> (Handle -> IO b) -> IO b
-query ch send receive = do
-	s <- readMVar ch
-	restartable s (send $ coProcessTo s) $ const $
-		restartable s (hFlush $ coProcessTo s) $ const $
-			restartable s (receive $ coProcessFrom s)
-				return
+query ch send receive = uninterruptibleMask $ \unmask ->
+	unmask (readMVar ch >>= restartable)
+		`catchAsync` forcerestart
   where
-	restartable s a cont
+	go s = do
+		void $ send $ coProcessTo s
+		hFlush $ coProcessTo s
+		receive $ coProcessFrom s
+	
+	restartable s
 		| coProcessNumRestarts (coProcessSpec s) > 0 =
-			maybe restart cont =<< catchMaybeIO a
-		| otherwise = cont =<< a
-	restart = do
-		s <- takeMVar ch
-		void $ catchMaybeIO $ do
+			catchMaybeIO (go s)
+				>>= maybe (restart s increstarts restartable) return
+		| otherwise = go s
+	
+	increstarts s = s { coProcessNumRestarts = coProcessNumRestarts s - 1 }
+
+	restart s f cont = do
+		void $ tryNonAsync $ do
 			hClose $ coProcessTo s
 			hClose $ coProcessFrom s
 		void $ waitForProcess $ coProcessPid s
-		s' <- start' $ (coProcessSpec s)
-			{ coProcessNumRestarts = coProcessNumRestarts (coProcessSpec s) - 1 }
-		putMVar ch s'
-		query ch send receive
+		s' <- withMVarMasked ch $ \_ -> start' (f (coProcessSpec s))
+		cont s'
+
+	forcerestart ex = do
+		s <- readMVar ch
+		terminateProcess (coProcessPid s)
+		restart s id $ \s' -> void $ swapMVar ch s'
+		either throwM throwM ex

Utility/Exception.hs

@@ -1,6 +1,6 @@
 {- Simple IO exception handling (and some more)
  -
- - Copyright 2011-2016 Joey Hess <id@joeyh.name>
+ - Copyright 2011-2020 Joey Hess <id@joeyh.name>
  -
  - License: BSD-2-clause
  -}
@@ -20,6 +20,7 @@ module Utility.Exception (
 	bracketIO,
 	catchNonAsync,
 	tryNonAsync,
+	catchAsync,
 	tryWhenExists,
 	catchIOErrorType,
 	IOErrorType(..),
@@ -87,6 +88,14 @@ catchNonAsync a onerr = a `catches`
 	, M.Handler (\ (e :: SomeException) -> onerr e)
 	]
 
+{- Catches only async exceptions. -}
+catchAsync :: MonadCatch m => m a -> (Either AsyncException SomeAsyncException -> m a) -> m a
+catchAsync a onerr = a `catches`
+	[ M.Handler (\ (e :: AsyncException) -> onerr (Left e))
+	, M.Handler (\ (e :: SomeAsyncException) -> onerr (Right e))
+	, M.Handler (\ (e :: SomeException) -> throwM e)
+	]
+
 tryNonAsync :: MonadCatch m => m a -> m (Either SomeException a)
 tryNonAsync a = go `catchNonAsync` (return . Left)
   where