Clarify hls-graph runtime restart keys

soulomoon · soulomoon · commit 027d58883c24 · 2026-05-08T05:58:03.000+08:00
diff --git a/ghcide/src/Development/IDE/Core/Shake.hs b/ghcide/src/Development/IDE/Core/Shake.hs
@@ -145,7 +145,8 @@ import           Development.IDE.GHC.Orphans            ()
 import           Development.IDE.Graph                  hiding (ShakeValue,
                                                          action)
 import qualified Development.IDE.Graph                  as Shake
-import           Development.IDE.Graph.Database         (ShakeDatabase,
+import           Development.IDE.Graph.Database         (RuntimeRestartKeys (..),
+                                                         ShakeDatabase,
                                                          instantiateDelayedAction,
                                                          shakeComputeToPreserve,
                                                          shakeGetActionQueueLength,
@@ -873,17 +874,16 @@ shakeRestart recorder IdeState{..} vfs reason acts ioActionBetweenShakeSession =
         prepareRuntimeRestart :: Bool -> [Key] -> IO (RuntimeKeysChanged, RuntimeRestartStats)
         prepareRuntimeRestart optRunSubset keys
           | optRunSubset = do
-              (affected, changedKeys, _lookupCount, _) <-
-                  shakeComputeToPreserve shakeDb $ fromListKeySet keys
-              logErrorAfter 10 $ shakeShutDatabase affected shakeDb
+              runtimeRestartKeys <- shakeComputeToPreserve shakeDb $ fromListKeySet keys
+              logErrorAfter 10 $ shakeShutDatabase (restartKillKeys runtimeRestartKeys) shakeDb
               surviving <- shakePeekAsyncsDelivers shakeDb
               queueCount <- shakeGetActionQueueLength shakeDb
               let preserved = fromListKeySet $ map GraphRuntime.deliverKey surviving
               pure
-                ( Just (changedKeys, preserved)
+                ( Just (runtimeRestartKeys, preserved)
                 , RuntimeRestartStats
                     { runtimeDirtyCount = length keys
-                    , runtimeAffectedCount = lengthKeySet affected
+                    , runtimeAffectedCount = lengthKeySet (restartKillKeys runtimeRestartKeys)
                     , runtimePreservedCount = lengthKeySet preserved
                     , runtimeActionQueueCount = queueCount
                     , runtimeSurvivingActions = map GraphRuntime.deliverName surviving
@@ -930,7 +930,7 @@ shakeEnqueue ShakeExtras{actionQueue, shakeRecorder} act = do
 
 data VFSModified = VFSUnmodified | VFSModified !VFS
 
-type RuntimeKeysChanged = Maybe (([Key], [Key]), KeySet)
+type RuntimeKeysChanged = Maybe (RuntimeRestartKeys, KeySet)
 
 -- | Set up a new 'ShakeSession' with a set of initial actions
 --   Will crash if there is an existing 'ShakeSession' running.
@@ -961,8 +961,8 @@ newSession recorder ShakeExtras{..} vfsMod shakeDb acts reason runtimeKeysChange
         workRun restore = withSpan "Shake session" $ \otSpan -> do
           setTag otSpan "reason" (fromString reason)
           setTag otSpan "queue" (fromString $ unlines $ map actionName reenqueued)
-          whenJust runtimeKeysChanged $ \((_, newKeys), _) ->
-              setTag otSpan "keys" (BS8.pack $ unlines $ map show newKeys)
+          whenJust runtimeKeysChanged $ \(runtimeRestartKeys, _) ->
+              setTag otSpan "keys" (BS8.pack $ unlines $ map show $ restartDirtyKeys runtimeRestartKeys)
           res <- try @SomeException $
             restore startDatabase
           return $ do
diff --git a/hls-graph/src/Development/IDE/Graph/Database.hs b/hls-graph/src/Development/IDE/Graph/Database.hs
@@ -16,6 +16,7 @@ module Development.IDE.Graph.Database(
     ,shakeGetBuildEdges,
     shakeShutDatabase,
     shakeGetActionQueueLength,
+    RuntimeRestartKeys(..),
     shakeComputeToPreserve,
     -- shakedatabaseRuntimeDep,
     shakePeekAsyncsDelivers,
@@ -91,7 +92,7 @@ unvoid = fmap undefined
 -- seperate incrementing the step from running the build.
 -- Also immediately enqueues upsweep actions for the newly dirty keys.
 shakeRunDatabaseForKeysSep
-    :: Maybe (([Key],[Key]),KeySet) -- ^ Set of keys changed since last run. 'Nothing' means everything has changed
+    :: Maybe (RuntimeRestartKeys, KeySet) -- ^ Set of keys changed since last run. 'Nothing' means everything has changed
     -> ShakeDatabase
     -> [Action a]
     -> IO (IO [Either SomeException a])
@@ -126,7 +127,7 @@ mkDelayedAction s p a = do
     u <- newUnique
     return $ DelayedAction (newDirectKey $ hashUnique u) s (toEnum (fromEnum p)) a
 
-shakeComputeToPreserve :: ShakeDatabase -> KeySet -> IO (KeySet, ([Key], [Key]), Int, [Key])
+shakeComputeToPreserve :: ShakeDatabase -> KeySet -> IO RuntimeRestartKeys
 shakeComputeToPreserve (ShakeDatabase _ _ db) ks = atomically (computeToPreserve db ks)
 
 shakeRunDatabaseForKeys
@@ -146,7 +147,12 @@ shakeRunDatabaseForKeysWithExceptions
     -> IO [Either SomeException a]
 shakeRunDatabaseForKeysWithExceptions Nothing sdb as2 = join $ shakeRunDatabaseForKeysSep Nothing sdb as2
 shakeRunDatabaseForKeysWithExceptions (Just x) sdb as2 =
-    let y = fromListKeySet x in join $ shakeRunDatabaseForKeysSep (Just (([], toListKeySet y), y)) sdb as2
+    let y = fromListKeySet x
+        restartKeys = RuntimeRestartKeys
+            { restartKillKeys = y
+            , restartDirtyKeys = toListKeySet y
+            }
+    in join $ shakeRunDatabaseForKeysSep (Just (restartKeys, y)) sdb as2
 
 
 shakePeekAsyncsDelivers :: ShakeDatabase -> IO [DeliverStatus]
diff --git a/hls-graph/src/Development/IDE/Graph/Internal/Database.hs b/hls-graph/src/Development/IDE/Graph/Internal/Database.hs
@@ -9,7 +9,7 @@
 {-# LANGUAGE RecordWildCards    #-}
 {-# LANGUAGE TypeFamilies       #-}
 
-module Development.IDE.Graph.Internal.Database (compute, newDatabase, incDatabase, build, getDirtySet, getKeysAndVisitAge, AsyncParentKill(..), computeToPreserve, getRunTimeRDeps, spawnAsyncWithDbRegistration) where
+module Development.IDE.Graph.Internal.Database (compute, newDatabase, incDatabase, build, getDirtySet, getKeysAndVisitAge, AsyncParentKill(..), RuntimeRestartKeys(..), computeToPreserve, getRunTimeRDeps, spawnAsyncWithDbRegistration) where
 
 import           Prelude                              hiding (unzip)
 
@@ -38,9 +38,8 @@ import qualified StmContainers.Map                    as SMap
 import           System.Time.Extra                    (duration)
 import           UnliftIO                             (MVar, atomically,
                                                        newEmptyMVar, putMVar,
-                                                       takeMVar)
+                                                       readMVar)
 
-import qualified Data.List                            as List
 import qualified UnliftIO.Exception                   as UE
 
 #if MIN_VERSION_base(4,19,0)
@@ -65,12 +64,11 @@ newDatabase dataBaseLogger databaseActionQueue databaseExtra databaseRules = do
 -- | Increment the step and mark dirty.
 --   Assumes that the database is not running a build
 -- only some keys are dirty
-incDatabase :: Database -> Maybe (([Key], [Key]), KeySet) -> IO KeySet
-incDatabase db (Just ((_oldKeys, newKeys), preserves)) = do
+incDatabase :: Database -> Maybe (RuntimeRestartKeys, KeySet) -> IO KeySet
+incDatabase db (Just (RuntimeRestartKeys{..}, preserves)) = do
     atomicallyNamed "incDatabase" $ modifyTVar' (databaseStep db) $ \(Step i) -> Step $ i + 1
-    forM_ newKeys $ \newKey -> atomically $ SMap.focus updateDirty newKey (databaseValues db)
-    -- only upsweep the keys that are not preserved
-    -- atomically $ writeUpsweepQueue (filter  (`notMemberKeySet` preserves) oldkeys ++ newKeys) db
+    forM_ restartDirtyKeys $ \newKey -> atomically $ SMap.focus updateDirty newKey (databaseValues db)
+    -- Only re-enqueue actions that were not preserved across the restart.
     return $ preserves
 
 -- all keys are dirty
@@ -82,10 +80,29 @@ incDatabase db Nothing = do
         SMap.focus updateDirty k (databaseValues db)
     return $ mempty
 
-computeToPreserve :: Database -> KeySet -> STM (KeySet, ([Key], [Key]), Int, [Key])
-computeToPreserve db dirtySet = do
-  (oldKeys, newKeys, affected) <- transitiveDirtyListBottomUpDiff db (toListKeySet dirtySet) []
-  pure (affected, (oldKeys, newKeys), length newKeys, [])
+data RuntimeRestartKeys = RuntimeRestartKeys
+  { restartKillKeys  :: !KeySet
+    -- ^ Keys used to select running runtime actions to stop before the next
+    -- session starts. This may include rule keys and delayed-action 'DirectKey's.
+  , restartDirtyKeys :: ![Key]
+    -- ^ Rule database keys to mark dirty before the next run. In the ghcide
+    -- restart path this is rule-key-only by construction; the raw hls-graph API
+    -- does not enforce that invariant by type.
+  } deriving Show
+
+-- Note [RuntimeRestartKeys]
+-- The restart plan intentionally keeps runtime cancellation separate from rule
+-- dirtiness. 'restartKillKeys' is consumed by shutdown and may include direct
+-- delayed-action keys. 'restartDirtyKeys' is consumed by the rule database and
+-- is expected to contain only rule keys that can be marked dirty.
+-- For the ghcide restart path, the initial dirty seeds come from rule keys
+-- ('toKey'/'toNoFileKey'), so 'restartDirtyKeys' can use the
+-- 'databaseRRuntimeDep' closure directly. Direct/root runtime edges are stored
+-- separately in 'databaseRRuntimeDepRoot' by 'insertdatabaseRuntimeDep' and are
+-- expanded only for 'restartKillKeys'. The raw hls-graph API does not enforce
+-- this seed invariant by type.
+computeToPreserve :: Database -> KeySet -> STM RuntimeRestartKeys
+computeToPreserve = transitiveDirtyKeysBottomUp
 
 updateDirty :: Monad m => Focus.Focus KeyDetails m ()
 updateDirty = Focus.adjust $ \(KeyDetails status rdeps) ->
@@ -130,7 +147,7 @@ interpreBuildContinue :: Database -> Key -> (Key, BuildContinue) -> IO (Key, Res
 interpreBuildContinue _db _pk (_kid, BCStop k v) = return (k, v)
 interpreBuildContinue db _pk (kid, BCContinue Nothing) = builderOneFinal db emptyStack kid
 interpreBuildContinue _db _pk (_kid, BCContinue (Just barrier)) =
-    takeMVar barrier >>= either throwIO pure
+    readMVar barrier >>= either throwIO pure
 
 
 builderOne :: Key -> Database -> Stack -> Key -> IO (Key, BuildContinue)
@@ -324,32 +341,22 @@ getRunTimeRDeps db k = SMap.lookup k (databaseRRuntimeDep db)
 getDeps :: SMap.Map Key KeySet -> Key -> STM (Maybe KeySet)
 getDeps m k = SMap.lookup k m
 
--- Edges in the reverse-dependency graph go from a child to its parents.
--- We perform a DFS and, after exploring all outgoing edges, cons the node onto
--- the accumulator. This yields children-before-parents order directly.
-
--- the lefts are keys that are no longer affected, we can try to mark them clean
--- the rights are new affected keys, we need to mark them dirty
-transitiveDirtyListBottomUpDiff :: Database -> [Key] -> [Key] -> STM ([Key], [Key], KeySet)
-transitiveDirtyListBottomUpDiff database seeds allOldKeys = do
-  (newKeys, seen) <- cacheTransitiveDirtyListBottomUpDFSWithRootKey database $ fromListKeySet seeds
-  let oldKeys = filter (`notMemberKeySet` seen) allOldKeys
-  return (oldKeys, newKeys, seen)
-
-cacheTransitiveDirtyListBottomUpDFSWithRootKey :: Database -> KeySet -> STM ([Key], KeySet)
-cacheTransitiveDirtyListBottomUpDFSWithRootKey db@Database{..} seeds = do
-  (newKeys, seen) <- cacheTransitiveDirtyListBottomUpDFS db seeds
-  -- we should put pump root keys back to seen
-  -- for each new key, get its root keys and put them back to seen
-  -- newKeys is for upsweep, databaseRRuntimeDepRoot only add new root keys which is not needed for upsweep
-  -- but seen is for thread filtering, we need to make sure all root keys are in seen
-  (_newKeys, newSeen) <- transitiveDirtyListBottomUpDFS databaseRRuntimeDepRoot seen
+transitiveDirtyKeysBottomUp :: Database -> KeySet -> STM RuntimeRestartKeys
+transitiveDirtyKeysBottomUp db@Database{..} seeds = do
+  TransitiveDirtyKeys dirtyKeys seen <- cacheTransitiveDirtyListBottomUpDFS db seeds
+  -- restartDirtyKeys should contain only rule keys. restartKillKeys also needs
+  -- the root/direct delayed-action keys, so expand through the root dependency
+  -- map only for the kill set.
+  TransitiveDirtyKeys _newKeys newSeen <- transitiveDirtyListBottomUpDFS databaseRRuntimeDepRoot seen
   let rootKey = newKey "root"
-  return $ (List.delete rootKey newKeys, deleteKeySet rootKey newSeen)
+  pure RuntimeRestartKeys
+    { restartDirtyKeys = dirtyKeys
+    , restartKillKeys = deleteKeySet rootKey newSeen
+    }
 
 
 
-cacheTransitiveDirtyListBottomUpDFS :: Database -> KeySet -> STM ([Key], KeySet)
+cacheTransitiveDirtyListBottomUpDFS :: Database -> KeySet -> STM TransitiveDirtyKeys
 cacheTransitiveDirtyListBottomUpDFS Database{..} seeds = do
     SMap.lookup seeds databaseTransitiveRRuntimeDepCache >>= \case
         Just v  -> return v
@@ -358,22 +365,25 @@ cacheTransitiveDirtyListBottomUpDFS Database{..} seeds = do
             SMap.insert r seeds databaseTransitiveRRuntimeDepCache
             return r
 
-transitiveDirtyListBottomUpDFS :: SMap.Map Key KeySet -> KeySet -> STM ([Key], KeySet)
+-- Edges in the reverse-dependency graph go from a child to its parents.
+-- We perform a DFS and, after exploring all outgoing edges, cons the node onto
+-- the accumulator. This yields children-before-parents order directly.
+transitiveDirtyListBottomUpDFS :: SMap.Map Key KeySet -> KeySet -> STM TransitiveDirtyKeys
 transitiveDirtyListBottomUpDFS database seeds = do
-  let go1 :: Key -> ([Key], KeySet) -> STM ([Key], KeySet)
-      go1 x acc@(dirties, seen) = do
+  let go1 :: Key -> TransitiveDirtyKeys -> STM TransitiveDirtyKeys
+      go1 x acc@TransitiveDirtyKeys{transitiveDirtySet = seen} = do
         if x `memberKeySet` seen
           then pure acc
           else do
-            let newAcc = (dirties, insertKeySet x seen)
+            let newAcc = acc{transitiveDirtySet = insertKeySet x seen}
             mnext <- getDeps database x
-            (newDirties, newSeen) <- foldrM go1 newAcc (maybe mempty toListKeySet mnext)
-            return (x:newDirties, newSeen)
-                -- if it is root key, we do not add it to the dirty list
-                -- since root key is not up for upsweep
-                -- but it would be in the seen list, so we would kill dirty root key async
+            childClosure <- foldrM go1 newAcc (maybe mempty toListKeySet mnext)
+            return childClosure{transitiveDirtyList = x : transitiveDirtyList childClosure}
+                -- Root keys are filtered out by 'transitiveDirtyKeysBottomUp'
+                -- for the dirty list, but kept in the set long enough to find
+                -- runtime roots that need shutdown.
   -- traverse all seeds
-  foldrM go1 ([], mempty) (toListKeySet seeds)
+  foldrM go1 (TransitiveDirtyKeys [] mempty) (toListKeySet seeds)
 
 -- | Original spawnRefresh using the general pattern
 -- inline
diff --git a/hls-graph/src/Development/IDE/Graph/Internal/Types.hs b/hls-graph/src/Development/IDE/Graph/Internal/Types.hs
@@ -283,7 +283,7 @@ data Database = Database {
     -- if not in any of the transitive reverse deps of a dirty node, it is clean
     -- we can skip clean the threads.
     -- this is update right before we query the database for the key result.
-    databaseTransitiveRRuntimeDepCache :: SMap.Map KeySet ([Key], KeySet),
+    databaseTransitiveRRuntimeDepCache :: SMap.Map KeySet TransitiveDirtyKeys,
     -- ^ this is a cache for transitive reverse deps if we have computed it before
     -- and the databaseRRuntimeDep did not change since last time
     -- it is very useful for large projects where many files depend on a few common files
@@ -307,6 +307,13 @@ data Database = Database {
 
     }
 
+data TransitiveDirtyKeys = TransitiveDirtyKeys
+    { transitiveDirtyList :: ![Key]
+      -- ^ Dirty keys in children-before-parents order.
+    , transitiveDirtySet  :: !KeySet
+      -- ^ Same transitive closure as a set, used for membership/filtering.
+    } deriving Show
+
 
 ---------------------------------------------------------------------
 -- | Remove finished asyncs from 'databaseThreads' (non-blocking).
diff --git a/hls-graph/test/ActionSpec.hs b/hls-graph/test/ActionSpec.hs
@@ -6,20 +6,30 @@ module ActionSpec where
 import           Control.Concurrent                      (MVar, readMVar)
 import qualified Control.Concurrent                      as C
 import           Control.Concurrent.STM
+import           Control.Exception                       (SomeException)
+import           Control.Monad                           (void)
 import           Control.Monad.IO.Class                  (MonadIO (..))
 import           Data.Typeable                           (Typeable)
 import           Development.IDE.Graph                   (RuleResult,
                                                           shakeOptions)
 import           Development.IDE.Graph.Classes           (Hashable)
-import           Development.IDE.Graph.Database          (shakeNewDatabase,
+import           Development.IDE.Graph.Database          (RuntimeRestartKeys (..),
+                                                          mkDelayedAction,
+                                                          shakeComputeToPreserve,
+                                                          shakeNewDatabase,
                                                           shakeRunDatabase,
-                                                          shakeRunDatabaseForKeys)
+                                                          shakeRunDatabaseForKeys,
+                                                          shakeShutDatabase)
+import           Development.IDE.Graph.Internal.Action   (actionCatch,
+                                                          actionFinally,
+                                                          pumpActionThreadReRun)
 import           Development.IDE.Graph.Internal.Database (build, incDatabase)
 import           Development.IDE.Graph.Internal.Key
 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
 
 
@@ -32,6 +42,13 @@ itInThread = it
 
 shakeRunDatabaseFromRight :: ShakeDatabase -> [Action a] -> IO [a]
 shakeRunDatabaseFromRight = shakeRunDatabase
+
+waitForRuntimeRootDep :: Database -> Key -> Key -> IO ()
+waitForRuntimeRootDep Database{..} child parent =
+  atomically $ do
+    deps <- STM.lookup child databaseRRuntimeDepRoot
+    check $ maybe False (memberKeySet parent) deps
+
 spec :: Spec
 spec = do
   describe "apply1" $ itInThread "Test build update, Buggy dirty mechanism in hls-graph #4237" $ do
@@ -105,6 +122,41 @@ spec = do
       db <- shakeNewDatabase shakeOptions $ addRule $ \(Rule :: Rule ()) _old _mode -> error "boom"
       let res = shakeRunDatabaseFromRight db $ pure $ apply1 (Rule @())
       res `shouldThrow` anyErrorCall
+    itInThread "restart kills a delayed action parked behind a caught producer failure" $ do
+      producerStarted <- C.newEmptyMVar
+      releaseProducer <- C.newEmptyMVar
+      producerCaught <- C.newEmptyMVar
+      waiterFinalized <- C.newEmptyMVar
+      sdb@(ShakeDatabase _ _ theDb) <- shakeNewDatabase shakeOptions $
+        addRule $ \(Rule :: Rule Int) _old _mode -> do
+          liftIO $ void $ C.tryPutMVar producerStarted ()
+          liftIO $ readMVar releaseProducer
+          error "boom"
+      producer <- mkDelayedAction "producer" Debug $
+        actionCatch @SomeException
+          (void $ apply1 (Rule @Int))
+          (\_ -> liftIO $ void $ C.tryPutMVar producerCaught ())
+      waiter <- mkDelayedAction "waiter" Debug $
+        actionFinally
+          (do
+            liftIO $ readMVar producerStarted
+            void $ apply1 (Rule @Int))
+          (void $ C.tryPutMVar waiterFinalized ())
+
+      _ <- shakeRunDatabaseForKeys Nothing sdb
+        [ pumpActionThreadReRun sdb producer
+        , pumpActionThreadReRun sdb waiter
+        ]
+      let dirtyKey = newKey (Rule @Int)
+      waitForRuntimeRootDep theDb dirtyKey (uniqueID waiter)
+      C.putMVar releaseProducer ()
+      readMVar producerCaught
+      C.tryReadMVar waiterFinalized >>= (`shouldBe` Nothing)
+
+      runtimeRestartKeys <- shakeComputeToPreserve sdb (singletonKeySet dirtyKey)
+      uniqueID waiter `memberKeySet` restartKillKeys runtimeRestartKeys `shouldBe` True
+      shakeShutDatabase (restartKillKeys runtimeRestartKeys) sdb
+      timeout 1000000 (readMVar waiterFinalized) >>= (`shouldBe` Just ())
     itInThread "computes a rule with branching dependencies does not invoke phantom dependencies #3423" $ do
       cond <- C.newMVar True
       count <- C.newMVar 0
