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/Action.hs

@@ -84,6 +84,8 @@ isAsyncException e
     | Just (_ :: AsyncCancelled) <- fromException e = True
     | Just (_ :: AsyncException) <- fromException e = True
     | Just (_ :: ExitCode) <- fromException e = True
+    -- See Note [Closing escaped rule computations].
+    | Just (_ :: ScopeClosed) <- fromException e = True
     | otherwise = False
 
 

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

@@ -8,7 +8,20 @@
 {-# 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
+  , Scope
+  , newScope
+  , scopeSize
+  , spawnInScope
+  , cleanupAsync
+  ) where
 
 import           Prelude                              hiding (unzip)
 
@@ -88,14 +101,17 @@ build
     => Database -> Stack -> f key -> IO (f Key, f value)
 -- build _ st k | traceShow ("build", st, k) False = undefined
 build db stack keys = do
-    built <- runAIO $ do
-        built <- builder db stack (fmap newKey keys)
-        case built of
-          Left clean  -> return clean
-          Right dirty -> liftIO dirty
+    -- 'repairRefusal' demoted the key before the refusal escaped, so the retry
+    -- recomputes it in a live scope. See Note [Closing escaped rule computations].
+    built <- attempt `catch` \ScopeClosed -> attempt
     let (ids, vs) = unzip built
     pure (ids, fmap (asV . resultValue) vs)
     where
+        attempt = runAIO $ do
+            built <- builder db stack (fmap newKey keys)
+            case built of
+              Left clean  -> return clean
+              Right dirty -> liftIO dirty
         asV :: Value -> value
         asV (Value x) = unwrapDynamic x
 
@@ -124,7 +140,7 @@ builder db@Database{..} stack keys = withRunInIO $ \(RunInIO run) -> do
                     pure val
                 Dirty s -> do
                     let act = run (refresh db stack id s)
-                        (force, val) = splitIO (join act)
+                        (force, val) = splitIO (repairRefusal db current id (join act))
                     SMap.focus (updateStatus $ Running current force val s) id databaseValues
                     modifyTVar' toForce (Spawn force:)
                     pure val
@@ -227,6 +243,27 @@ updateStatus res = Focus.alter
     (Just . maybe (KeyDetails res mempty)
     (\it -> it{keyStatus = res}))
 
+-- | Drop a refused key's 'Running' status before rethrowing, so its next build
+-- recomputes it in a live scope. Refusing the spawn stops the thread leaking.
+-- The entry it leaves behind is bound to the dead scope, so every later force
+-- in this step refuses too.
+--
+-- See Note [Closing escaped rule computations].
+repairRefusal :: Database -> Step -> Key -> IO a -> IO a
+repairRefusal db step key act = act `catch` \e@ScopeClosed ->
+    -- Masking stops a second exception from skipping the repair and leaving the
+    -- poisoned entry installed.
+    mask_ $ do
+        -- By the time the repair runs, 'compute' may have written 'Clean' or a
+        -- newer step may have taken the key over.
+        let demote = Focus.adjust $ \it -> case keyStatus it of
+                Running s _ _ prev | s == step -> it{keyStatus = Dirty prev}
+                _                              -> it
+
+        atomicallyNamed "builder repair refused" $
+            SMap.focus demote key (databaseValues db)
+        throwIO e
+
 -- | Returns the set of dirty keys annotated with their age (in # of builds)
 getDirtySet :: Database -> IO [(Key, Int)]
 getDirtySet db = do
@@ -300,27 +337,90 @@ 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 Scope IO a }
   deriving newtype (Applicative, Functor, Monad, MonadIO)
 
+-- | The threads a scope owns, or 'Nothing' once it has closed.
+-- See Note [Closing escaped rule computations].
+data Scope = Scope
+    { scopeClosing :: !(IORef Bool)
+    , scopeAsyncs  :: !(MVar (Maybe [Async ()]))
+    }
+
+newScope :: IO Scope
+newScope = Scope <$> newIORef False <*> newMVar (Just [])
+
+-- | How many asyncs the scope owns, or 'Nothing' if it has closed.
+scopeSize :: Scope -> IO (Maybe Int)
+scopeSize = fmap (fmap length) . readMVar . scopeAsyncs
+
 -- | Run the monadic computation, cancelling all the spawned asyncs if an exception arises
 runAIO :: AIO a -> IO a
 runAIO (AIO act) = do
