WIP: Ensure restarts cancel outstanding shake threads

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

@@ -17,7 +17,8 @@ import           Control.Concurrent.Extra
 import           Control.Concurrent.STM.Stats         (STM, atomically,
                                                        atomicallyNamed,
                                                        modifyTVar', newTVarIO,
-                                                       readTVarIO)
+                                                       readTVar, readTVarIO,
+                                                       stateTVar)
 import           Control.Exception
 import           Control.Monad
 import           Control.Monad.IO.Class               (MonadIO (liftIO))
@@ -27,6 +28,7 @@ import qualified Control.Monad.Trans.State.Strict     as State
 import           Data.Dynamic
 import           Data.Either
 import           Data.Foldable                        (for_, traverse_)
+import qualified Data.IntMap.Strict                   as IntMap
 import           Data.IORef.Extra
 import           Data.Maybe
 import           Data.Traversable                     (for)
@@ -52,6 +54,7 @@ import           Data.List.NonEmpty                   (unzip)
 newDatabase :: Dynamic -> TheRules -> IO Database
 newDatabase databaseExtra databaseRules = do
     databaseStep <- newTVarIO $ Step 0
+    databaseAsyncs <- newTVarIO (0, IntMap.empty)
     databaseValues <- atomically SMap.new
     pure Database{..}
 
@@ -60,6 +63,9 @@ newDatabase databaseExtra databaseRules = do
 incDatabase :: Database -> Maybe [Key] -> IO ()
 -- only some keys are dirty
 incDatabase db (Just kk) = do
+    -- Cancel threads that escaped the previous scope before the step bump.
+    -- See Note [Cancelling escaped rule computations].
+    cancelTrackedAsyncs db
     atomicallyNamed "incDatabase" $ modifyTVar'  (databaseStep db) $ \(Step i) -> Step $ i + 1
     transitiveDirtyKeys <- transitiveDirtySet db kk
     for_ (toListKeySet transitiveDirtyKeys) $ \k ->
@@ -70,6 +76,7 @@ incDatabase db (Just kk) = do
 
 -- all keys are dirty
 incDatabase db Nothing = do
+    cancelTrackedAsyncs db
     atomically $ modifyTVar'  (databaseStep db) $ \(Step i) -> Step $ i + 1
     let list = SMap.listT (databaseValues db)
     atomicallyNamed "incDatabase - all " $ flip ListT.traverse_ list $ \(k,_) ->
@@ -132,7 +139,7 @@ builder db@Database{..} stack keys = withRunInIO $ \(RunInIO run) -> do
             pure (id, val)
 
     toForceList <- liftIO $ readTVarIO toForce
-    let waitAll = run $ waitConcurrently_ toForceList
+    let waitAll = run $ waitConcurrently_ db toForceList
     case toForceList of
         [] -> return $ Left results
         _ -> return $ Right $ do
@@ -176,9 +183,9 @@ refresh :: Database -> Stack -> Key -> Maybe Result -> AIO (IO 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}) -> asyncWithCleanUp db $ refreshDeps mempty db stack key me (reverse deps)
     (Right stack, _) ->
-        asyncWithCleanUp $ liftIO $ compute db stack key RunDependenciesChanged result
+        asyncWithCleanUp db $ liftIO $ compute db stack key RunDependenciesChanged result
 
 -- | Compute a key.
 compute :: Database -> Stack -> Key -> RunMode -> Maybe Result -> IO Result
@@ -310,15 +317,63 @@ runAIO (AIO act) = do
     asyncs <- newIORef []
     runReaderT act asyncs `onException` cleanupAsync asyncs
 
