Guard against leaking threads while restarting shake

ghcide-test/exe/DiagnosticTests.hs

@@ -15,7 +15,8 @@ import           Development.IDE.Test            (diagnostic,
                                                   expectDiagnostics,
                                                   expectDiagnosticsWithTags,
                                                   expectNoMoreDiagnostics,
-                                                  flushMessages, waitForAction)
+                                                  flushMessages, waitForAction,
+                                                  waitForBuildQueue)
 import           Development.IDE.Types.Location
 import qualified Language.LSP.Protocol.Lens      as L
 import           Language.LSP.Protocol.Message
@@ -54,6 +55,20 @@ tests = testGroup "diagnostics"
               }
       changeDoc doc [change]
       expectDiagnostics [("Testing.hs", [])]
+  , testWithDummyPluginEmpty "rapid edits then save does not strand a stale diagnostic" $ do
+      let v rhs = T.unlines ["module Testing where", "foo :: Int", "foo = " <> rhs]
+          whole rhs = TextDocumentContentChangeEvent . InR . TextDocumentContentChangeWholeDocument $ v rhs
+      doc <- createDoc "Testing.hs" "haskell" (v "()")
+      expectDiagnostics [("Testing.hs", [(DiagnosticSeverity_Error, (2, 6), "Couldn't match expected type 'Int' with actual type '()'", Just "GHC-83865")])]
+      changeDoc doc [whole "()"]
+      changeDoc doc [whole "'a'"]
+      changeDoc doc [whole "True"]
+      changeDoc doc [whole "0"]
+      sendNotification SMethod_TextDocumentDidSave (DidSaveTextDocumentParams doc Nothing)
+      waitForBuildQueue
+      liftIO $ sleep 0.2
+      flushMessages
+      expectCurrentDiagnostics doc []
   , testWithDummyPluginEmpty "introduce syntax error" $ do
       let content = T.unlines [ "module Testing where" ]
       doc <- createDoc "Testing.hs" "haskell" content

hls-graph/hls-graph.cabal

@@ -129,6 +129,7 @@ test-suite tests
     -threaded -rtsopts -with-rtsopts=-N -fno-ignore-asserts
 
   build-depends:
+    , async
     , base
     , extra
     , hls-graph

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