-    asyncs <- newIORef []
-    runReaderT act asyncs `onException` cleanupAsync asyncs
+    scope <- newScope
+    -- Close on every exit, so no later force can spawn into a scope that ended.
+    -- Cancel only on an exception, since a normal return has already waited on
+    -- everything it spawned.
+    (runReaderT act scope `onException` cleanupAsync scope)
+        `finally` void (closeScope scope)
+
+{- Note [Closing escaped rule computations]
+
+A 'Running' status memoizes a 'splitIO' thunk bound to the AIO scope that
+created it. Forcing that thunk after its scope has ended spawns an async with 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.
+
+How the thunk outlives its scope, all at one step S:
+  1. A build opens scope-1 and installs 'Running S' for the key.
+  2. That build throws before forcing the thunk, so scope-1 spawned nothing for
+     the key and its teardown finds nothing to cancel. The 'Running S' entry
+     survives.
+  3. A second build at step S opens scope-2, sees that entry, waits on the
+     thunk, and forces it on scope-2's thread.
+
+'spawnInScope' decides whether to spawn while holding the scope open,
+so a refused computation is never started and the forcing thread raises
+'ScopeClosed'.
+  * 'repairRefusal' demotes the key to 'Dirty' before re-throwing.
+  * 'build' retries once, so the demoted key recomputes in a live scope.
+  * 'isAsyncException' classes it async, so it isn't swallowed.
+-}
+
+-- | Spawn one async into the scope, or refuse if the scope has closed.
+--
+-- The wait is returned rather than run, so the lock is free before anyone
+-- blocks on a result.
+--
+-- See Note [Closing escaped rule computations].
+spawnInScope :: Scope -> IO a -> IO (IO a)
+spawnInScope Scope{scopeClosing, scopeAsyncs} io =
+    mask_ $ modifyMVar scopeAsyncs $ \case
+        Nothing -> pure (Nothing, refuse)
+        Just as -> do
+            closing <- readIORef scopeClosing
+            if closing
+                -- Teardown is on its way and will take this list.
+                then pure (Just as, refuse)
+                else do
+                    a <- asyncWithUnmask $ \unmask -> unmask io
+                    pure (Just (void a : as), wait a)
+  where
+    refuse = throwIO ScopeClosed
 
 -- | Like 'async' but with built-in cancellation.
---   Returns an IO action to wait on the result.
+-- 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
+    scope <- 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
+    liftIO $ spawnInScope scope io
 
 unliftAIO :: AIO a -> AIO (IO a)
 unliftAIO act = do
@@ -334,46 +434,47 @@ withRunInIO k = do
     st <- AIO ask
     k $ RunInIO (\aio -> runReaderT (unAIO aio) st)
 
-cleanupAsync :: IORef [Async a] -> IO ()
+-- | Close the scope so no later force can spawn into it, handing back whatever
+-- asyncs it still owns.
+closeScope :: Scope -> IO [Async ()]
+closeScope Scope{scopeClosing, scopeAsyncs} = do
+    writeIORef scopeClosing True
+    mask_ $ modifyMVar scopeAsyncs $ \m -> pure (Nothing, fromMaybe [] m)
+
+cleanupAsync :: Scope -> IO ()
 -- mask to make sure we interrupt all the asyncs
-cleanupAsync ref = uninterruptibleMask $ \unmask -> do
-    asyncs <- atomicModifyIORef' ref ([],)
-    -- interrupt all the asyncs without waiting
-    mapM_ (\a -> throwTo (asyncThreadId a) AsyncCancelled) asyncs
-    -- Wait until all the asyncs are done
-    -- But if it takes more than 10 seconds, log to stderr
-    unless (null asyncs) $ do
-        let warnIfTakingTooLong = unmask $ forever $ do
-                sleep 10
-                traceM "cleanupAsync: waiting for asyncs to finish"
-        withAsync warnIfTakingTooLong $ \_ ->
-            mapM_ waitCatch asyncs
+cleanupAsync scope = uninterruptibleMask $ \unmask -> do
+    let warnIfTakingTooLong = unmask $ forever $ do
+            sleep 10
+            traceM "cleanupAsync: waiting for asyncs to finish"
+    -- Armed before the close so a stall taking the lock is reported too.
+    withAsync warnIfTakingTooLong $ \_ -> do
+        asyncs <- closeScope scope
+        mapConcurrently_ cancel asyncs
 
 data Wait
-    = Wait {justWait :: !(IO ())}
-    | Spawn {justWait :: !(IO ())}
+    = Wait !(IO ())
+    | Spawn !(IO ())
 
-fmapWait :: (IO () -> IO ()) -> Wait -> Wait
-fmapWait f (Wait io)  = Wait (f io)
-fmapWait f (Spawn io) = Spawn (f io)
+partitionWaits :: [Wait] -> ([IO ()], [IO ()])
+partitionWaits = partitionEithers . map toEither
+  where
+    toEither (Wait io)  = Left io
+    toEither (Spawn io) = Right io
 
