Guard against leaking threads while restarting Shake

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

@@ -124,7 +124,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 act
                     SMap.focus (updateStatus $ Running current force val s) id databaseValues
                     modifyTVar' toForce (Spawn force:)
                     pure val
@@ -172,13 +172,13 @@ refreshDeps visited db stack key result = \case
                     else refreshDeps newVisited db stack key result deps
 
 -- | Refresh a key:
-refresh :: Database -> Stack -> Key -> Maybe Result -> AIO (IO Result)
+refresh :: Database -> Stack -> Key -> Maybe Result -> AIO Result
 -- refresh _ st k _ | traceShow ("refresh", st, k) False = undefined
 refresh db stack key result = case (addStack key stack, result) of
     (Left e, _) -> throw e
-    (Right stack, Just me@Result{resultDeps = ResultDeps deps}) -> asyncWithCleanUp $ refreshDeps mempty db stack key me (reverse deps)
+    (Right stack, Just me@Result{resultDeps = ResultDeps deps}) -> refreshDeps mempty db stack key me (reverse deps)
     (Right stack, _) ->
-        asyncWithCleanUp $ liftIO $ compute db stack key RunDependenciesChanged result
+        liftIO $ compute db stack key RunDependenciesChanged result
 
 -- | Compute a key.
 compute :: Database -> Stack -> Key -> RunMode -> Maybe Result -> IO Result
@@ -300,32 +300,72 @@ 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
 
--- | Like 'async' but with built-in cancellation.
---   Returns an IO action to wait on the result.
-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
-
-unliftAIO :: AIO a -> AIO (IO a)
-unliftAIO act = do
-    st <- AIO ask
-    return $ runReaderT (unAIO act) st
+{- Note [Closing escaped rule computations]
+
+Rule computations run as asyncs inside a per-'build' AIO scope, which
+'cleanupAsync' drains when the scope is cancelled. A force-runner may be spawned
+while that scope is open but still be parked before registration when cleanup
+closes the scope. Without an admission gate, that runner can execute a memoized
+'splitIO' force after cleanup has already drained the registry. Once the build
+step advances, 'viewDirty' prevents a new build from deliberately selecting the
+old 'Running' force; the race is with a runner that was already spawned.
+
+The admission protocol gives every thread exactly one owner:
+  - A registered thread may execute its force or rule body, and is cancelled by
+    'cleanupAsync'.
+  - A refused thread observes the closed admission, skips its body, and finishes
+    normally. It is never registered, so cleanup does not cancel it.
+  - A registered parent that observes a refused child waits at an interruptible
+    point for cleanup to cancel it. It does not cancel itself or the child.
+
+This prevents either a force or rule body from escaping teardown. See
+https://github.com/haskell/haskell-language-server/issues/4985.
+-}
+
+-- | Register threads into a scope and report whether they were admitted.
+-- Refused threads are not owned by the scope; their admission gates must be
+-- released with 'False' so they can finish without executing their bodies.
+--
+-- See Note [Closing escaped rule computations].
+registerAsyncs :: IORef (Maybe [Async ()]) -> [Async ()] -> IO Bool
+registerAsyncs st as =
+    atomicModifyIORef' st $ \case
+        Nothing -> (Nothing, False)
+        Just xs -> (Just (as ++ xs), True)
+
+-- | Spawn a parked thread. An admitted thread runs the body and returns its
+-- result in 'Just'; a refused thread skips the body and finishes with 'Nothing'.
+runAsyncIfRegistered :: MVar Bool -> IO a -> IO (Async (Maybe a))
+runAsyncIfRegistered gate io =
+  asyncWithUnmask $ \unmask -> unmask $ do
+    registered <- readMVar gate
+    -- Only if the thread was successfully registered do we run the computation.
+    if registered then Just <$> io else pure Nothing
+
+-- | Wait for a child's admission result. A refused child has finished without
+-- doing work, while its already-admitted parent remains owned by cleanup and
+-- must wait for that external cancellation rather than cancelling itself.
+waitForRegisteredAsync :: Async (Maybe a) -> IO a
+waitForRegisteredAsync a = wait a >>= maybe waitForScopeCancellation pure
+
+-- This loop is deliberately interruptible. In the rejected-child path the
+-- current thread is already registered, so cleanup is its sole canceller.
+waitForScopeCancellation :: IO a
+waitForScopeCancellation = forever $ do
+    allowInterrupt
+    sleep 3600
 
 newtype RunInIO = RunInIO (forall a. AIO a -> IO a)
 
@@ -334,46 +374,48 @@ 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 ([],)
-    -- 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
+    -- 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)
+    -- Cancel the asyncs concurrently. If teardown takes over 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
+            mapConcurrently_ cancel asyncs
 
 data Wait
     = Wait {justWait :: !(IO ())}
     | Spawn {justWait :: !(IO ())}
 