@@ -8,7 +8,9 @@
 {-# LANGUAGE RecordWildCards    #-}
 {-# LANGUAGE TypeFamilies       #-}
 
-module Development.IDE.Graph.Internal.Database (compute, newDatabase, incDatabase, build, getDirtySet, getKeysAndVisitAge) where
+module Development.IDE.Graph.Internal.Database (compute, newDatabase, incDatabase, build, getDirtySet, getKeysAndVisitAge,
+    -- * Exposed for testing. See Note [Closing escaped rule computations].
+    registerAsyncs, cleanupAsync) where
 
 import           Prelude                              hiding (unzip)
 
@@ -300,27 +302,57 @@ transitiveDirtySet database = flip State.execStateT mempty . traverse_ loop
 -- Asynchronous computations with cancellation
 
 -- | A simple monad to implement cancellation on top of 'Async',
---   generalizing 'withAsync' to monadic scopes.
-newtype AIO a = AIO { unAIO :: ReaderT (IORef [Async ()]) IO a }
+-- generalizing 'withAsync' to monadic scopes.
+--
+-- See Note [Closing escaped rule computations].
+newtype AIO a = AIO { unAIO :: ReaderT (IORef (Maybe [Async ()])) IO a }
   deriving newtype (Applicative, Functor, Monad, MonadIO)
 
 -- | Run the monadic computation, cancelling all the spawned asyncs if an exception arises
 runAIO :: AIO a -> IO a
 runAIO (AIO act) = do
-    asyncs <- newIORef []
+    asyncs <- newIORef (Just [])
     runReaderT act asyncs `onException` cleanupAsync asyncs
 
+{- Note [Closing escaped rule computations]
+
+Rule computations run as asyncs inside a per-'build' AIO scope, which
+'cleanupAsync' drains when the scope ends. One kind of async escapes that drain.
+
+  - A memoized 'splitIO' thunk in a 'Running' status can be forced by a later
+    build.
+  - That force runs in the original scope that produced the thunk.
+  - If that scope has already ended, the async it spawns has no parent to cancel
+    it, so on a restart it escapes the step bump and leaks. See
+    https://github.com/haskell/haskell-language-server/issues/4985.
+
+We close the leak at registration rather than teardown. The scope registry is
+a 'Maybe'. 'cleanupAsync' sets it to 'Nothing', and 'registerAsyncs' refuses a
+closed scope and cancels it, so the thunk can't spawn a surviving thread.
+-}
+
+-- | Register asyncs into a scope, or, if the scope has already closed, cancel
+--   them and report failure. See Note [Closing escaped rule computations].
+registerAsyncs :: IORef (Maybe [Async ()]) -> [Async ()] -> IO Bool
+registerAsyncs st as = do
+    registered <- atomicModifyIORef' st $ \case
+        Nothing -> (Nothing, False)
+        Just xs -> (Just (as ++ xs), True)
+    unless registered $ mapM_ uninterruptibleCancel as
+    pure registered
+
 -- | Like 'async' but with built-in cancellation.
 --   Returns an IO action to wait on the result.
+--   See Note [Closing escaped rule computations].
 asyncWithCleanUp :: AIO a -> AIO (IO a)
 asyncWithCleanUp act = do
     st <- AIO ask
     io <- unliftAIO act
     -- mask to make sure we keep track of the spawned async
     liftIO $ uninterruptibleMask $ \restore -> do
         a <- async $ restore io
-        atomicModifyIORef'_ st (void a :)
-        return $ wait a
+        registered <- registerAsyncs st [void a]
+        return $ if registered then wait a else throwIO AsyncCancelled
 
 unliftAIO :: AIO a -> AIO (IO a)
 unliftAIO act = do
@@ -334,10 +366,12 @@ withRunInIO k = do
     st <- AIO ask
     k $ RunInIO (\aio -> runReaderT (unAIO aio) st)
 
-cleanupAsync :: IORef [Async a] -> IO ()
+cleanupAsync :: IORef (Maybe [Async a]) -> IO ()
 -- mask to make sure we interrupt all the asyncs
 cleanupAsync ref = uninterruptibleMask $ \unmask -> do
-    asyncs <- atomicModifyIORef' ref ([],)
+    -- Close the scope so no later force can register, and take the live asyncs
+    -- to interrupt. See Note [Closing escaped rule computations].
+    asyncs <- atomicModifyIORef' ref $ \m -> (Nothing, fromMaybe [] m)
     -- interrupt all the asyncs without waiting
     mapM_ (\a -> throwTo (asyncThreadId a) AsyncCancelled) asyncs
     -- Wait until all the asyncs are done
@@ -368,12 +402,16 @@ waitConcurrently_ many = do
     ref <- AIO ask
     -- spawn the async computations.
     -- mask to make sure we keep track of all the asyncs.
-    (asyncs, syncs) <- liftIO $ uninterruptibleMask $ \unmask -> do
+    (asyncs, syncs, registered) <- liftIO $ uninterruptibleMask $ \unmask -> do
         waits <- liftIO $ traverse (waitOrSpawn . fmapWait unmask) many
         let (syncs, asyncs) = partitionEithers waits
-        liftIO $ atomicModifyIORef'_ ref (asyncs ++)
-        return (asyncs, syncs)
-    -- work on the sync computations
-    liftIO $ sequence_ syncs
-    -- wait for the async computations before returning
-    liftIO $ traverse_ wait asyncs
+        registered <- liftIO $ registerAsyncs ref asyncs
+        return (asyncs, syncs, registered)
+    -- A closed scope means our asyncs were cancelled, so abort the superseded work.
+    if not registered
+      then liftIO $ throwIO AsyncCancelled
+      else do
+        -- work on the sync computations
+        liftIO $ sequence_ syncs
+        -- wait for the async computations before returning
+        liftIO $ traverse_ wait asyncs

hls-graph/test/ActionSpec.hs

@@ -5,13 +5,18 @@ module ActionSpec where
 
 import           Control.Concurrent                      (MVar, readMVar)
 import qualified Control.Concurrent                      as C
+import           Control.Concurrent.Async                (async, poll)
 import           Control.Concurrent.STM
 import           Control.Monad.IO.Class                  (MonadIO (..))
+import           Data.IORef                              (newIORef, readIORef)
+import           Data.Maybe                              (isJust, isNothing)
 import           Development.IDE.Graph                   (shakeOptions)
 import           Development.IDE.Graph.Database          (shakeNewDatabase,
                                                           shakeRunDatabase,
                                                           shakeRunDatabaseForKeys)
-import           Development.IDE.Graph.Internal.Database (build, incDatabase)
+import           Development.IDE.Graph.Internal.Database (build, cleanupAsync,
+                                                          incDatabase,
+                                                          registerAsyncs)
 import           Development.IDE.Graph.Internal.Key
 import           Development.IDE.Graph.Internal.Types
 import           Development.IDE.Graph.Rule
@@ -129,3 +134,18 @@ spec = do
     res `shouldBe` [[True]]
     Just (Clean res) <- lookup (newKey theKey) <$> getDatabaseValues theDb
     resultDeps res `shouldBe` UnknownDeps
+
+  describe "Closing escaped rule computations" $ do
+    it "tracks an async registered while the scope is open" $ do
+      scope <- newIORef (Just [])
+      a <- async $ C.threadDelay maxBound
+      registerAsyncs scope [a] `shouldReturn` True
+    it "refuses to register into a closed scope and cancels the late async" $ do
+      scope <- newIORef (Just [])
+      cleanupAsync scope
+      readIORef scope >>= \m -> isNothing m `shouldBe` True
+      late <- async $ C.threadDelay maxBound
+      -- Registration fails rather than leaking the async past teardown.
+      registerAsyncs scope [late] `shouldReturn` False
+      -- The refused async was cancelled, not left running.
+      poll late >>= \res -> isJust res `shouldBe` True