From cf3b74dba275cfbc13accc1361201d62d59aced3 Mon Sep 17 00:00:00 2001
From: Jens Nolte <git@queezle.net>
Date: Sat, 12 Feb 2022 00:56:37 +0100
Subject: [PATCH] Implement registerNewResource for new resource implementation

---
 src/Quasar/Awaitable.hs  |  6 ++++++
 src/Quasar/Exceptions.hs | 31 ++++++++++++++++++++++++++++-
 src/Quasar/Resources.hs  | 43 +++++++++++++++++++++++++++++++++++-----
 3 files changed, 74 insertions(+), 6 deletions(-)

diff --git a/src/Quasar/Awaitable.hs b/src/Quasar/Awaitable.hs
index ddbd649..47dd6b7 100644
--- a/src/Quasar/Awaitable.hs
+++ b/src/Quasar/Awaitable.hs
@@ -2,6 +2,7 @@ module Quasar.Awaitable (
   -- * MonadAwaitable
   MonadAwait(..),
   peekAwaitable,
+  peekAwaitableSTM,
   awaitSTM,
 
   -- * Awaitable
@@ -105,6 +106,11 @@ instance MonadAwait m => MonadAwait (MaybeT m) where
 peekAwaitable :: MonadIO m => Awaitable r -> m (Maybe r)
 peekAwaitable awaitable = liftIO $ atomically $ (Just <$> awaitSTM awaitable) `orElse` pure Nothing
 
+-- | Returns the result (in a `Just`) when the awaitable is completed, throws an `Exception` when the awaitable is
+-- failed and returns `Nothing` otherwise.
+peekAwaitableSTM :: Awaitable r -> STM (Maybe r)
+peekAwaitableSTM awaitable = (Just <$> awaitSTM awaitable) `orElse` pure Nothing
+
 
 
 class IsAwaitable r a | a -> r where
diff --git a/src/Quasar/Exceptions.hs b/src/Quasar/Exceptions.hs
index 1c2fe45..572d484 100644
--- a/src/Quasar/Exceptions.hs
+++ b/src/Quasar/Exceptions.hs
@@ -12,6 +12,10 @@ module Quasar.Exceptions (
   isAsyncDisposed,
   DisposeException(..),
   isDisposeException,
+  FailedToAttachResource(..),
+  isFailedToAttachResource,
+  AlreadyDisposing(..),
+  isAlreadyDisposing,
 ) where
 
 import Control.Concurrent.STM
@@ -62,7 +66,7 @@ isAsyncDisposed (fromException @AsyncDisposed -> Just _) = True
 isAsyncDisposed _ = False
 
 
-  
+
 data DisposeException = DisposeException SomeException
   deriving stock Show
 
@@ -72,3 +76,28 @@ instance Exception DisposeException where
 isDisposeException :: SomeException -> Bool
 isDisposeException (fromException @DisposeException -> Just _) = True
 isDisposeException _ = False
+
+
+
+data FailedToAttachResource = FailedToAttachResource
+  deriving stock (Eq, Show)
+
+instance Exception FailedToAttachResource where
+  displayException FailedToAttachResource =
+    "FailedToRegisterResource: Failed to attach a resource to a resource manager. This might result in leaked resources if left unhandled."
+
+isFailedToAttachResource :: SomeException -> Bool
+isFailedToAttachResource (fromException @FailedToAttachResource -> Just _) = True
+isFailedToAttachResource _ = False
+
+
+data AlreadyDisposing = AlreadyDisposing
+  deriving stock (Eq, Show)
+
+instance Exception AlreadyDisposing where
+  displayException AlreadyDisposing =
+    "AlreadyDisposing: Failed to create a resource because the resource manager it should be attached to is already disposing."
+
+isAlreadyDisposing :: SomeException -> Bool
+isAlreadyDisposing (fromException @AlreadyDisposing -> Just _) = True
+isAlreadyDisposing _ = False
diff --git a/src/Quasar/Resources.hs b/src/Quasar/Resources.hs
index b6c8e73..3bf9dca 100644
--- a/src/Quasar/Resources.hs
+++ b/src/Quasar/Resources.hs
@@ -2,21 +2,28 @@ module Quasar.Resources (
   -- * Resources
   Resource(..),
   dispose,
-  disposeEventuallySTM,
-  disposeEventuallySTM_,
+  isDisposing,
   isDisposed,
 
-  -- * Monadic resource management
+  -- * Resource management in the `Quasar` monad
   registerResource,
+  registerNewResource,
   registerDisposeAction,
   registerDisposeTransaction,
+  disposeEventually,
+  disposeEventually_,
+
+  -- * STM
+  disposeEventuallySTM,
+  disposeEventuallySTM_,
 
-  -- * Disposer
+  -- * Types to implement resources
+  -- ** Disposer
   Disposer,
   newIODisposer,
   newSTMDisposer,
 
-  -- * Resource manager
+  -- ** Resource manager
   ResourceManager,
   newResourceManagerSTM,
   attachResource,
@@ -24,6 +31,8 @@ module Quasar.Resources (
 
 
 import Control.Concurrent.STM
+import Control.Monad.Catch
+import Quasar.Awaitable
 import Quasar.Async.STMHelper
 import Quasar.Exceptions
 import Quasar.Monad
@@ -56,3 +65,27 @@ registerDisposeTransaction fn = do
   exChan <- askExceptionChannel
   rm <- askResourceManager
   runSTM $ attachResource rm =<< newSTMDisposer worker exChan fn
+
+registerNewResource :: forall a m. (Resource a, MonadQuasar m) => m a -> m a
+registerNewResource fn = do
+  rm <- askResourceManager
+  disposing <- isJust <$> runSTM (peekAwaitableSTM (isDisposing rm))
+  -- Bail out before creating the resource _if possible_
+  when disposing $ throwM AlreadyDisposing
+
+  maskIfRequired do
+    resource <- fn
+    registerResource resource `catchAll` \ex -> do
+      -- When the resource cannot be registered (because resource manager is now disposing), destroy it to prevent leaks
+      disposeEventually_ resource
+      case ex of
+        (fromException -> Just FailedToAttachResource) -> throwM AlreadyDisposing
+        _ -> throwM ex
+    pure resource
+
+
+disposeEventually :: (Resource r, MonadQuasar m) => r -> m (Awaitable ())
+disposeEventually res = runSTM $ disposeEventuallySTM res
+
+disposeEventually_ :: (Resource r, MonadQuasar m) => r -> m ()
+disposeEventually_ res = runSTM $ disposeEventuallySTM_ res
-- 
GitLab