Forward exceptions from async tasks to resource manager

src/Quasar/Async.hs

@@ -5,18 +5,19 @@ module Quasar.Async (
   asyncWithUnmask,
   asyncWithUnmask_,
 
-  -- ** Task exceptions
-  CancelTask(..),
-  TaskDisposed(..),
+  -- ** Async exceptions
+  CancelAsync(..),
+  AsyncDisposed(..),
+  AsyncException(..),
 ) where
 
+import Control.Monad.Catch
 import Control.Monad.Reader
 import Quasar.Async.Unmanaged
 import Quasar.Awaitable
 import Quasar.Prelude
 import Quasar.ResourceManager
 
-
 -- | TODO: Documentation
 --
 -- The action will be run with asynchronous exceptions unmasked.
@@ -30,12 +31,18 @@ asyncWithUnmask :: MonadResourceManager m => ((ResourceManagerIO a -> ResourceMa
 asyncWithUnmask action = do
   resourceManager <- askResourceManager
   toAwaitable <$> registerNewResource do
-    unmanagedAsyncWithUnmask (\unmask -> runReaderT (action (liftUnmask unmask)) resourceManager)
+    coreAsyncImplementation (handler resourceManager) \unmask ->
+      onResourceManager resourceManager (action (liftUnmask unmask))
   where
+    handler :: ResourceManager -> SomeException -> IO ()
+    handler resourceManager ex = when (fromException ex /= Just AsyncDisposed) do
+      -- Throwing to the resource manager is safe because the handler runs on the async thread the resource manager
+      -- cannot reach disposed state until the thread exits
+      throwToResourceManager resourceManager ex
     liftUnmask :: (forall b. IO b -> IO b) -> ResourceManagerIO a -> ResourceManagerIO a
     liftUnmask unmask innerAction = do
       resourceManager <- askResourceManager
-      liftIO $ unmask $ runReaderT innerAction resourceManager
+      liftIO $ unmask $ onResourceManager resourceManager innerAction
 
 async_ :: MonadResourceManager m => (ResourceManagerIO ()) -> m ()
 async_ action = void $ async action

src/Quasar/Async/Unmanaged.hs

@@ -7,9 +7,13 @@ module Quasar.Async.Unmanaged (
   unmanagedAsyncWithUnmask_,
 
   -- ** Task exceptions
-  CancelTask(..),
-  TaskDisposed(..),
-)where
+  CancelAsync(..),
+  AsyncDisposed(..),
+  AsyncException(..),
+
+  -- ** Implementation internals
+  coreAsyncImplementation
+) where
 
 
 import Control.Concurrent (ThreadId, forkIOWithUnmask, throwTo)
@@ -40,7 +44,7 @@ instance IsDisposable (Task r) where
         TaskStateRunning threadId -> do
           writeTVar stateVar TaskStateThrowing
           pure do
-            throwTo threadId $ CancelTask key
+            throwTo threadId $ CancelAsync key
             atomically $ writeTVar stateVar TaskStateCompleted
         TaskStateThrowing -> pure $ pure ()
         TaskStateCompleted -> pure $ pure ()
@@ -56,25 +60,26 @@ instance Functor Task where
   fmap fn (Task key actionVar finalizerVar resultAwaitable) = Task key actionVar finalizerVar (fn <$> resultAwaitable)
 
 
-data CancelTask = CancelTask Unique
-instance Show CancelTask where
-  show _ = "CancelTask"
-instance Exception CancelTask where
+data CancelAsync = CancelAsync Unique
+  deriving stock Eq
+instance Show CancelAsync where
+  show _ = "CancelAsync"
+instance Exception CancelAsync where
 
-data TaskDisposed = TaskDisposed
-  deriving stock Show
-instance Exception TaskDisposed where
+data AsyncDisposed = AsyncDisposed
+  deriving stock (Eq, Show)
+instance Exception AsyncDisposed where
 
+-- TODO Needs a descriptive name. This is similar in functionality to `ExceptionThrownInLinkedThread`
+data AsyncException = AsyncException SomeException
+  deriving stock Show
+  deriving anyclass Exception
 
 
-unmanagedAsync :: MonadIO m => IO a -> m (Task a)
-unmanagedAsync action = unmanagedAsyncWithUnmask \unmask -> unmask action
 
-unmanagedAsync_ :: MonadIO m => IO () -> m Disposable
-unmanagedAsync_ action = toDisposable <$> unmanagedAsync action
-
-unmanagedAsyncWithUnmask :: MonadIO m => ((forall b. IO b -> IO b) -> IO a) -> m (Task a)
-unmanagedAsyncWithUnmask action = do
+-- | Base implementation for the `unmanagedAsync`- and `Quasar.Async.async`-class of functions.
+coreAsyncImplementation :: MonadIO m => (SomeException -> IO ()) -> ((forall b. IO b -> IO b) -> IO a) -> m (Task a)
+coreAsyncImplementation handler action = do
   liftIO $ mask_ do
     key <- newUnique
     resultVar <- newAsyncVar
@@ -83,18 +88,19 @@ unmanagedAsyncWithUnmask action = do
 
     threadId <- forkIOWithUnmask \unmask ->
       handleAll
-        do \ex -> fail $ "unmanagedAsyncWithUnmask thread failed: " <> displayException ex
+        do \ex -> fail $ "coreAsyncImplementation thread failed: " <> displayException ex
         do
-          result <- try $ handleIf
-            do \(CancelTask exKey) -> key == exKey
-            do \_ -> throwIO TaskDisposed
-            do
-              action unmask
+          result <- try $ catchAll
+            do action unmask
+            \ex -> do
+              -- Rewrite exception if its the cancel exception for this async
+              when (fromException ex == Just (CancelAsync key)) $ throwIO AsyncDisposed
+              throwIO $ AsyncException ex
 
           -- The `action` has completed its work.
           -- "disarm" dispose:
           handleIf
-            do \(CancelTask exKey) -> key == exKey
+            do \(CancelAsync exKey) -> key == exKey
             do mempty -- ignore exception if it matches; this can only happen once
             do
               atomically $ readTVar stateVar >>= \case
@@ -103,14 +109,29 @@ unmanagedAsyncWithUnmask action = do
                 TaskStateThrowing -> retry -- Could not disarm so we have to wait for the exception to arrive
                 TaskStateCompleted -> pure ()
 
+          catchAll
+            case result of
+              Left ex -> when (fromException ex /= Just AsyncDisposed) $ handler ex
+              Right _ -> pure ()
+            \ex -> undefined
+
           atomically do
             putAsyncVarEitherSTM_ resultVar result
             defaultRunFinalizers finalizers
 
-
     atomically $ writeTVar stateVar $ TaskStateRunning threadId
 
     pure $ Task key stateVar finalizers (toAwaitable resultVar)
 
+
+unmanagedAsync :: MonadIO m => IO a -> m (Task a)
+unmanagedAsync action = unmanagedAsyncWithUnmask \unmask -> unmask action
+
+unmanagedAsync_ :: MonadIO m => IO () -> m Disposable
+unmanagedAsync_ action = toDisposable <$> unmanagedAsync action
+
+unmanagedAsyncWithUnmask :: MonadIO m => ((forall b. IO b -> IO b) -> IO a) -> m (Task a)
+unmanagedAsyncWithUnmask = coreAsyncImplementation (traceIO . ("Unhandled exception in unmanaged async: " <>) . displayException)
+
 unmanagedAsyncWithUnmask_ :: MonadIO m => ((forall b. IO b -> IO b) -> IO ()) -> m Disposable
 unmanagedAsyncWithUnmask_ action = toDisposable <$> unmanagedAsyncWithUnmask action

src/Quasar/Timer.hs

@@ -204,7 +204,7 @@ newtype Delay = Delay (Task ())
   deriving newtype IsDisposable
 
 instance IsAwaitable () Delay where
-  toAwaitable (Delay task) = toAwaitable task `catch` \TaskDisposed -> throwM TimerCancelled
+  toAwaitable (Delay task) = toAwaitable task `catch` \AsyncDisposed -> throwM TimerCancelled
 
 newDelay :: MonadResourceManager m => Int -> m Delay
 newDelay microseconds = registerNewResource $ newUnmanagedDelay microseconds