+{- Note [Cancelling escaped rule computations]
+
+Rule computations run as asyncs inside a per-'build' AIO scope, and that scope's
+'cleanupAsync' drains them when it ends. Two kinds of thread escape that drain:
+
+  - one registered after the scope had already drained,
+  - one forced by a later build.
+
+On a restart 'incDatabase' bumps the step, and an escaped thread left running can
+read across the bump or leak.
+
+'trackedAsync' registers every spawned async on the 'Database', keyed by a
+monotonic id and self-removing when it finishes. The registry outlives any
+single scope, so 'cancelTrackedAsyncs' can reach and kill every escaped thread
+before the next step bump.
+-}
+
+-- | Spawn an async and register it on the 'Database' for its whole lifetime,
+--   self-removing when it finishes.
+--   See Note [Cancelling escaped rule computations].
+trackedAsync :: Database -> IO a -> IO (Async a)
+trackedAsync Database{..} act = mask_ $ do
+    -- Park the child until it is registered, so its self-removal 'finally'
+    -- cannot fire before its entry exists.
+    gate <- newEmptyMVar
+    ident <- atomically $ stateTVar databaseAsyncs $ \(next, m) -> (next, (next + 1, m))
+    -- Wrap the gate wait too. A restart can cancel the child while it is still
+    -- parked, and it must still drop its entry, or 'cancelTrackedAsyncs' spins
+    -- on an entry that never deregisters.
+    a <- async $
+            (takeMVar gate >> act)
+              `finally` atomically (modifyTVar' databaseAsyncs (second (IntMap.delete ident)))
+    atomically $ modifyTVar' databaseAsyncs (second (IntMap.insert ident (void a)))
+    putMVar gate ()
+    pure a
+
+-- | Cancel every tracked async and wait for it to die. Loops until the registry
+-- is empty to catch threads a dying thread spawns.
+cancelTrackedAsyncs :: Database -> IO ()
+cancelTrackedAsyncs Database{..} =
+  let go = do
+        as <- atomically $ IntMap.elems . snd <$> readTVar databaseAsyncs
+        unless (null as) $ do
+          mapM_ (\a -> throwTo (asyncThreadId a) AsyncCancelled) as
+          mapM_ waitCatch as
+          go
+  in go
+
 -- | 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
+asyncWithCleanUp :: Database -> AIO a -> AIO (IO a)
+asyncWithCleanUp db 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
+        a <- trackedAsync db $ restore io
         atomicModifyIORef'_ st (void a :)
         return $ wait a
 
@@ -357,19 +412,19 @@ 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 :: Database -> Wait -> IO (Either (IO ()) (Async ()))
+waitOrSpawn _  (Wait io)  = pure $ Left io
+waitOrSpawn db (Spawn io) = Right <$> trackedAsync db io
 
-waitConcurrently_ :: [Wait] -> AIO ()
-waitConcurrently_ [] = pure ()
-waitConcurrently_ [one] = liftIO $ justWait one
-waitConcurrently_ many = do
+waitConcurrently_ :: Database -> [Wait] -> AIO ()
+waitConcurrently_ _ [] = pure ()
+waitConcurrently_ _ [one] = liftIO $ justWait one
+waitConcurrently_ db 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
+        waits <- liftIO $ traverse (waitOrSpawn db . fmapWait unmask) many
         let (syncs, asyncs) = partitionEithers waits
         liftIO $ atomicModifyIORef'_ ref (asyncs ++)
         return (asyncs, syncs)

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

@@ -5,6 +5,7 @@
 
 module Development.IDE.Graph.Internal.Types where
 
+import           Control.Concurrent.Async           (Async)
 import           Control.Concurrent.STM             (STM)
 import           Control.Monad                      ((>=>))
 import           Control.Monad.Catch
@@ -16,6 +17,7 @@ import qualified Data.ByteString                    as BS
 import           Data.Dynamic
 import           Data.Foldable                      (fold)
 import qualified Data.HashMap.Strict                as Map
+import           Data.IntMap.Strict                 (IntMap)
 import           Data.IORef
 import           Data.List                          (intercalate)
 import           Data.Maybe
@@ -112,6 +114,10 @@ data Database = Database {
     databaseExtra  :: Dynamic,
     databaseRules  :: TheRules,
     databaseStep   :: !(TVar Step),
+    -- | Every in-flight rule computation, keyed by a monotonic id, drained by
+    -- 'cancelTrackedAsyncs'. See Note [Cancelling escaped rule computations]
+    -- in Development.IDE.Graph.Internal.Database.
+    databaseAsyncs :: !(TVar (Int, IntMap (Async ()))),
     databaseValues :: !(Map Key KeyDetails)
     }
 
@@ -152,7 +158,7 @@ data Result = Result {
     resultValue     :: !Value,
     resultBuilt     :: !Step, -- ^ the step when it was last recomputed
     resultChanged   :: !Step, -- ^ the step when it last changed
-    resultVisited   :: !Step, -- ^ the step when it was last looked up
+    resultVisited   :: !Step, -- ^ the step when it was last looked up or produced
     resultDeps      :: !ResultDeps,
     resultExecution :: !Seconds, -- ^ How long it took, last time it ran
     resultData      :: !BS.ByteString

hls-graph/test/ActionSpec.hs

@@ -5,7 +5,10 @@ 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.Exception                       (catch, onException)
+import           Control.Monad                           (void)
 import           Control.Monad.IO.Class                  (MonadIO (..))
 import           Development.IDE.Graph                   (shakeOptions)
 import           Development.IDE.Graph.Database          (shakeNewDatabase,
@@ -17,6 +20,7 @@ import           Development.IDE.Graph.Internal.Types
 import           Development.IDE.Graph.Rule
 import           Example
 import qualified StmContainers.Map                       as STM
+import           System.Timeout                          (timeout)
 import           Test.Hspec
 
 
@@ -129,3 +133,58 @@ spec = do
     res `shouldBe` [[True]]
     Just (Clean res) <- lookup (newKey theKey) <$> getDatabaseValues theDb
     resultDeps res `shouldBe` UnknownDeps
+
+  describe "Async registry" $ do
+    it "registers a running computation and drops it on normal completion" $ do
+      started <- C.newEmptyMVar
+      proceed <- C.newEmptyMVar :: IO (MVar ())
+      done    <- C.newEmptyMVar
+      db@(ShakeDatabase _ _ theDb) <- shakeNewDatabase shakeOptions $
+        addRule $ \(Rule :: Rule ()) _old _mode -> do
+          liftIO $ C.putMVar started () >> C.takeMVar proceed
+          return $ RunResult ChangedRecomputeDiff "" () (return ())
+      forkSwallowCancel (shakeRunDatabase db [apply1 (Rule @())] >> C.putMVar done ())
+      C.takeMVar started
+      runningCount theDb `shouldReturn` 1
+      C.putMVar proceed ()
+      _ <- C.takeMVar done
+      runningCount theDb `shouldReturn` 0
+    it "cancels an in-flight computation on restart so no thread leaks" $ do
+      started   <- C.newEmptyMVar
+      proceed   <- C.newEmptyMVar :: IO (MVar ())  -- never filled, so the rule blocks here
+      cancelled <- C.newEmptyMVar
+      db@(ShakeDatabase _ _ theDb) <- shakeNewDatabase shakeOptions $
+        addRule $ \(Rule :: Rule ()) _old _mode -> do
+          liftIO $ (C.putMVar started () >> C.takeMVar proceed)
+                     `onException` C.putMVar cancelled ()
+          return $ RunResult ChangedRecomputeDiff "" () (return ())
+      forkSwallowCancel (shakeRunDatabase db [apply1 (Rule @())])
+      C.takeMVar started
+      runningCount theDb `shouldReturn` 1
+      incDatabase theDb (Just [])
+      -- The rule's thread received the async exception
+      got <- timeout 5_000_000 (C.takeMVar cancelled) -- 5s
+      got `shouldBe` Just ()
+      runningCount theDb `shouldReturn` 0
+    it "drains nested in-flight computations on restart" $ do
+      started <- C.newEmptyMVar
+      proceed <- C.newEmptyMVar :: IO (MVar ())
+      db@(ShakeDatabase _ _ theDb) <- shakeNewDatabase shakeOptions $ do
+        addRule $ \(Rule :: Rule ()) _old _mode -> do
+          liftIO $ C.putMVar started () >> C.takeMVar proceed
+          return $ RunResult ChangedRecomputeDiff "" () (return ())
+        addRule $ \(Rule :: Rule Bool) _old _mode -> do
+          () <- apply1 (Rule @())
+          return $ RunResult ChangedRecomputeDiff "" True (return ())
+      forkSwallowCancel (shakeRunDatabase db [apply1 (Rule @Bool)])
+      C.takeMVar started
+      -- Parent `Rule Bool` blocks on the leaf `Rule ()`.
+      n <- runningCount theDb
+      n `shouldSatisfy` (>= 2)
+      -- The drain loop reaps the whole chain, not just the roots.
+      incDatabase theDb (Just [])
+      runningCount theDb `shouldReturn` 0
+  where
+    forkSwallowCancel act = void . C.forkIO $ void act `catch` \AsyncCancelled -> pure ()
+    runningCount :: Database -> IO Int
+    runningCount theDb = length . snd <$> readTVarIO (databaseAsyncs theDb)