-waitOrSpawn :: Wait -> IO (Either (IO ()) (Async ()))
-waitOrSpawn (Wait io)  = pure $ Left io
-waitOrSpawn (Spawn io) = Right <$> async io
+justWait :: Wait -> IO ()
+justWait (Wait io)  = io
+justWait (Spawn io) = io
 
 waitConcurrently_ :: [Wait] -> AIO ()
 waitConcurrently_ [] = pure ()
-waitConcurrently_ [one] = liftIO $ justWait one
-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
-        waits <- liftIO $ traverse (waitOrSpawn . fmapWait unmask) many
-        let (syncs, asyncs) = partitionEithers waits
-        liftIO $ atomicModifyIORef'_ ref (asyncs ++)
-        return (asyncs, syncs)
-    -- work on the sync computations
+waitConcurrently_ [one] = liftIO $ justWait one -- Avoid spawning when only a single action
+waitConcurrently_ waits = do
+    scope <- AIO ask
+    let (syncs, spawns) = partitionWaits waits
+    waitAll <- liftIO $ case spawns of
+        []  -> pure $ pure ()
+        [s] -> spawnInScope scope s
+        ss  -> spawnInScope scope (mapConcurrently_ id ss)
     liftIO $ sequence_ syncs
-    -- wait for the async computations before returning
-    liftIO $ traverse_ wait asyncs
+    liftIO waitAll

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

@@ -256,6 +256,16 @@ instance Exception StackException where
     toException this@(StackException (Stack stack _)) = toException $
         GraphException (show$ last stack) (map show stack) this
 
+-- | A rule computation was asked to start in a scope that had already closed.
+-- Distinct from 'AsyncCancelled' so 'build' can retry a refusal in a fresh
+-- scope instead of unwinding.
+--
+-- See Note [Closing escaped rule computations].
+data ScopeClosed = ScopeClosed
+  deriving stock (Show)
+
+instance Exception ScopeClosed
+
 addStack :: Key -> Stack -> Either StackException Stack
 addStack k (Stack ks is)
     | k `memberKeySet` is = Left $ StackException stack2

hls-graph/test/ActionSpec.hs

@@ -5,13 +5,19 @@ module ActionSpec where
 
 import           Control.Concurrent                      (MVar, readMVar)
 import qualified Control.Concurrent                      as C
+import           Control.Concurrent.Async                (AsyncCancelled (..))
 import           Control.Concurrent.STM
 import           Control.Monad.IO.Class                  (MonadIO (..))
+import           Data.IORef                              (newIORef, readIORef,
+                                                          writeIORef)
 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, newScope,
+                                                          scopeSize,
+                                                          spawnInScope)
 import           Development.IDE.Graph.Internal.Key
 import           Development.IDE.Graph.Internal.Types
 import           Development.IDE.Graph.Rule
@@ -129,3 +135,36 @@ spec = do
     res `shouldBe` [[True]]
     Just (Clean res) <- lookup (newKey theKey) <$> getDatabaseValues theDb
     resultDeps res `shouldBe` UnknownDeps
+
+  describe "Closing escaped rule computations" $ do
+    it "runs a spawned body and cancels it at teardown" $ do
+      scope <- newScope
+      started <- C.newEmptyMVar
+      waitForIt <- spawnInScope scope $ do
+        C.putMVar started ()
+        C.threadDelay maxBound
+      scopeSize scope `shouldReturn` Just 1
+      -- The signal only arrives if the body really started running.
+      C.takeMVar started
+      cleanupAsync scope
+      scopeSize scope `shouldReturn` Nothing
+      waitForIt `shouldThrow` \(_ :: AsyncCancelled) -> True
+    it "spawns nothing into a closed scope" $ do
+      scope <- newScope
+      cleanupAsync scope
+      ran <- newIORef False
+      waitForIt <- spawnInScope scope $ writeIORef ran True
+      -- Still closed, so nothing was added behind teardown's back.
+      scopeSize scope `shouldReturn` Nothing
+      -- Nobody gets a result, and the body never ran.
+      waitForIt `shouldThrow` \ScopeClosed -> True
+      readIORef ran `shouldReturn` False
+    it "recomputes a key whose scope died before forcing it" $ do
+      (ShakeDatabase _ _ theDb) <- shakeNewDatabase shakeOptions ruleCycleAfterVictim
+      -- The cycle tears down the inner scope while 'CycleRule 1' sits 'Running'
+      -- with a thunk that scope never forced.
+      build theDb emptyStack [CycleRule 0] `shouldThrow` \StackException{} -> True
+      -- Same step, so the stale entry is still visible and gets waited on rather
+      -- than respawned.
+      res <- build theDb emptyStack [CycleRule 1]
+      snd res `shouldBe` [1 :: Int]