From 91c00c5b33848918d4db136caab4bf82a709884c Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Sat, 10 May 2025 10:28:07 +0400
Subject: [PATCH 1/2] rethrow feat: support for exception context propagation

We specialize the `throwIO` call using a newly implemented `rethrowIO'`
which behaves as `rethrowIO` from base 4.21 when available or like the
previous `throw` implementation.

In short:

- Before `base-4.21`, the code is exactly as before
- After `base-4.21`, the code does not override the backtrace
  annotations and instead uses `rethrowIO`.

Example of usage / changes:

The following code:

```haskell
{-# LANGUAGE DeriveAnyClass #-}
import Control.Concurrent.Async
import Control.Exception
import Control.Exception.Context
import Control.Exception.Annotation
import Data.Typeable
import Data.Traversable
import GHC.Stack

data Ann = Ann String
  deriving (Show, ExceptionAnnotation)

asyncTask :: HasCallStack => IO ()
asyncTask = annotateIO (Ann "bonjour") $ do
  error "yoto"

asyncTask' :: HasCallStack => IO ()
asyncTask' = annotateIO (Ann "bonjour2") $ do
  error "yutu"

main = do
  -- withAsync asyncTask wait
  concurrently asyncTask asyncTask'
  -- race asyncTask asyncTask'
```

When run without this commit leads to:

```
ASyncException.hs: Uncaught exception ghc-internal:GHC.Internal.Exception.ErrorCall:

yoto

HasCallStack backtrace:
  throwIO, called at ./Control/Concurrent/Async/Internal.hs:630:24 in async-2.2.5-50rpfAJ7BEc1o5OswtTMUN:Control.Concurrent.Async.Internal
```

When run with this commit:

```
*** Exception: yoto

Ann "bonjour"
HasCallStack backtrace:
  error, called at /home/guillaume//ASyncException.hs:15:3 in async-2.2.5-inplace:Main
  asyncTask, called at /home/guillaume//ASyncException.hs:23:16 in async-2.2.5-inplace:Main
```
---
 Control/Concurrent/Async/Internal.hs | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index e099fd9..f07b389 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -7,6 +7,7 @@
 {-# LANGUAGE DeriveDataTypeable #-}
 #endif
 {-# OPTIONS -Wall #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -182,10 +183,23 @@ withAsyncUsing doFork = \action inner -> do
     let a = Async t (readTMVar var)
     r <- restore (inner a) `catchAll` \e -> do
       uninterruptibleCancel a
-      throwIO e
+      rethrowIO' e
     uninterruptibleCancel a
     return r
 
+
+-- | This function attempts at rethrowing while keeping the context
+-- This is internal and only working with GHC >=9.12
+rethrowIO' :: SomeException -> IO a
+#if MIN_VERSION_base(4,21,0)
+rethrowIO' e = 
+  case fromException e of
+    Just (e' :: ExceptionWithContext SomeException) -> rethrowIO e'
+    Nothing -> throwIO e
+#else
+rethrowIO' = throwIO
+#endif
+
 -- | Wait for an asynchronous action to complete, and return its
 -- value.  If the asynchronous action threw an exception, then the
 -- exception is re-thrown by 'wait'.
@@ -613,7 +627,7 @@ race left right = concurrently' left right collect
     collect m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right r -> return r
 
 -- race_ :: IO a -> IO b -> IO ()
@@ -627,7 +641,7 @@ concurrently left right = concurrently' left right (collect [])
     collect xs m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right r -> collect (r:xs) m
 
 -- concurrentlyE :: IO (Either e a) -> IO (Either e b) -> IO (Either e (a, b))
@@ -640,7 +654,7 @@ concurrentlyE left right = concurrently' left right (collect [])
     collect xs m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right r -> collect (r:xs) m
 
 concurrently' :: IO a -> IO b
@@ -699,7 +713,7 @@ concurrently_ left right = concurrently' left right (collect 0)
     collect i m = do
         e <- m
         case e of
-            Left ex -> throwIO ex
+            Left ex -> rethrowIO' ex
             Right _ -> collect (i + 1 :: Int) m
 
 

From d374da2922beb63dcef61c40caae423345ac73ab Mon Sep 17 00:00:00 2001
From: Guillaume Bouchard <guillaum.bouchard@gmail.com>
Date: Mon, 12 May 2025 17:02:51 +0400
Subject: [PATCH 2/2] Handle the case for waitSTM

---
 Control/Concurrent/Async/Internal.hs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/Control/Concurrent/Async/Internal.hs b/Control/Concurrent/Async/Internal.hs
index f07b389..13812c4 100644
--- a/Control/Concurrent/Async/Internal.hs
+++ b/Control/Concurrent/Async/Internal.hs
@@ -242,7 +242,12 @@ poll = atomically . pollSTM
 waitSTM :: Async a -> STM a
 waitSTM a = do
    r <- waitCatchSTM a
-   either throwSTM return r
+   either (rethrowSTM) return r
+
+rethrowSTM e = 
+  case fromException e of
+    Just (e' :: ExceptionWithContext SomeException) -> throwSTM (NoBacktrace e')
+    Nothing -> throwSTM e
 
 -- | A version of 'waitCatch' that can be used inside an STM transaction.
 --