Integrate TIOWorker into disposer (new resource management)

Co-authored-by: Jan Beinke <git@janbeinke.com>

src/Quasar/Resources.hs

@@ -5,11 +5,15 @@ module Quasar.Resources (
   dispose,
   disposeEventuallySTM,
   disposeEventuallySTM_,
+  isDisposed,
+  newPrimitiveDisposer,
 ) where
 
 
+import Control.Concurrent (forkIO)
 import Control.Concurrent.STM
 import Control.Monad.Catch
+import Quasar.Async.STMHelper
 import Quasar.Awaitable
 import Quasar.Exceptions
 import Quasar.Prelude
@@ -23,22 +27,22 @@ class Resource a where
 type DisposerState = TOnce DisposeFn (Awaitable ())
 
 data Disposer
-  = FnDisposer ExceptionChannel DisposerState Finalizers
+  = FnDisposer Unique TIOWorker ExceptionChannel DisposerState Finalizers
   | ResourceManagerDisposer ResourceManager
 
-data DisposeFn
-  = IODisposeFn (IO ())
-  | STMDisposeFn (STM ())
+type DisposeFn = IO (Awaitable ())
 
-newDisposer :: ExceptionChannel -> IO () -> STM Disposer
-newDisposer exChan disposeFn = newFnDisposer exChan (IODisposeFn disposeFn)
+newShortDisposer :: TIOWorker -> ExceptionChannel -> IO () -> STM Disposer
+newShortDisposer worker exChan disposeFn = newPrimitiveDisposer worker exChan (pure <$> disposeFn)
 
-newSTMDisposer :: ExceptionChannel -> STM () -> STM Disposer
-newSTMDisposer exChan disposeFn = newFnDisposer exChan (STMDisposeFn disposeFn)
+newShortSTMDisposer :: TIOWorker -> ExceptionChannel -> STM () -> STM Disposer
+newShortSTMDisposer worker exChan disposeFn = newShortDisposer worker exChan (atomically disposeFn)
 
-newFnDisposer :: ExceptionChannel -> DisposeFn -> STM Disposer
-newFnDisposer exChan fn =
-  FnDisposer exChan <$> newTOnce fn <*> newFinalizersSTM
+-- TODO document: IO has to be "short"
+newPrimitiveDisposer :: TIOWorker -> ExceptionChannel -> IO (Awaitable ()) -> STM Disposer
+newPrimitiveDisposer worker exChan fn = do
+  key <- newUniqueSTM
+  FnDisposer key worker exChan <$> newTOnce fn <*> newFinalizers
 
 
 dispose :: (MonadIO m, Resource r) => r -> m ()
@@ -47,10 +51,9 @@ dispose resource = liftIO $ await =<< atomically (disposeEventuallySTM resource)
 disposeEventuallySTM :: Resource r => r -> STM (Awaitable ())
 disposeEventuallySTM resource =
   case getDisposer resource of
-    FnDisposer channel state finalizers -> do
-      beginDispose channel state finalizers
-    ResourceManagerDisposer resourceManager ->
-      beginDisposeResourceManager resourceManager
+    FnDisposer _ worker exChan state finalizers -> do
+      beginDisposeFnDisposer worker exChan state finalizers
+    ResourceManagerDisposer resourceManager -> undefined
 
 disposeEventuallySTM_ :: Resource r => r -> STM ()
 disposeEventuallySTM_ resource = void $ disposeEventuallySTM resource
@@ -63,13 +66,28 @@ isDisposed resource =
     ResourceManagerDisposer _resourceManager -> undefined -- resource manager
 
 
-beginDispose :: ExceptionChannel -> DisposerState -> Finalizers -> STM (Awaitable ())
-beginDispose channel disposeState finalizers =
+beginDisposeFnDisposer :: TIOWorker -> ExceptionChannel -> DisposerState -> Finalizers -> STM (Awaitable ())
+beginDisposeFnDisposer worker exChan disposeState finalizers =
   mapFinalizeTOnce disposeState startDisposeFn
   where
     startDisposeFn :: DisposeFn -> STM (Awaitable ())
-    startDisposeFn = undefined -- launch dispose thread
-
+    startDisposeFn disposeFn = do
+      awaitableVar <- newAsyncVarSTM
+      startTrivialIO_ worker exChan (runDisposeFn awaitableVar disposeFn)
+      pure $ join (toAwaitable awaitableVar)
+
+    runDisposeFn :: AsyncVar (Awaitable ()) -> DisposeFn -> IO ()
+    runDisposeFn awaitableVar disposeFn = mask_ $ handleAll exceptionHandler do
+      awaitable <- disposeFn 
+      putAsyncVar_ awaitableVar awaitable
+      runFinalizersAfter finalizers awaitable
+      where
+        exceptionHandler :: SomeException -> IO ()
+        exceptionHandler ex = do
+          -- In case of an exception mark disposable as completed to prevent resource managers from being stuck indefinitely
+          putAsyncVar_ awaitableVar (pure ())
+          atomically $ runFinalizers finalizers
+          throwIO $ DisposeException ex
 
 
 data ResourceManager = ResourceManager
@@ -92,22 +110,32 @@ data ResourceManagerResult = ResourceManagerResult Unique (Awaitable [ResourceMa
 
 newtype Finalizers = Finalizers (TMVar [STM ()])
 
-newFinalizers :: IO Finalizers
-newFinalizers = Finalizers <$> newTMVarIO []
-
-newFinalizersSTM :: STM Finalizers
-newFinalizersSTM = Finalizers <$> newTMVar []
+newFinalizers :: STM Finalizers
+newFinalizers = do
+  Finalizers <$> newTMVar []
 
-defaultRegisterFinalizer :: Finalizers -> STM () -> STM Bool
-defaultRegisterFinalizer (Finalizers finalizerVar) finalizer =
+registerFinalizer :: Finalizers -> STM () -> STM Bool
+registerFinalizer (Finalizers finalizerVar) finalizer =
   tryTakeTMVar finalizerVar >>= \case
     Just finalizers -> do
       putTMVar finalizerVar (finalizer : finalizers)
       pure True
     Nothing -> pure False
 
-defaultRunFinalizers :: Finalizers -> STM ()
-defaultRunFinalizers (Finalizers finalizerVar) = do
+runFinalizers :: Finalizers -> STM ()
+runFinalizers (Finalizers finalizerVar) = do
   tryTakeTMVar finalizerVar >>= \case
     Just finalizers -> sequence_ finalizers
-    Nothing -> throwM $ userError "defaultRunFinalizers was called multiple times (it must only be run once)"
+    Nothing -> throwM $ userError "runFinalizers was called multiple times (it must only be run once)"
+
+runFinalizersAfter :: Finalizers -> Awaitable () -> IO ()
+runFinalizersAfter finalizers awaitable = do
+  -- Peek awaitable to ensure trivial disposables always run without forking
+  isCompleted <- isJust <$> peekAwaitable awaitable 
+  if isCompleted
+    then
+      atomically $ runFinalizers finalizers
+    else
+      void $ forkIO do
+        await awaitable
+        atomically $ runFinalizers finalizers