From e94dc46b4baa23d384f6e60415912539a9c32794 Mon Sep 17 00:00:00 2001
From: Jens Nolte <git@queezle.net>
Date: Thu, 25 Nov 2021 19:34:50 +0100
Subject: [PATCH] Add STM variants of sone resource management functions

---
 src/Quasar/Async.hs           |  4 +-
 src/Quasar/Disposable.hs      | 17 ++++++---
 src/Quasar/ResourceManager.hs | 69 ++++++++++++++++++++++-------------
 src/Quasar/Timer.hs           | 14 +++----
 4 files changed, 64 insertions(+), 40 deletions(-)

diff --git a/src/Quasar/Async.hs b/src/Quasar/Async.hs
index f4d2e2c..1419123 100644
--- a/src/Quasar/Async.hs
+++ b/src/Quasar/Async.hs
@@ -27,7 +27,7 @@ import Quasar.ResourceManager
 -- | TODO: Documentation
 --
 -- The action will be run with asynchronous exceptions unmasked.
-async :: MonadResourceManager m => (ResourceManagerIO a) -> m (Awaitable a)
+async :: MonadResourceManager m => ResourceManagerIO a -> m (Awaitable a)
 async action = asyncWithUnmask \unmask -> unmask action
 
 -- | TODO: Documentation
@@ -38,7 +38,7 @@ asyncWithUnmask action = do
   resourceManager <- askResourceManager
   asyncWithHandlerAndUnmask (throwToResourceManager resourceManager . AsyncException) action
 
-async_ :: MonadResourceManager m => (ResourceManagerIO ()) -> m ()
+async_ :: MonadResourceManager m => ResourceManagerIO () -> m ()
 async_ action = void $ async action
 
 asyncWithUnmask_ :: MonadResourceManager m => ((ResourceManagerIO a -> ResourceManagerIO a) -> ResourceManagerIO ()) -> m ()
diff --git a/src/Quasar/Disposable.hs b/src/Quasar/Disposable.hs
index 40a3105..a79a00d 100644
--- a/src/Quasar/Disposable.hs
+++ b/src/Quasar/Disposable.hs
@@ -23,6 +23,7 @@ module Quasar.Disposable (
   ResourceManagerResult(..),
   DisposableFinalizers,
   newDisposableFinalizers,
+  newDisposableFinalizersSTM,
   defaultRegisterFinalizer,
   defaultRunFinalizers,
   awaitResourceManagerResult,
@@ -35,6 +36,7 @@ import Control.Monad.Reader
 import Data.List.NonEmpty (nonEmpty)
 import Data.HashSet (HashSet)
 import Data.HashSet qualified as HashSet
+import GHC.Conc (unsafeIOToSTM)
 import Quasar.Awaitable
 import Quasar.Prelude
 import Quasar.Utils.Exceptions
@@ -65,6 +67,8 @@ class IsDisposable a where
   isDisposed = isDisposed . toDisposable
 
   -- | Finalizers MUST NOT throw exceptions.
+  --
+  -- The boolean returned by register finalizer indicates if the operation was successful.
   registerFinalizer :: a -> STM () -> STM Bool
   registerFinalizer = registerFinalizer . toDisposable
 
@@ -203,7 +207,7 @@ instance IsDisposable STMDisposable where
 -- when `dispose` is called multiple times).
 --
 -- The action must not block (retry) for an unbound time.
