Move worker restart ref into WorkerThread

crtschin · crtschin · commit 224e74f5a332 · 2026-05-13T22:38:28.000+02:00
diff --git a/ghcide-test/exe/ShakeRestartTests.hs b/ghcide-test/exe/ShakeRestartTests.hs
@@ -4,8 +4,6 @@
 module ShakeRestartTests (tests) where
 
 import           Control.Concurrent.STM
-import           Data.IORef
-import           Data.IORef.Extra           (atomicModifyIORef'_)
 import           Development.IDE.Core.Shake
 import           Development.IDE.Graph      (newKey)
 import           Language.LSP.VFS
@@ -20,37 +18,18 @@ tests = testGroup "shake restart merging"
         newestVFSModified vfs1 VFSUnmodified @?= vfs1
         newestVFSModified VFSUnmodified vfs1 @?= vfs1
 
-    , testCase "mergePendingRestart Nothing" $ do
-        let p = PendingRestart VFSUnmodified [] ["reason"] [] []
-        if mergePendingRestart p Nothing == p
-          then pure ()
-          else assertFailure "merging with nothing should get new"
-
-    , testCase "mergePendingRestart Just" $ do
+    , testCase "<>" $ do
         done1 <- newEmptyTMVarIO
         done2 <- newEmptyTMVarIO
         let key1 = newKey ("1" :: String)
             key2 = newKey ("2" :: String)
             p1 = PendingRestart VFSUnmodified [pure [key1]] ["r1"] [] [done1]
             p2 = PendingRestart VFSUnmodified [pure [key2]] ["r2"] [] [done2]
-            merged = mergePendingRestart p1 (Just p2)
+            merged = p1 <> p2
 
         pendingRestartReasons merged @?= ["r1", "r2"]
         keys <- sequence $ reverse $ pendingRestartActionBetweenSessions merged
         concat keys @?= [key2, key1]
-
-    , testCase "RestartSlot coalescing" $ do
-        slot <- newRestartSlot
-        let p1 = PendingRestart VFSUnmodified [] ["r1"] [] []
-            p2 = PendingRestart VFSUnmodified [] ["r2"] [] []
-
-        atomicModifyIORef'_ (queuedRestart slot) $ Just . mergePendingRestart p1
-        atomicModifyIORef'_ (queuedRestart slot) $ Just . mergePendingRestart p2
-
-        res <- atomicModifyIORef' (queuedRestart slot) (Nothing,)
-        case res of
-            Nothing -> assertFailure "Should have a pending restart"
-            Just p  -> pendingRestartReasons p @?= ["r2", "r1"]
     ]
 
 instance Eq VFSModified where
diff --git a/ghcide/session-loader/Development/IDE/Session.hs b/ghcide/session-loader/Development/IDE/Session.hs
@@ -90,6 +90,7 @@ import           Data.HashMap.Strict                 (HashMap)
 import           Data.HashSet                        (HashSet)
 import qualified Data.HashSet                        as Set
 import           Database.SQLite.Simple
+import qualified Development.IDE.Core.Shake          as Shake
 import           Development.IDE.Core.Tracing        (withTrace)
 import           Development.IDE.Core.WorkerThread
 import           Development.IDE.Session.Dependency
@@ -136,6 +137,7 @@ data Log
   | LogLookupSessionCache !FilePath
   | LogTime !String
   | LogSessionGhc Ghc.Log
+  | LogShake Shake.Log
 deriving instance Show Log
 
 instance Pretty Log where
@@ -209,6 +211,7 @@ instance Pretty Log where
     LogSessionGhc msg -> pretty msg
     LogSessionLoadingChanged ->
       "Session Loading config changed, reloading the full session."
+    LogShake msg -> pretty msg
 
 -- | Bump this version number when making changes to the format of the data stored in hiedb
 hiedbDataVersion :: String
@@ -633,7 +636,7 @@ newSessionState = do
 -- components mapping to the same hie.yaml file are mapped to the same
 -- HscEnv which is updated as new components are discovered.
 
