From 3f5dd23feeeb41ebac00ad099f465ad2ac945642 Mon Sep 17 00:00:00 2001
From: Jens Nolte <git@queezle.net>
Date: Tue, 31 Aug 2021 19:16:31 +0200
Subject: [PATCH] Add registerDisposable to MonadResourceManager

Co-authored-by: Jan Beinke <git@janbeinke.com>
---
 src/Quasar/Async.hs      |  6 +++---
 src/Quasar/Disposable.hs | 32 ++++++++++++++++++++++++++------
 src/Quasar/Observable.hs |  2 +-
 src/Quasar/Timer.hs      | 10 ++++------
 4 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/src/Quasar/Async.hs b/src/Quasar/Async.hs
index 4c013ac..42d892c 100644
--- a/src/Quasar/Async.hs
+++ b/src/Quasar/Async.hs
@@ -86,9 +86,9 @@ newtype UnlimitedAsync r = UnlimitedAsync { unUnlimitedAsync :: (ReaderT Resourc
 instance MonadAsync UnlimitedAsync where
   asyncWithUnmask action = do
     resourceManager <- askResourceManager
-    liftIO $ mask_ $ do
-      task <- forkTaskWithUnmask (\unmask -> runReaderT (unUnlimitedAsync (action (liftUnmask unmask))) resourceManager)
-      attachDisposable resourceManager task
+    mask_ $ do
+      task <- liftIO $ forkTaskWithUnmask (\unmask -> runReaderT (unUnlimitedAsync (action (liftUnmask unmask))) resourceManager)
+      registerDisposable task
       pure $ toAwaitable task
     where
       liftUnmask :: (forall b. IO b -> IO b) -> UnlimitedAsync a -> UnlimitedAsync a
diff --git a/src/Quasar/Disposable.hs b/src/Quasar/Disposable.hs
index 1494c18..d17a282 100644
--- a/src/Quasar/Disposable.hs
+++ b/src/Quasar/Disposable.hs
@@ -209,20 +209,40 @@ class HasResourceManager a where
 instance HasResourceManager ResourceManager where
   getResourceManager = id
 
-class (MonadMask m, MonadIO m) => MonadResourceManager m where
-  askResourceManager :: m ResourceManager
+class (MonadAwait m, MonadMask m, MonadIO m) => MonadResourceManager m where
+  registerDisposable :: IsDisposable a => a -> m ()
+
+  registerDisposeAction :: IO (Awaitable ()) -> m ()
+  registerDisposeAction disposeAction = mask_ $ registerDisposable =<< newDisposable disposeAction
+
   localResourceManager :: ResourceManager -> m a -> m a
 
-instance (MonadMask m, MonadIO m) => MonadResourceManager (ReaderT ResourceManager m) where
+  askResourceManager :: m ResourceManager
+
+  -- TODO askResourceManager could maybe be replaced with
+  --withRunResourceContextInIO :: (((forall f. MonadResourceManager f => f a) -> IO a) -> m b) -> m b
+
+
+instance (MonadAwait m, MonadMask m, MonadIO m) => MonadResourceManager (ReaderT ResourceManager m) where
+  registerDisposable disposable = do
+    resourceManager <- ask
+    attachDisposable resourceManager disposable
+
+  localResourceManager resourceManager = local (const resourceManager)
+
   askResourceManager = ask
-  localResourceManager = local . const
+
 
 instance {-# OVERLAPPABLE #-} MonadResourceManager m => MonadResourceManager (ReaderT r m) where
-  askResourceManager = lift askResourceManager
+  registerDisposable disposable = lift $ registerDisposable disposable
+
   localResourceManager resourceManager action = do
     x <- ask
     lift $ localResourceManager resourceManager $ runReaderT action x
 
+  askResourceManager = lift askResourceManager
+
+
 
 onResourceManager :: (HasResourceManager a) => a -> ReaderT ResourceManager m r -> m r
 onResourceManager target action = runReaderT action (getResourceManager target)
@@ -362,7 +382,7 @@ attachDisposable resourceManager disposable = liftIO $ mask \unmask -> do
 
 -- | Creates an `Disposable` that is bound to a ResourceManager. It will automatically be disposed when the resource manager is disposed.
 attachDisposeAction :: MonadIO m => ResourceManager -> IO (Awaitable ()) -> m Disposable
-attachDisposeAction resourceManager action = do
+attachDisposeAction resourceManager action = liftIO $ mask_ $ do
   disposable <- newDisposable action
   attachDisposable resourceManager disposable
   pure disposable
diff --git a/src/Quasar/Observable.hs b/src/Quasar/Observable.hs
index 5b302a3..7d2f82b 100644
--- a/src/Quasar/Observable.hs
+++ b/src/Quasar/Observable.hs
@@ -80,7 +80,7 @@ class IsRetrievable v a | a -> v where
 retrieveIO :: IsRetrievable v a => a -> IO v
 retrieveIO x = withOnResourceManager $ await =<< retrieve x
 
-type MonadObserve m = (MonadAwait m, MonadResourceManager m)
+type MonadObserve m = MonadResourceManager m
 
 {-# DEPRECATED unsafeAsyncObserveIO "Old implementation of `observe`." #-}
 class IsRetrievable v o => IsObservable v o | o -> v where
diff --git a/src/Quasar/Timer.hs b/src/Quasar/Timer.hs
index 76511d5..9030405 100644
--- a/src/Quasar/Timer.hs
+++ b/src/Quasar/Timer.hs
@@ -190,12 +190,10 @@ instance IsAwaitable () Delay where
   toAwaitable (Delay task) = toAwaitable task `catch` \TaskDisposed -> throwM TimerCancelled
 
 newDelay :: MonadResourceManager m => Int -> m Delay
-newDelay microseconds = do
-  resourceManager <- askResourceManager
-  mask_ do
-    delay <- Delay <$> forkTask (liftIO (threadDelay microseconds))
-    attachDisposable resourceManager delay
-    pure delay
+newDelay microseconds = mask_ do
+  delay <- Delay <$> forkTask (liftIO (threadDelay microseconds))
+  registerDisposable delay
+  pure delay
 
 
 
-- 
GitLab