-fmapWait :: (IO () -> IO ()) -> Wait -> Wait
-fmapWait f (Wait io)  = Wait (f io)
-fmapWait f (Spawn io) = Spawn (f io)
-
-waitOrSpawn :: Wait -> IO (Either (IO ()) (Async ()))
-waitOrSpawn (Wait io)  = pure $ Left io
-waitOrSpawn (Spawn io) = Right <$> async io
+waitOrSpawn :: MVar Bool -> Wait -> IO (Either (IO ()) (Async (Maybe ())))
+waitOrSpawn _mv (Wait io) = pure $ Left io
+waitOrSpawn mv (Spawn io) = Right <$> runAsyncIfRegistered mv 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
+    -- Mask so spawn + register stay atomic.
+    (asyncs, syncs, registered) <- liftIO $ uninterruptibleMask_ $ do
+        gate <- newEmptyMVar
+        waits <- liftIO $ traverse (waitOrSpawn gate) 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 (map void asyncs)
+        putMVar gate registered
+        return (asyncs, syncs, registered)
+    -- Refused children finish without running their forces. This parent was
+    -- already admitted, so cleanup remains its sole canceller.
+    if not registered
+      then liftIO $ traverse_ wait asyncs >> waitForScopeCancellation
+      else do
+        -- work on the sync computations
+        liftIO $ sequence_ syncs
+        -- wait for the async computations before returning
+        liftIO $ traverse_ waitForRegisteredAsync asyncs

hls-graph/test/ActionSpec.hs

@@ -3,10 +3,14 @@
 
 module ActionSpec where
 
-import           Control.Concurrent                      (MVar, readMVar)
+import           Control.Concurrent                      (MVar, newEmptyMVar,
+                                                          readMVar)
 import qualified Control.Concurrent                      as C
+import           Control.Concurrent.Async                (cancel, withAsync)
 import           Control.Concurrent.STM
+import           Control.Exception                       (finally)
 import           Control.Monad.IO.Class                  (MonadIO (..))
+import           Data.Maybe                              (isJust, isNothing)
 import           Development.IDE.Graph                   (shakeOptions)
 import           Development.IDE.Graph.Database          (shakeNewDatabase,
                                                           shakeRunDatabase,
@@ -17,9 +21,20 @@
 import           Development.IDE.Graph.Rule
 import           Example
 import qualified StmContainers.Map                       as STM
+import           System.IO.Unsafe                        (unsafePerformIO)
+import           System.Timeout                          (timeout)
 import           Test.Hspec
 
 
+-- Park traversal after the preceding key has been handled. The test releases
+-- this key normally; cancellation never crosses this unsafe test-only barrier.
+{-# NOINLINE blockSecondKey #-}
+blockSecondKey :: MVar () -> MVar () -> Rule ()
+blockSecondKey reached release = unsafePerformIO $ do
+  C.putMVar reached ()
+  C.takeMVar release
+  pure Rule
+
 
 spec :: Spec
 spec = do
@@ -129,3 +144,51 @@
     res `shouldBe` [[True]]
     Just (Clean res) <- lookup (newKey theKey) <$> getDatabaseValues theDb
     resultDeps res `shouldBe` UnknownDeps
+
+  describe "Closing escaped rule computations" $
+    it "does not leave a late async alive outside its closed AIO scope" $ do
+      bodyStarted <- newEmptyMVar
+      releaseBody <- newEmptyMVar
+      bodyFinished <- newEmptyMVar
+      ShakeDatabase _ _ theDb@Database{..} <- shakeNewDatabase shakeOptions $
+        addRule $ \(Rule :: Rule ()) _old _mode -> do
+          liftIO $
+            (C.putMVar bodyStarted () >> C.takeMVar releaseBody)
+              `finally` C.putMVar bodyFinished ()
+          pure $ RunResult ChangedRecomputeDiff "" () (pure ())
+
+      -- The first build publishes a lazy 'Running' force but cannot reach its
+      -- forcing phase while traversing this infinite list.
+      withAsync (build theDb emptyStack $ repeat (Rule @())) $ \firstBuild -> do
+        atomically $ do
+          details <- STM.lookup (newKey (Rule @())) databaseValues
+          case details of
+            Just KeyDetails{keyStatus = Running{}} -> pure ()
+            _                                      -> retry
+
+        -- A concurrent build selects that Running force, then parks while the
+        -- first build still owns an open AIO scope.
+        reachedSecondKey <- newEmptyMVar
+        releaseSecondKey <- newEmptyMVar
+        withAsync
+          (build theDb emptyStack
+            [Rule @(), blockSecondKey reachedSecondKey releaseSecondKey]) $ \secondBuild -> do
+              C.takeMVar reachedSecondKey
+              -- Close the force's original scope before the already-running
+              -- second build proceeds to force what it selected.
+              cancel firstBuild
+              C.putMVar releaseSecondKey ()
+              started <- timeout 1000000 $ C.takeMVar bodyStarted
+              -- Cancelling the waiter must not reveal a child that survived
+              -- teardown in the already-closed first scope.
+              cancel secondBuild
+              finished <- C.tryReadMVar bodyFinished
+              let outOfScopeAlive = isJust started && isNothing finished
+              -- Let an upstream orphan finish before reporting RED, so the
+              -- regression test itself never leaks a thread.
+              case (started, finished) of
+                (Just (), Nothing) -> do
+                  C.putMVar releaseBody ()
+                  C.takeMVar bodyFinished
+                _ -> pure ()
+              outOfScopeAlive `shouldBe` False