Add async variant that additionally accepts an exception handler

src/Quasar/Async.hs

@@ -5,6 +5,12 @@ module Quasar.Async (
   asyncWithUnmask,
   asyncWithUnmask_,
 
+  -- ** Async with explicit error handling
+  handleAsync,
+  handleAsync_,
+  handleAsyncWithUnmask,
+  handleAsyncWithUnmask_,
+
   -- ** Async exceptions
   CancelAsync(..),
   AsyncDisposed(..),
@@ -29,23 +35,39 @@ async action = asyncWithUnmask \unmask -> unmask action
 -- The action will be run with asynchronous exceptions masked and will be passed an action that can be used to unmask.
 asyncWithUnmask :: MonadResourceManager m => ((ResourceManagerIO a -> ResourceManagerIO a) -> ResourceManagerIO r) -> m (Awaitable r)
 asyncWithUnmask action = do
+  resourceManager <- askResourceManager
+  handleAsyncWithUnmask (throwToResourceManager resourceManager) action
+
+async_ :: MonadResourceManager m => (ResourceManagerIO ()) -> m ()
+async_ action = void $ async action
+
+asyncWithUnmask_ :: MonadResourceManager m => ((ResourceManagerIO a -> ResourceManagerIO a) -> ResourceManagerIO ()) -> m ()
+asyncWithUnmask_ action = void $ asyncWithUnmask action
+
+-- | TODO: Documentation
+--
+-- The action will be run with asynchronous exceptions unmasked. When an exception is thrown that is not caused from
+-- the disposable instance (i.e. the task being canceled), the handler is called with that exception.
+handleAsyncWithUnmask :: MonadResourceManager m => (SomeException -> IO ()) -> ((ResourceManagerIO a -> ResourceManagerIO a) -> ResourceManagerIO r) -> m (Awaitable r)
+handleAsyncWithUnmask handler action = do
   resourceManager <- askResourceManager
   toAwaitable <$> registerNewResource do
-    coreAsyncImplementation (handler resourceManager) \unmask ->
+    coreAsyncImplementation wrappedHandler \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
+    wrappedHandler :: SomeException -> IO ()
+    wrappedHandler (fromException -> Just AsyncDisposed) = pure ()
+    wrappedHandler ex = handler ex
     liftUnmask :: (forall b. IO b -> IO b) -> ResourceManagerIO a -> ResourceManagerIO a
     liftUnmask unmask innerAction = do
       resourceManager <- askResourceManager
       liftIO $ unmask $ onResourceManager resourceManager innerAction
 
-async_ :: MonadResourceManager m => (ResourceManagerIO ()) -> m ()
-async_ action = void $ async action
+handleAsyncWithUnmask_ :: MonadResourceManager m => (SomeException -> IO ()) -> ((ResourceManagerIO a -> ResourceManagerIO a) -> ResourceManagerIO r) -> m ()
+handleAsyncWithUnmask_ handler action = void $ handleAsyncWithUnmask handler action
 
-asyncWithUnmask_ :: MonadResourceManager m => ((ResourceManagerIO a -> ResourceManagerIO a) -> ResourceManagerIO ()) -> m ()
-asyncWithUnmask_ action = void $ asyncWithUnmask action
+handleAsync :: MonadResourceManager m => (SomeException -> IO ()) -> ResourceManagerIO r -> m (Awaitable r)
+handleAsync handler action = handleAsyncWithUnmask handler \unmask -> unmask action
+
+handleAsync_ :: MonadResourceManager m => (SomeException -> IO ()) -> ResourceManagerIO r -> m ()
+handleAsync_ handler action = void $ handleAsync handler action

src/Quasar/Async/Unmanaged.hs

@@ -99,21 +99,22 @@ coreAsyncImplementation handler action = do
 
           -- The `action` has completed its work.
           -- "disarm" dispose:
-          handleIf
+          catchIf
             do \(CancelAsync exKey) -> key == exKey
-            do mempty -- ignore exception if it matches; this can only happen once
             do
               atomically $ readTVar stateVar >>= \case
                 TaskStateInitializing -> retry
                 TaskStateRunning _ -> writeTVar stateVar TaskStateCompleted
                 TaskStateThrowing -> retry -- Could not disarm so we have to wait for the exception to arrive
                 TaskStateCompleted -> pure ()
+            do mempty -- ignore exception if it matches; this can only happen once (see TaskStateThrowing above)
 
           catchAll
             case result of
-              Left ex -> when (fromException ex /= Just AsyncDisposed) $ handler ex
-              Right _ -> pure ()
-            \ex -> undefined
+              Left (fromException -> Just AsyncDisposed) -> pure ()
+              Left ex -> handler ex
+              _ -> pure ()
+            \ex -> traceIO $ "An exception was thrown while handling an async exception: " <> displayException ex
 
           atomically do
             putAsyncVarEitherSTM_ resultVar result