-loadSessionWithOptions :: Recorder (WithPriority Log) -> SessionLoadingOptions -> FilePath -> TaskQueue (IO ()) -> IO (Action IdeGhcSession)
+loadSessionWithOptions :: Recorder (WithPriority Log) -> SessionLoadingOptions -> FilePath -> WorkerTasks STM (IO ()) -> IO (Action IdeGhcSession)
 loadSessionWithOptions recorder SessionLoadingOptions{..} rootDir que = do
   let toAbsolutePath = toAbsolute rootDir -- see Note [Root Directory]
 
@@ -663,7 +666,7 @@ loadSessionWithOptions recorder SessionLoadingOptions{..} rootDir que = do
     -- see Note [Serializing runs in separate thread]
     -- Start the 'getOptionsLoop' if the queue is empty
     liftIO $ atomically $
-      Extra.whenM (isEmptyTaskQueue que) $ do
+      Extra.whenM (nullWorkerTasks que) $ do
         let newSessionLoadingOptions = SessionLoadingOptions
               { findCradle = cradleLoc
               , ..
@@ -683,7 +686,7 @@ loadSessionWithOptions recorder SessionLoadingOptions{..} rootDir que = do
               , sessionLoadingOptions = newSessionLoadingOptions
               }
 
-        writeTaskQueue que (runReaderT (getOptionsLoop recorder sessionShake sessionState knownTargetsVar) sessionEnv)
+        addWorkerTask que (runReaderT (getOptionsLoop recorder sessionShake sessionState knownTargetsVar) sessionEnv)
 
     -- Each one of deps will be registered as a FileSystemWatcher in the GhcSession action
     -- so that we can get a workspace/didChangeWatchedFiles notification when a dep changes.
diff --git a/ghcide/src/Development/IDE/Core/Compile.hs b/ghcide/src/Development/IDE/Core/Compile.hs
@@ -75,7 +75,7 @@ import           Development.IDE.Core.Preprocessor
 import           Development.IDE.Core.ProgressReporting       (progressUpdate)
 import           Development.IDE.Core.RuleTypes
 import           Development.IDE.Core.Shake
-import           Development.IDE.Core.WorkerThread            (writeTaskQueue)
+import           Development.IDE.Core.WorkerThread            (WorkerTasks (..))
 import           Development.IDE.Core.Tracing                 (withTrace)
 import qualified Development.IDE.GHC.Compat                   as Compat
 import qualified Development.IDE.GHC.Compat                   as GHC
@@ -942,7 +942,7 @@ indexHieFile se mod_summary srcPath !hash hf = do
       -- hiedb doesn't use the Haskell src, so we clear it to avoid unnecessarily keeping it around
       let !hf' = hf{hie_hs_src = mempty}
       modifyTVar' indexPending $ HashMap.insert srcPath hash
-      writeTaskQueue indexQueue $ \withHieDb -> do
+      addWorkerTask indexQueue $ \withHieDb -> do
         -- We are now in the worker thread
         -- Check if a newer index of this file has been scheduled, and if so skip this one
         newerScheduled <- atomically $ do
diff --git a/ghcide/src/Development/IDE/Core/FileStore.hs b/ghcide/src/Development/IDE/Core/FileStore.hs
@@ -304,7 +304,8 @@ typecheckParentsAction recorder nfp = do
 setSomethingModified :: VFSModified -> IdeState -> T.Text -> IO [Key] -> IO ()
 setSomethingModified vfs state reason actionBetweenSession = do
     -- Update database to remove any files that might have been renamed/deleted
-    atomically $ writeTaskQueue (indexQueue $ hiedbWriter $ shakeExtras state) (\withHieDb -> withHieDb deleteMissingRealFiles)
+    let indexQueue' = indexQueue $ hiedbWriter $ shakeExtras state
+    atomically $ addWorkerTask indexQueue' (\withHieDb -> withHieDb deleteMissingRealFiles)
     void $ restartShakeSession (shakeExtras state) vfs reason [] actionBetweenSession
 
 registerFileWatches :: [String] -> LSP.LspT Config IO Bool
