waitAnyCancel should be using restoreM

maoe · Dec 22, 2014 · a9ba1f4 · a9ba1f4
1 parent 91e0575
commit a9ba1f4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/src/Control/Concurrent/Async/Lifted.hs b/src/Control/Concurrent/Async/Lifted.hs
@@ -98,8 +98,14 @@ waitAnyCatch as = do
   return (a, r)
 
 -- | Generalized version of 'A.waitAnyCancel'.
-waitAnyCancel :: MonadBase IO m => [Async a] -> m (Async a, a)
-waitAnyCancel = liftBase . A.waitAnyCancel
+waitAnyCancel
+  :: MonadBaseControl IO m
+  => [Async (StM m a)]
+  -> m (Async (StM m a), a)
+waitAnyCancel as = do
+  (a, s) <- liftBase $ A.waitAnyCancel as
+  r <- restoreM s
+  return (a, r)
 
 -- | Generalized version of 'A.waitAnyCatchCancel'.
 waitAnyCatchCancel

diff --git a/src/Control/Concurrent/Async/Lifted/Safe.hs b/src/Control/Concurrent/Async/Lifted/Safe.hs
@@ -204,8 +204,14 @@ waitAnyCatch as = do
   return (a, r)
 
 -- | Generalized version of 'A.waitAnyCancel'.
-waitAnyCancel :: MonadBase IO m => [Async a] -> m (Async a, a)
-waitAnyCancel = liftBase . A.waitAnyCancel
+waitAnyCancel
+  :: (MonadBaseControl IO m, StM m a ~ a)
+  => [Async (StM m a)]
+  -> m (Async (StM m a), a)
+waitAnyCancel as = do
+  (a, s) <- liftBase $ A.waitAnyCancel as
+  r <- restoreM s
+  return (a, r)
 
 -- | Generalized version of 'A.waitAnyCatchCancel'.
 waitAnyCatchCancel