-newSTMDisposable :: MonadIO m => STM () -> m Disposable
+newSTMDisposable :: STM () -> STM Disposable
 newSTMDisposable disposeAction = toDisposable <$> newSTMDisposable' disposeAction
 
 -- | Create a new disposable from an STM action. Is is guaranteed, that the STM action will only be called once (even
@@ -213,10 +217,10 @@ newSTMDisposable disposeAction = toDisposable <$> newSTMDisposable' disposeActio
 --
 -- This variant of `newSTMDisposable` returns an unboxed `STMDisposable` which can be disposed from `STM` by using
 -- `disposeSTMDisposable`.
-newSTMDisposable' :: MonadIO m => STM () -> m STMDisposable
-newSTMDisposable' disposeAction = liftIO do
-  key <- newUnique
-  STMDisposable key <$> newTMVarIO disposeAction <*> newDisposableFinalizers <*> newAsyncVar
+newSTMDisposable' :: STM () -> STM STMDisposable
+newSTMDisposable' disposeAction = do
+  key <- unsafeIOToSTM newUnique
+  STMDisposable key <$> newTMVar disposeAction <*> newDisposableFinalizersSTM <*> newAsyncVarSTM
 
 disposeSTMDisposable :: STMDisposable -> STM ()
 disposeSTMDisposable (STMDisposable key actionVar finalizers resultVar) = do
@@ -250,6 +254,9 @@ newtype DisposableFinalizers = DisposableFinalizers (TMVar [STM ()])
 newDisposableFinalizers :: IO DisposableFinalizers
 newDisposableFinalizers = DisposableFinalizers <$> newTMVarIO []
 
+newDisposableFinalizersSTM :: STM DisposableFinalizers
+newDisposableFinalizersSTM = DisposableFinalizers <$> newTMVar []
+
 defaultRegisterFinalizer :: DisposableFinalizers -> STM () -> STM Bool
 defaultRegisterFinalizer (DisposableFinalizers finalizerVar) finalizer =
   tryTakeTMVar finalizerVar >>= \case
diff --git a/src/Quasar/ResourceManager.hs b/src/Quasar/ResourceManager.hs
index aefeb80..4445a8d 100644
--- a/src/Quasar/ResourceManager.hs
+++ b/src/Quasar/ResourceManager.hs
@@ -25,6 +25,7 @@ module Quasar.ResourceManager (
   IsResourceManager(..),
   ResourceManager,
   newResourceManager,
+  newResourceManagerSTM,
   attachDisposeAction,
   attachDisposeAction_,
 
@@ -48,12 +49,16 @@ import Data.HashMap.Strict qualified as HM
 import Data.List.NonEmpty (NonEmpty(..), (<|), nonEmpty)
 import Data.Sequence (Seq(..), (|>))
 import Data.Sequence qualified as Seq
+import GHC.Conc (unsafeIOToSTM)
 import Quasar.Async.Unmanaged
 import Quasar.Awaitable
 import Quasar.Disposable
 import Quasar.Prelude
 import Quasar.Utils.Exceptions
 
+-- TODO: Merge `DefaultResourceManager` and `RootResourceManager` as `ResourceManager`
+-- This allows to remove functions other than `toResourceManager` from the `IsResourceManager` class.
+
 
 data DisposeException = DisposeException SomeException
   deriving stock Show
@@ -85,6 +90,12 @@ class IsDisposable a => IsResourceManager a where
   attachDisposable :: (IsDisposable b, MonadIO m) => a -> b -> m ()
   attachDisposable self = attachDisposable (toResourceManager self)
 
+  -- | Attaches an `Disposable` to a ResourceManager. It will automatically be disposed when the resource manager is disposed.
+  --
+  -- May throw an `FailedToRegisterResource` if the resource manager is disposing/disposed.
+  attachDisposableSTM :: IsDisposable b => a -> b -> STM ()
+  attachDisposableSTM self = attachDisposableSTM (toResourceManager self)
+
   lockResourceManagerImpl :: (MonadIO m, MonadMask m) => a -> m b -> m b
   lockResourceManagerImpl self = lockResourceManagerImpl (toResourceManager self)
 
@@ -92,13 +103,14 @@ class IsDisposable a => IsResourceManager a where
   throwToResourceManager :: Exception e => a -> e -> IO ()
   throwToResourceManager = throwToResourceManager . toResourceManager
 
-  {-# MINIMAL toResourceManager | (attachDisposable, lockResourceManagerImpl, throwToResourceManager) #-}
+  {-# MINIMAL toResourceManager | (attachDisposable, attachDisposableSTM, lockResourceManagerImpl, throwToResourceManager) #-}
 
 
 data ResourceManager = forall a. IsResourceManager a => ResourceManager a
 instance IsResourceManager ResourceManager where
   toResourceManager = id
   attachDisposable (ResourceManager x) = attachDisposable x
+  attachDisposableSTM (ResourceManager x) = attachDisposableSTM x
   lockResourceManagerImpl (ResourceManager x) = lockResourceManagerImpl x
   throwToResourceManager (ResourceManager x) = throwToResourceManager x
 instance IsDisposable ResourceManager where
@@ -223,6 +235,7 @@ data RootResourceManager
 
 instance IsResourceManager RootResourceManager where
   attachDisposable (RootResourceManager internal _ _ _) = attachDisposable internal
+  attachDisposableSTM (RootResourceManager internal _ _ _) = attachDisposableSTM internal
   lockResourceManagerImpl (RootResourceManager internal _ _ _) = lockResourceManagerImpl internal
   throwToResourceManager (RootResourceManager _ _ exceptionsVar _) ex = do
     -- TODO only log exceptions after a timeout
@@ -254,7 +267,7 @@ newUnmanagedRootResourceManagerInternal = liftIO do
   mfix \root -> do
     -- TODO reevaluate if using unmanagedAsync and voiding the result is correct
     void $ unmanagedAsync (disposeThread root)
-    internal <- newUnmanagedDefaultResourceManagerInternal (toResourceManager root)
+    internal <- atomically $ newUnmanagedDefaultResourceManagerInternal (toResourceManager root)
     pure $ RootResourceManager internal disposingVar exceptionsVar finalExceptionsVar
 
   where
@@ -314,17 +327,17 @@ data ResourceManagerState
 instance IsResourceManager DefaultResourceManager where
   throwToResourceManager DefaultResourceManager{throwToHandler} = throwToHandler . toException
 
-  attachDisposable DefaultResourceManager{stateVar, disposablesVar} disposable = liftIO $ mask_ do
-    key <- newUnique
-    join $ atomically do
-      state <- readTVar stateVar
-      case state of
-        ResourceManagerNormal -> do
-          disposables <- takeTMVar disposablesVar
-          putTMVar disposablesVar (HM.insert key (toDisposable disposable) disposables)
-          registerFinalizer disposable (finalizer key)
-          pure $ pure @IO ()
-        _ -> pure $ throwM @IO FailedToRegisterResource
+  attachDisposable self disposable = liftIO $ atomically $ attachDisposableSTM self disposable
+
+  attachDisposableSTM DefaultResourceManager{stateVar, disposablesVar} disposable = do
+    key <- unsafeIOToSTM newUnique
+    state <- readTVar stateVar
+    case state of
+      ResourceManagerNormal -> do
+        disposables <- takeTMVar disposablesVar
+        putTMVar disposablesVar (HM.insert key (toDisposable disposable) disposables)
+        void $ registerFinalizer disposable (finalizer key)
+      _ -> throwM FailedToRegisterResource
     where
       finalizer :: Unique -> STM ()
       finalizer key =
@@ -445,18 +458,16 @@ defaultResourceManagerDisposeResult :: DefaultResourceManager -> DisposeResult
 defaultResourceManagerDisposeResult DefaultResourceManager{resourceManagerKey, resultVar} =
   DisposeResultResourceManager $ ResourceManagerResult resourceManagerKey $ join $ toAwaitable resultVar
 
-newUnmanagedDefaultResourceManager :: MonadIO m => ResourceManager -> m ResourceManager
-newUnmanagedDefaultResourceManager parentResourceManager = liftIO do
-  toResourceManager <$> newUnmanagedDefaultResourceManagerInternal parentResourceManager
-
-newUnmanagedDefaultResourceManagerInternal :: MonadIO m => ResourceManager -> m DefaultResourceManager
-newUnmanagedDefaultResourceManagerInternal parentResourceManager = liftIO do
-  resourceManagerKey <- newUnique
-  stateVar <- newTVarIO ResourceManagerNormal
-  disposablesVar <- newTMVarIO HM.empty
-  lockVar <- newTVarIO 0
-  finalizers <- newDisposableFinalizers
-  resultVar <- newAsyncVar
+-- | Internal constructor. The resulting resource manager is not attached to it's parent, which is required internally
+-- to implement the root resource manager.
+newUnmanagedDefaultResourceManagerInternal :: ResourceManager -> STM DefaultResourceManager
+newUnmanagedDefaultResourceManagerInternal parentResourceManager = do
+  resourceManagerKey <- unsafeIOToSTM newUnique
+  stateVar <- newTVar ResourceManagerNormal
+  disposablesVar <- newTMVar HM.empty
+  lockVar <- newTVar 0
+  finalizers <- newDisposableFinalizersSTM
+  resultVar <- newAsyncVarSTM
 
   pure DefaultResourceManager {
     resourceManagerKey,
@@ -471,10 +482,16 @@ newUnmanagedDefaultResourceManagerInternal parentResourceManager = liftIO do
 newResourceManager :: MonadResourceManager m => m ResourceManager
 newResourceManager = mask_ do
   parent <- askResourceManager
-  resourceManager <- newUnmanagedDefaultResourceManager parent
+  resourceManager <- liftIO $ atomically $ toResourceManager <$> newUnmanagedDefaultResourceManagerInternal parent
   registerDisposable resourceManager
   pure resourceManager
 
+newResourceManagerSTM :: ResourceManager -> STM ResourceManager
+newResourceManagerSTM parent = do
+  resourceManager <- toResourceManager <$> newUnmanagedDefaultResourceManagerInternal parent
+  attachDisposableSTM parent resourceManager
+  pure resourceManager
+
 
 -- * Utilities
 
diff --git a/src/Quasar/Timer.hs b/src/Quasar/Timer.hs
index ac626ca..caa0714 100644
--- a/src/Quasar/Timer.hs
+++ b/src/Quasar/Timer.hs
@@ -178,18 +178,18 @@ newUnmanagedTimer :: MonadIO m => TimerScheduler -> UTCTime -> m Timer
 newUnmanagedTimer scheduler time = liftIO do
   key <- newUnique
   completed <- newAsyncVar
-  disposable <- newSTMDisposable' do
-    cancelled <- failAsyncVarSTM completed TimerCancelled
-    when cancelled do
-      modifyTVar (activeCount scheduler) (+ (-1))
-      modifyTVar (cancelledCount scheduler) (+ 1)
-  let timer = Timer { key, time, completed, disposable, scheduler }
   atomically do
+    disposable <- newSTMDisposable' do
+      cancelled <- failAsyncVarSTM completed TimerCancelled
+      when cancelled do
+        modifyTVar (activeCount scheduler) (+ (-1))
+        modifyTVar (cancelledCount scheduler) (+ 1)
+    let timer = Timer { key, time, completed, disposable, scheduler }
     tryTakeTMVar (heap scheduler) >>= \case
       Just timers -> putTMVar (heap scheduler) (insert timer timers)
       Nothing -> throwM TimerSchedulerDisposed
     modifyTVar (activeCount scheduler) (+ 1)
-  pure timer
+    pure timer
 
 
 sleepUntil :: MonadIO m => TimerScheduler -> UTCTime -> m ()
-- 
GitLab