diff --git a/ghcide/src/Development/IDE/Core/Shake.hs b/ghcide/src/Development/IDE/Core/Shake.hs
@@ -26,10 +26,10 @@ module Development.IDE.Core.Shake(
     IdeState, shakeSessionInit, shakeExtras, shakeDb, rootDir,
     ShakeExtras(..), getShakeExtras, getShakeExtrasRules,
     KnownTargets(..), Target(..), toKnownFiles, unionKnownTargets, mkKnownTargets,
-    IdeRule, IdeResult, RestartQueue,
+    IdeRule, IdeResult,
     GetModificationTime(GetModificationTime, GetModificationTime_, missingFileDiagnostics),
     shakeOpen, shakeShut,
-    withRestartWorker, newRestartSlot,
+    newRestartSlot,
     shakeEnqueue,
     newSession,
     use, useNoFile, uses, useWithStaleFast, useWithStaleFast', delayedAction,
@@ -76,14 +76,14 @@ module Development.IDE.Core.Shake(
     RestartSlot(..),
     addPersistentRule,
     newestVFSModified,
-    mergePendingRestart,
     garbageCollectDirtyKeys,
     garbageCollectDirtyKeysOlderThan,
     Log(..),
     VFSModified(..), getClientConfigAction,
     ThreadQueue(..),
     runWithSignal,
-    askShake
+    askShake,
+    processPendingRestart
     ) where
 
 import           Control.Concurrent.Async
@@ -149,7 +149,6 @@ import           Development.IDE.GHC.Compat             (NameCache,
                                                          initNameCache,
                                                          knownKeyNames)
 #endif
-import           Data.IORef.Extra                       (atomicModifyIORef'_)
 import qualified Data.Text.Encoding                     as T
 import           Development.IDE.GHC.Orphans            ()
 import           Development.IDE.Graph                  hiding (ShakeValue,
@@ -192,10 +191,7 @@ import qualified StmContainers.Map                      as STM
 import           System.FilePath                        hiding (makeRelative)
 import           System.IO.Unsafe                       (unsafePerformIO)
 import           System.Time.Extra
-import           UnliftIO                               (IORef,
-                                                         MonadUnliftIO (withRunInIO),
-                                                         atomicModifyIORef',
-                                                         newIORef)
+import           UnliftIO                               (MonadUnliftIO (..))
 
 
 data Log
@@ -277,16 +273,15 @@ data HieDbWriter
 -- | Actions to queue up on the index worker thread
 -- The inner `(HieDb -> IO ()) -> IO ()` wraps `HieDb -> IO ()`
 -- with (currently) retry functionality
-type IndexQueue = TaskQueue (((HieDb -> IO ()) -> IO ()) -> IO ())
-type RestartQueue = TaskQueue (IO ())
-type LoaderQueue = TaskQueue (IO ())
-
-
-data ThreadQueue = ThreadQueue {
-    tIndexQueue    :: IndexQueue
-    , tRestartSlot :: RestartSlot
-    , tLoaderQueue :: LoaderQueue
-}
+type IndexQueue = WorkerTasks STM (((HieDb -> IO ()) -> IO ()) -> IO ())
+type RestartRef = WorkerTasks STM PendingRestart
+type LoaderQueue = WorkerTasks STM (IO ())
+
+data ThreadQueue = ThreadQueue
+  { tIndexQueue  :: IndexQueue
+  , tRestartSlot :: RestartSlot
+  , tLoaderQueue :: LoaderQueue
+  }
 
 -- Note [Semantic Tokens Cache Location]
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -826,23 +821,21 @@ data PendingRestart = PendingRestart
     , pendingRestartDoneSignals           :: ![TMVar ()]
     }
 
-newestVFSModified :: VFSModified -> VFSModified -> VFSModified
-newestVFSModified VFSUnmodified old     = old
-newestVFSModified new@(VFSModified _) _ = new
-
-mergePendingRestart :: PendingRestart -> Maybe PendingRestart -> PendingRestart
-mergePendingRestart new Nothing = new
-mergePendingRestart new (Just old) = PendingRestart
+instance Semigroup PendingRestart where
+  new <> old = PendingRestart
     { pendingRestartVFS = newestVFSModified (pendingRestartVFS new) (pendingRestartVFS old)
     , pendingRestartReasons = pendingRestartReasons new ++ pendingRestartReasons old
     , pendingRestartActions = pendingRestartActions new ++ pendingRestartActions old
     , pendingRestartActionBetweenSessions = pendingRestartActionBetweenSessions new ++ pendingRestartActionBetweenSessions old
     , pendingRestartDoneSignals = pendingRestartDoneSignals new ++ pendingRestartDoneSignals old
     }
 
+newestVFSModified :: VFSModified -> VFSModified -> VFSModified
+newestVFSModified VFSUnmodified old     = old
+newestVFSModified new@(VFSModified _) _ = new
+
 data RestartSlot = RestartSlot
-    { queuedRestart      :: IORef (Maybe PendingRestart)
-    , restartSignal      :: MVar ()
+    { restartRef         :: WorkerTasks STM PendingRestart
     , lastRestartBarrier :: TVar (TMVar ())
     -- ^ A barrier that is filled when the most recent shake restart completes.
     --
@@ -851,77 +844,62 @@ data RestartSlot = RestartSlot
     -- restart can then wait on this.
     }
 
-newRestartSlot :: IO RestartSlot
-newRestartSlot = do
+newRestartSlot :: RestartRef -> IO RestartSlot
+newRestartSlot queuedRestart = do
     initialBarrier <- newTMVarIO ()  -- starts filled (no pending restart)
-    RestartSlot <$> newIORef Nothing <*> newEmptyMVar <*> newTVarIO initialBarrier
+    RestartSlot <$> pure queuedRestart <*> newTVarIO initialBarrier
 
 -- | Restart the current 'ShakeSession' with the given system actions.
 --
 -- Any actions running in the current session will be aborted, but actions added
 -- via 'shakeEnqueue' will be requeued.
 shakeRestart :: IdeState -> VFSModified -> T.Text -> [DelayedAction ()] -> IO [Key] -> IO ()
 shakeRestart IdeState{..} vfs reason acts ioActionBetweenShakeSession = do
-    restartDone <- newEmptyTMVarIO
-    let slot = restartSlot shakeExtras
-    -- Publish this restart's barrier, that dependents LSP requests can wait on.
-    atomically $ writeTVar (lastRestartBarrier slot) restartDone
-    atomicModifyIORef'_ (queuedRestart slot) $ Just . mergePendingRestart PendingRestart
-        { pendingRestartVFS = vfs
-        , pendingRestartActionBetweenSessions = [ioActionBetweenShakeSession]
-        , pendingRestartReasons = [reason]
-        , pendingRestartActions = acts
-        , pendingRestartDoneSignals = [restartDone]
-        }
-    void $ tryPutMVar (restartSignal slot) ()
-    -- Block until the restart (including ioActionBetweenShakeSession) completes.
-    -- This preserves the invariant from the original synchronous shakeRestart:
-    -- callers (e.g. the session loader) must not proceed until their
-    -- between-session actions have run, otherwise downstream rules can observe
-    -- stale results (see Note at Session.hs restartSession call site).
-    atomically $ readTMVar restartDone
-
--- | Run a worker that asynchronously processes shake restart requests. Will
--- only ever queue upto 1 additional restart, accumulating data while processing
--- any restart.
-withRestartWorker :: MVar IdeState -> IO r -> IO r
-withRestartWorker ideMVar action = do
-    let restartWorkerAction = do
-          ide@IdeState{..} <- readMVar ideMVar
-          forever (processPendingRestart (shakeRecorder shakeExtras) ide)
-            `catch` \(e :: SomeException) ->
-                case fromException e of
-                    Just AsyncCancelled -> throwIO e
-                    _ -> logWith (shakeRecorder shakeExtras) Error (LogRestartWorkerException e)
-
-    withAsync restartWorkerAction $ \_ -> action
-
-processPendingRestart :: Recorder (WithPriority Log) -> IdeState -> IO ()
-processPendingRestart recorder IdeState{..} = do
-    takeMVar (restartSignal (restartSlot shakeExtras))
-    pendingRestart <- atomicModifyIORef' (queuedRestart (restartSlot shakeExtras)) (Nothing,)
-    void $ forM pendingRestart $ \PendingRestart {..} -> do
-        flip finally (atomically $ traverse (flip tryPutTMVar ()) (reverse pendingRestartDoneSignals)) $ do
-            let sessionAction runner = do
-                    (stopTime,()) <- duration $ logErrorAfter 10 $ cancelShakeSession runner
-                    keys <- fmap concat (sequence (reverse pendingRestartActionBetweenSessions))
-                    -- it is every important to update the dirty keys after we enter the critical section
-                    -- see Note [Housekeeping rule cache and dirty key outside of hls-graph]
-                    atomically $ modifyTVar' (dirtyKeys shakeExtras) $ \x -> foldl' (flip insertKeySet) x keys
-                    res <- shakeDatabaseProfile shakeDb
-                    backlog <- readTVarIO $ dirtyKeys shakeExtras
-                    queue <- atomicallyNamed "actionQueue - peek" $ peekInProgress $ actionQueue shakeExtras
-
-                    -- this log is required by tests
-                    logWith recorder Debug $ LogBuildSessionRestart (reverse pendingRestartReasons) queue backlog stopTime res
-
-            withMVar' shakeSession sessionAction $ \() ->
-                -- It is crucial to be masked here, otherwise we can get killed
-                -- between spawning the new thread and updating shakeSession.
-                -- See https://github.com/haskell/ghcide/issues/79
-                (,()) <$> newSession recorder shakeExtras pendingRestartVFS shakeDb
-                    (reverse pendingRestartActions)
-                    (reverse pendingRestartReasons)
+  restartDone <- newEmptyTMVarIO
+  let RestartSlot {..} = restartSlot shakeExtras
+  -- Publish this restart's barrier, that dependents LSP requests can wait on.
+  atomically $ do
+    writeTVar lastRestartBarrier restartDone
+    addWorkerTask restartRef $ PendingRestart
+      { pendingRestartVFS = vfs
+      , pendingRestartActionBetweenSessions = [ioActionBetweenShakeSession]
+      , pendingRestartReasons = [reason]
+      , pendingRestartActions = acts
+      , pendingRestartDoneSignals = [restartDone]
+      }
+
+processPendingRestart :: Recorder (WithPriority Log) -> MVar IdeState -> PendingRestart -> IO ()
+processPendingRestart recorder ideMVar pendingRestart = do
+  processPendingRestart' recorder ideMVar pendingRestart
+    `catch` \(e :: SomeException) ->
+        case fromException e of
+            Just AsyncCancelled -> throwIO e
+            _ -> logWith recorder Error (LogRestartWorkerException e)
+
+processPendingRestart' :: Recorder (WithPriority Log) -> MVar IdeState -> PendingRestart -> IO ()
+processPendingRestart' recorder ideMVar PendingRestart{..} = do
+    IdeState{..} <- readMVar ideMVar
+    flip finally (atomically $ traverse (flip tryPutTMVar ()) (reverse pendingRestartDoneSignals)) $ do
+        let sessionAction runner = do
+                (stopTime,()) <- duration $ logErrorAfter 10 $ cancelShakeSession runner
+                keys <- fmap concat (sequence (reverse pendingRestartActionBetweenSessions))
+                -- it is every important to update the dirty keys after we enter the critical section
+                -- see Note [Housekeeping rule cache and dirty key outside of hls-graph]
+                atomically $ modifyTVar' (dirtyKeys shakeExtras) $ \x -> foldl' (flip insertKeySet) x keys
+                res <- shakeDatabaseProfile shakeDb
+                backlog <- readTVarIO $ dirtyKeys shakeExtras
+                queue <- atomicallyNamed "actionQueue - peek" $ peekInProgress $ actionQueue shakeExtras
+
+                -- this log is required by tests
+                logWith recorder Debug $ LogBuildSessionRestart (reverse pendingRestartReasons) queue backlog stopTime res
+
+        withMVar' shakeSession sessionAction $ \() ->
+            -- It is crucial to be masked here, otherwise we can get killed
+            -- between spawning the new thread and updating shakeSession.
+            -- See https://github.com/haskell/ghcide/issues/79
+            (,()) <$> newSession recorder shakeExtras pendingRestartVFS shakeDb
+                (reverse pendingRestartActions)
+                (reverse pendingRestartReasons)
   where
     logErrorAfter :: Seconds -> IO () -> IO ()
     logErrorAfter seconds action = flip withAsync (const action) $ do
diff --git a/ghcide/src/Development/IDE/Core/WorkerThread.hs b/ghcide/src/Development/IDE/Core/WorkerThread.hs
diff --git a/ghcide/src/Development/IDE/LSP/LanguageServer.hs b/ghcide/src/Development/IDE/LSP/LanguageServer.hs
