From 3e01da22b32b9fe488eb4d61efe91168fa6f270f Mon Sep 17 00:00:00 2001
From: parsonsmatt <parsonsmatt@gmail.com>
Date: Fri, 27 May 2022 16:10:25 -0600
Subject: [PATCH 1/3] Test demonstrating problem

---
 unliftio/test/UnliftIO/ExceptionSpec.hs | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/unliftio/test/UnliftIO/ExceptionSpec.hs b/unliftio/test/UnliftIO/ExceptionSpec.hs
index 82a7772..2e9964b 100644
--- a/unliftio/test/UnliftIO/ExceptionSpec.hs
+++ b/unliftio/test/UnliftIO/ExceptionSpec.hs
@@ -3,7 +3,7 @@
 module UnliftIO.ExceptionSpec (spec) where
 
 import qualified Control.Exception
-import Control.Monad (void, (<=<))
+import Control.Monad (void, (<=<), when)
 import Data.Bifunctor (first)
 import Test.Hspec
 import UnliftIO
@@ -18,7 +18,7 @@ cancelled = Control.Exception.ThreadKilled
 #endif
 
 spec :: Spec
-spec = do
+spec = focus $ do
   let -- The callback will run in a thread that gets cancelled immediately,
       -- then get Exception2 thrown synchronously after 1 second.
       withAsyncExceptionThrown :: (IO a -> IO b) -> IO b
@@ -79,6 +79,27 @@ spec = do
       result <- withWrappedAsyncExceptionThrown $ \m -> trySyncOrAsync (void m)
       first fromExceptionUnwrap result `shouldBe` Left (Just Exception1)
 
+  describe "withException" $ do
+    it "should work when withAsync is in the handler" $ do
+      let
+        action =
+          error "oops"
+            `onException` do
+              let
+                timerAction n = do
+                  threadDelay 1000000
+                  when (n < 10) $ do
+                    timerAction (n + 1)
+              withAsync (timerAction 0) $ \a -> do
+                cancel a
+      eresult <-
+        race
+          (action `shouldThrow` errorCall "oops")
+          (do
+            threadDelay 1000000
+            pure 10)
+      eresult `shouldBe` Left ()
+
   describe "fromExceptionUnwrap" $ do
     it "should be the inverse of toAsyncException" $ do
       fromExceptionUnwrap (toAsyncException Exception1) `shouldBe` Just Exception1

From 4890a736906b4bff465f63385fceeb3e4682a44e Mon Sep 17 00:00:00 2001
From: parsonsmatt <parsonsmatt@gmail.com>
Date: Fri, 27 May 2022 16:18:57 -0600
Subject: [PATCH 2/3] remove focus

---
 unliftio/src/UnliftIO/Exception.hs      | 2 +-
 unliftio/test/UnliftIO/ExceptionSpec.hs | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/unliftio/src/UnliftIO/Exception.hs b/unliftio/src/UnliftIO/Exception.hs
index 493a475..c7509b5 100644
--- a/unliftio/src/UnliftIO/Exception.hs
+++ b/unliftio/src/UnliftIO/Exception.hs
@@ -422,7 +422,7 @@ withException thing after = withRunInIO $ \run -> EUnsafe.uninterruptibleMask $
     case res1 of
         Left e1 -> do
             -- see explanation in bracket
-            _ :: Either SomeException b <- EUnsafe.try $ run $ after e1
+            _ :: Either SomeException b <- EUnsafe.try $ restore $ run $ after e1
             EUnsafe.throwIO e1
         Right x -> return x
 
diff --git a/unliftio/test/UnliftIO/ExceptionSpec.hs b/unliftio/test/UnliftIO/ExceptionSpec.hs
index 2e9964b..583fe4d 100644
--- a/unliftio/test/UnliftIO/ExceptionSpec.hs
+++ b/unliftio/test/UnliftIO/ExceptionSpec.hs
@@ -18,7 +18,7 @@ cancelled = Control.Exception.ThreadKilled
 #endif
 
 spec :: Spec
-spec = focus $ do
+spec = do
   let -- The callback will run in a thread that gets cancelled immediately,
       -- then get Exception2 thrown synchronously after 1 second.
       withAsyncExceptionThrown :: (IO a -> IO b) -> IO b

From f63417570064531f639333d6fe8b7008e7836016 Mon Sep 17 00:00:00 2001
From: parsonsmatt <parsonsmatt@gmail.com>
Date: Fri, 27 May 2022 16:51:59 -0600
Subject: [PATCH 3/3] alternate fix

---
 unliftio/src/UnliftIO/Exception.hs       | 2 +-
 unliftio/src/UnliftIO/Internals/Async.hs | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/unliftio/src/UnliftIO/Exception.hs b/unliftio/src/UnliftIO/Exception.hs
index c7509b5..493a475 100644
--- a/unliftio/src/UnliftIO/Exception.hs
+++ b/unliftio/src/UnliftIO/Exception.hs
@@ -422,7 +422,7 @@ withException thing after = withRunInIO $ \run -> EUnsafe.uninterruptibleMask $
     case res1 of
         Left e1 -> do
             -- see explanation in bracket
-            _ :: Either SomeException b <- EUnsafe.try $ restore $ run $ after e1
+            _ :: Either SomeException b <- EUnsafe.try $ run $ after e1
             EUnsafe.throwIO e1
         Right x -> return x
 
diff --git a/unliftio/src/UnliftIO/Internals/Async.hs b/unliftio/src/UnliftIO/Internals/Async.hs
index 4b1a292..2130163 100644
--- a/unliftio/src/UnliftIO/Internals/Async.hs
+++ b/unliftio/src/UnliftIO/Internals/Async.hs
@@ -75,7 +75,13 @@ asyncOnWithUnmask i m =
 --
 -- @since 0.1.0.0
 withAsync :: MonadUnliftIO m => m a -> (Async a -> m b) -> m b
-withAsync a b = withRunInIO $ \run -> A.withAsync (run a) (run . b)
+withAsync a b = withRunInIO $ \run -> do
+  maskingState <- E.getMaskingState
+  case maskingState of
+    E.MaskedUninterruptible ->
+      A.withAsyncWithUnmask (\unmask -> unmask (run a)) (run . b)
+    _ ->
+      A.withAsync (run a) (run . b)
 
 -- | Unlifted 'A.withAsyncBound'.
 --