Refactor async handling in Database module to use IORef instead of TVar

Author: soulomoon

hls-graph/src/Development/IDE/Graph/Internal/Database.hs

@@ -12,22 +12,13 @@ module Development.IDE.Graph.Internal.Database (compute, newDatabase, incDatabas
 
 import           Prelude                              hiding (unzip)
 
-import           Control.Concurrent.Async
-import           Control.Concurrent.Extra
-import           Control.Concurrent.STM.Stats         (STM, TVar, atomically,
-                                                       atomicallyNamed,
-                                                       modifyTVar', newTVarIO,
-                                                       readTVar, readTVarIO,
-                                                       retry)
-import           Control.Exception
+import           Control.Concurrent.STM.Stats         (atomicallyNamed, retry)
 import           Control.Monad
-import           Control.Monad.IO.Class               (MonadIO (liftIO))
 import           Control.Monad.Trans.Class            (lift)
 import           Control.Monad.Trans.Reader
 import qualified Control.Monad.Trans.State.Strict     as State
 import           Data.Dynamic
 import           Data.Foldable                        (for_, traverse_)
-import           Data.IORef.Extra
 import           Data.Maybe
 import           Data.Traversable                     (for)
 import           Data.Tuple.Extra
@@ -40,11 +31,14 @@ import qualified Focus
 import qualified ListT
 import qualified StmContainers.Map                    as SMap
 import           System.Time.Extra                    (duration, sleep)
-import           UnliftIO                             (MonadUnliftIO (withRunInIO))
+import           UnliftIO
 import qualified UnliftIO.Exception                   as UE
 
 #if MIN_VERSION_base(4,19,0)
+import           Control.Exception                    (throw)
 import           Data.Functor                         (unzip)
+import           Data.IORef.Extra                     (atomicModifyIORef'_)
+import           UnliftIO.Concurrent
 #else
 import           Data.List.NonEmpty                   (unzip)
 #endif
@@ -280,7 +274,7 @@ transitiveDirtySet database = flip State.execStateT mempty . traverse_ loop
 
 -- | A simple monad to implement cancellation on top of 'Async',
 --   generalizing 'withAsync' to monadic scopes.
-newtype AIO a = AIO { unAIO :: ReaderT (TVar [Async ()]) IO a }
+newtype AIO a = AIO { unAIO :: ReaderT (IORef [Async ()]) IO a }
   deriving newtype (Applicative, Functor, Monad, MonadIO)
 
 data AsyncParentKill = AsyncParentKill ThreadId Step
@@ -293,14 +287,11 @@ instance Exception AsyncParentKill where
 -- | Run the monadic computation, cancelling all the spawned asyncs if an exception arises
 runAIO :: Step -> AIO a -> IO a
 runAIO s (AIO act) = do
-    asyncsRef <- newTVarIO []
+    asyncsRef <- newIORef []
     -- Log the exact exception (including async exceptions) before cleanup,
     -- then rethrow to preserve previous semantics.
     runReaderT act asyncsRef `onException` do
-        asyncs <- atomically $ do
-            r <- readTVar asyncsRef
-            modifyTVar' asyncsRef $ const []
-            return r
+        asyncs <- atomicModifyIORef' asyncsRef ([],)
         tid <- myThreadId
         cleanupAsync asyncs tid s
 
@@ -313,7 +304,8 @@ asyncWithCleanUp act = do
     -- mask to make sure we keep track of the spawned async
     liftIO $ uninterruptibleMask $ \restore -> do
         a <- async $ restore io
-        atomically $ modifyTVar' st (void a :)
+        -- atomically $ modifyTVar' st (void a :)
+        atomicModifyIORef'_ st (void a:)
         return $ wait a
 
 unliftAIO :: AIO a -> AIO (IO a)