Avoid race conditions with VFS and VFS versions (haskell#2789)

* Avoid race conditions with VFS and VFS version

We need to take VFS snapshots as soon as we get a change notification.

Consider the following interleaving of events:

1. Change Notification A (updates LSP VFS)
2. Restart Shake Session (A changed) initiated
3. Change Notification B (updates LSP VFS)
4. Restart Shake Session (A changed) takes VFS snapshot and possibly performs more computation
5. Restart Shake Session (B changed)

In particular, between step 3 and 5, we took a snapshot for a previous build,
but this snapshot included changes from a newer VFS state that the build should
not have seen.

To fix this, we need to take snapshots as soon as a notification handler is
called, before forking any threads. This works because LSP calls all handlers
in a single threaded fashion and these handlers block message processing. It
is essential to this on the LSP handler thread rather than the reactor thread
that GHCIDE sets up in order to maintin the property.

* Disable flaky test 'add missing module (non workspace)'

exe/Main.hs

@@ -46,7 +46,7 @@ main = do
     -- This plugin just installs a handler for the `initialized` notification, which then
     -- picks up the LSP environment and feeds it to our recorders
     let lspRecorderPlugin = (defaultPluginDescriptor "LSPRecorderCallback")
-          { pluginNotificationHandlers = mkPluginNotificationHandler LSP.SInitialized $ \_ _ _ -> do
+          { pluginNotificationHandlers = mkPluginNotificationHandler LSP.SInitialized $ \_ _ _ _ -> do
               env <- LSP.getLspEnv
               liftIO $ (cb1 <> cb2) env
           }

ghcide/exe/Main.hs

@@ -94,7 +94,7 @@ main = withTelemetryLogger $ \telemetryLogger -> do
     -- This plugin just installs a handler for the `initialized` notification, which then
     -- picks up the LSP environment and feeds it to our recorders
     let lspRecorderPlugin = (defaultPluginDescriptor "LSPRecorderCallback")
-          { pluginNotificationHandlers = mkPluginNotificationHandler LSP.SInitialized $ \_ _ _ -> do
+          { pluginNotificationHandlers = mkPluginNotificationHandler LSP.SInitialized $ \_ _ _ _ -> do
               env <- LSP.getLspEnv
               liftIO $ (cb1 <> cb2) env
           }

ghcide/session-loader/Development/IDE/Session.hs

@@ -581,7 +581,9 @@ loadSessionWithOptions recorder SessionLoadingOptions{..} dir = do
 
           -- Invalidate all the existing GhcSession build nodes by restarting the Shake session
           invalidateShakeCache
-          restartShakeSession "new component" []
+
+          -- The VFS doesn't change on cradle edits, re-use the old one.
+          restartShakeSession VFSUnmodified "new component" []
 
           -- Typecheck all files in the project on startup
           checkProject <- getCheckProject

ghcide/src/Development/IDE.hs

@@ -40,7 +40,8 @@ import           Development.IDE.Core.Shake            as X (FastResult (..),
                                                              useWithStaleFast,
                                                              useWithStaleFast',
                                                              useWithStale_,
-                                                             use_, uses, uses_)
+                                                             use_, uses, uses_,
+                                                             VFSModified(..))
 import           Development.IDE.GHC.Compat            as X (GhcVersion (..),
                                                              ghcVersion)
 import           Development.IDE.GHC.Error             as X

ghcide/src/Development/IDE/Core/FileStore.hs

@@ -223,19 +223,20 @@ fileStoreRules recorder isWatched = do
 -- | Note that some buffer for a specific file has been modified but not
 -- with what changes.
 setFileModified :: Recorder (WithPriority Log)
+                -> VFSModified
                 -> IdeState
                 -> Bool -- ^ Was the file saved?
                 -> NormalizedFilePath
                 -> IO ()
-setFileModified recorder state saved nfp = do
+setFileModified recorder vfs state saved nfp = do
     ideOptions <- getIdeOptionsIO $ shakeExtras state
     doCheckParents <- optCheckParents ideOptions
     let checkParents = case doCheckParents of
           AlwaysCheck -> True
           CheckOnSave -> saved
           _           -> False
     join $ atomically $ recordDirtyKeys (shakeExtras state) GetModificationTime [nfp]
-    restartShakeSession (shakeExtras state) (fromNormalizedFilePath nfp ++ " (modified)") []
+    restartShakeSession (shakeExtras state) vfs (fromNormalizedFilePath nfp ++ " (modified)") []
     when checkParents $
       typecheckParents recorder state nfp
 
@@ -256,14 +257,14 @@ typecheckParentsAction recorder nfp = do
 -- | Note that some keys have been modified and restart the session
 --   Only valid if the virtual file system was initialised by LSP, as that
 --   independently tracks which files are modified.
-setSomethingModified :: IdeState -> [Key] -> String -> IO ()
-setSomethingModified state keys reason = do
+setSomethingModified :: VFSModified -> IdeState -> [Key] -> String -> IO ()
+setSomethingModified vfs state keys reason = do
     -- Update database to remove any files that might have been renamed/deleted
     atomically $ do
         writeTQueue (indexQueue $ hiedbWriter $ shakeExtras state) (\withHieDb -> withHieDb deleteMissingRealFiles)
         modifyTVar' (dirtyKeys $ shakeExtras state) $ \x ->
             foldl' (flip HSet.insert) x keys
-    void $ restartShakeSession (shakeExtras state) reason []
+    void $ restartShakeSession (shakeExtras state) vfs reason []
 
 registerFileWatches :: [String] -> LSP.LspT Config IO Bool
 registerFileWatches globs = do

ghcide/src/Development/IDE/Core/Rules.hs

@@ -558,7 +558,7 @@ getHieAstsRule recorder =
 persistentHieFileRule :: Recorder (WithPriority Log) -> Rules ()
 persistentHieFileRule recorder = addPersistentRule GetHieAst $ \file -> runMaybeT $ do
   res <- readHieFileForSrcFromDisk recorder file
-  vfsRef <- asks vfs
+  vfsRef <- asks vfsVar
   vfsData <- liftIO $ vfsMap <$> readTVarIO vfsRef
   (currentSource, ver) <- liftIO $ case M.lookup (filePathToUri' file) vfsData of
     Nothing -> (,Nothing) . T.decodeUtf8 <$> BS.readFile (fromNormalizedFilePath file)

ghcide/src/Development/IDE/Core/Shake.hs

@@ -75,7 +75,8 @@ module Development.IDE.Core.Shake(
     addPersistentRule,
     garbageCollectDirtyKeys,
     garbageCollectDirtyKeysOlderThan,
-    Log(..)
+    Log(..),
+    VFSModified(..)
     ) where
 
 import           Control.Concurrent.Async
@@ -253,7 +254,8 @@ data ShakeExtras = ShakeExtras
     ,ideTesting :: IdeTesting
     -- ^ Whether to enable additional lsp messages used by the test suite for checking invariants
     ,restartShakeSession
-        :: String
+        :: VFSModified
+        -> String
         -> [DelayedAction ()]
         -> IO ()
     ,ideNc :: IORef NameCache
@@ -269,7 +271,7 @@ data ShakeExtras = ShakeExtras
     , persistentKeys :: TVar (HMap.HashMap Key GetStalePersistent)
       -- ^ Registery for functions that compute/get "stale" results for the rule
       -- (possibly from disk)
-    , vfs :: TVar VFS
+    , vfsVar :: TVar VFS
     -- ^ A snapshot of the current state of the virtual file system. Updated on shakeRestart
     -- VFS state is managed by LSP. However, the state according to the lsp library may be newer than the state of the current session,
     -- leaving us vulnerable to suble race conditions. To avoid this, we take a snapshot of the state of the VFS on every
@@ -318,7 +320,7 @@ class Typeable a => IsIdeGlobal a where
 -- | Read a virtual file from the current snapshot
 getVirtualFile :: NormalizedFilePath -> Action (Maybe VirtualFile)
 getVirtualFile nf = do
-  vfs <- fmap vfsMap . liftIO . readTVarIO . vfs =<< getShakeExtras
+  vfs <- fmap vfsMap . liftIO . readTVarIO . vfsVar =<< getShakeExtras
   pure $! Map.lookup (filePathToUri' nf) vfs -- Don't leak a reference to the entire map
 
 -- Take a snapshot of the current LSP VFS
@@ -598,7 +600,7 @@ shakeOpen recorder lspEnv defaultConfig logger debouncer
 
         dirtyKeys <- newTVarIO mempty
         -- Take one VFS snapshot at the start
-        vfs <- newTVarIO =<< vfsSnapshot lspEnv
+        vfsVar <- newTVarIO =<< vfsSnapshot lspEnv
         pure ShakeExtras{..}
     shakeDb  <-
         shakeNewDatabase
@@ -640,7 +642,10 @@ startTelemetry db extras@ShakeExtras{..}
 -- | Must be called in the 'Initialized' handler and only once
 shakeSessionInit :: Recorder (WithPriority Log) -> IdeState -> IO ()
 shakeSessionInit recorder ide@IdeState{..} = do
-    initSession <- newSession recorder shakeExtras shakeDb [] "shakeSessionInit"
+    -- Take a snapshot of the VFS - it should be empty as we've recieved no notifications
+    -- till now, but it can't hurt to be in sync with the `lsp` library.
+    vfs <- vfsSnapshot (lspEnv shakeExtras)
+    initSession <- newSession recorder shakeExtras (VFSModified vfs) shakeDb [] "shakeSessionInit"
     putMVar shakeSession initSession
     logDebug (ideLogger ide) "Shake session initialized"
 
@@ -679,8 +684,8 @@ delayedAction a = do
 -- | 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 :: Recorder (WithPriority Log) -> IdeState -> String -> [DelayedAction ()] -> IO ()
-shakeRestart recorder IdeState{..} reason acts =
+shakeRestart :: Recorder (WithPriority Log) -> IdeState -> VFSModified -> String -> [DelayedAction ()] -> IO ()
+shakeRestart recorder IdeState{..} vfs reason acts =
     withMVar'
         shakeSession
         (\runner -> do
@@ -707,7 +712,7 @@ shakeRestart recorder IdeState{..} reason acts =
         -- between spawning the new thread and updating shakeSession.
         -- See https://github.com/haskell/ghcide/issues/79
         (\() -> do
-          (,()) <$> newSession recorder shakeExtras shakeDb acts reason)
+          (,()) <$> newSession recorder shakeExtras vfs shakeDb acts reason)
     where
         logErrorAfter :: Seconds -> Recorder (WithPriority Log) -> IO () -> IO ()
         logErrorAfter seconds recorder action = flip withAsync (const action) $ do
@@ -743,19 +748,24 @@ shakeEnqueue ShakeExtras{actionQueue, logger} act = do
               ]
     return (wait' b >>= either throwIO return)
 
+data VFSModified = VFSUnmodified | VFSModified !VFS
+
 -- | Set up a new 'ShakeSession' with a set of initial actions
 --   Will crash if there is an existing 'ShakeSession' running.
 newSession
     :: Recorder (WithPriority Log)
     -> ShakeExtras
+    -> VFSModified
     -> ShakeDatabase
     -> [DelayedActionInternal]
     -> String
     -> IO ShakeSession
-newSession recorder extras@ShakeExtras{..} shakeDb acts reason = do
+newSession recorder extras@ShakeExtras{..} vfsMod shakeDb acts reason = do
 
     -- Take a new VFS snapshot
-    atomically . writeTVar vfs =<< vfsSnapshot lspEnv
+    case vfsMod of
+      VFSUnmodified -> pure ()
+      VFSModified vfs -> atomically $ writeTVar vfsVar vfs
 
     IdeOptions{optRunSubset} <- getIdeOptionsIO extras
     reenqueued <- atomicallyNamed "actionQueue - peek" $ peekInProgress actionQueue

ghcide/src/Development/IDE/LSP/Notifications.hs

@@ -55,41 +55,41 @@ whenUriFile uri act = whenJust (LSP.uriToFilePath uri) $ act . toNormalizedFileP
 descriptor :: Recorder (WithPriority Log) -> PluginId -> PluginDescriptor IdeState
 descriptor recorder plId = (defaultPluginDescriptor plId) { pluginNotificationHandlers = mconcat
   [ mkPluginNotificationHandler LSP.STextDocumentDidOpen $
-      \ide _ (DidOpenTextDocumentParams TextDocumentItem{_uri,_version}) -> liftIO $ do
+      \ide vfs _ (DidOpenTextDocumentParams TextDocumentItem{_uri,_version}) -> liftIO $ do
       atomically $ updatePositionMapping ide (VersionedTextDocumentIdentifier _uri (Just _version)) (List [])
       whenUriFile _uri $ \file -> do
           -- We don't know if the file actually exists, or if the contents match those on disk
           -- For example, vscode restores previously unsaved contents on open
           addFileOfInterest ide file Modified{firstOpen=True}
-          setFileModified (cmapWithPrio LogFileStore recorder) ide False file
+          setFileModified (cmapWithPrio LogFileStore recorder) (VFSModified vfs) ide False file
           logDebug (ideLogger ide) $ "Opened text document: " <> getUri _uri
 
   , mkPluginNotificationHandler LSP.STextDocumentDidChange $
-      \ide _ (DidChangeTextDocumentParams identifier@VersionedTextDocumentIdentifier{_uri} changes) -> liftIO $ do
+      \ide vfs _ (DidChangeTextDocumentParams identifier@VersionedTextDocumentIdentifier{_uri} changes) -> liftIO $ do
         atomically $ updatePositionMapping ide identifier changes
         whenUriFile _uri $ \file -> do
           addFileOfInterest ide file Modified{firstOpen=False}
-          setFileModified (cmapWithPrio LogFileStore recorder) ide False file
+          setFileModified (cmapWithPrio LogFileStore recorder) (VFSModified vfs) ide False file
         logDebug (ideLogger ide) $ "Modified text document: " <> getUri _uri
 
   , mkPluginNotificationHandler LSP.STextDocumentDidSave $
-      \ide _ (DidSaveTextDocumentParams TextDocumentIdentifier{_uri} _) -> liftIO $ do
+      \ide vfs _ (DidSaveTextDocumentParams TextDocumentIdentifier{_uri} _) -> liftIO $ do
         whenUriFile _uri $ \file -> do
             addFileOfInterest ide file OnDisk
-            setFileModified (cmapWithPrio LogFileStore recorder) ide True file
+            setFileModified (cmapWithPrio LogFileStore recorder) (VFSModified vfs) ide True file
         logDebug (ideLogger ide) $ "Saved text document: " <> getUri _uri
 
   , mkPluginNotificationHandler LSP.STextDocumentDidClose $
-        \ide _ (DidCloseTextDocumentParams TextDocumentIdentifier{_uri}) -> liftIO $ do
+        \ide vfs _ (DidCloseTextDocumentParams TextDocumentIdentifier{_uri}) -> liftIO $ do
           whenUriFile _uri $ \file -> do
               deleteFileOfInterest ide file
               let msg = "Closed text document: " <> getUri _uri
               scheduleGarbageCollection ide
-              setSomethingModified ide [] $ Text.unpack msg
+              setSomethingModified (VFSModified vfs) ide [] $ Text.unpack msg
               logDebug (ideLogger ide) msg
 
   , mkPluginNotificationHandler LSP.SWorkspaceDidChangeWatchedFiles $
-      \ide _ (DidChangeWatchedFilesParams (List fileEvents)) -> liftIO $ do
+      \ide vfs _ (DidChangeWatchedFilesParams (List fileEvents)) -> liftIO $ do
         -- See Note [File existence cache and LSP file watchers] which explains why we get these notifications and
         -- what we do with them
         -- filter out files of interest, since we already know all about those
@@ -106,24 +106,24 @@ descriptor recorder plId = (defaultPluginDescriptor plId) { pluginNotificationHa
             logDebug (ideLogger ide) $ "Watched file events: " <> Text.pack msg
             modifyFileExists ide fileEvents'
             resetFileStore ide fileEvents'
-            setSomethingModified ide [] msg
+            setSomethingModified (VFSModified vfs) ide [] msg
 
   , mkPluginNotificationHandler LSP.SWorkspaceDidChangeWorkspaceFolders $
-      \ide _ (DidChangeWorkspaceFoldersParams events) -> liftIO $ do
+      \ide _ _ (DidChangeWorkspaceFoldersParams events) -> liftIO $ do
         let add       = S.union
             substract = flip S.difference
         modifyWorkspaceFolders ide
           $ add       (foldMap (S.singleton . parseWorkspaceFolder) (_added   events))
           . substract (foldMap (S.singleton . parseWorkspaceFolder) (_removed events))
 
   , mkPluginNotificationHandler LSP.SWorkspaceDidChangeConfiguration $
-      \ide _ (DidChangeConfigurationParams cfg) -> liftIO $ do
+      \ide vfs _ (DidChangeConfigurationParams cfg) -> liftIO $ do
         let msg = Text.pack $ show cfg
         logDebug (ideLogger ide) $ "Configuration changed: " <> msg
         modifyClientSettings ide (const $ Just cfg)
-        setSomethingModified ide [toKey GetClientSettings emptyFilePath] "config change"
+        setSomethingModified (VFSModified vfs) ide [toKey GetClientSettings emptyFilePath] "config change"
 
-  , mkPluginNotificationHandler LSP.SInitialized $ \ide _ _ -> do
+  , mkPluginNotificationHandler LSP.SInitialized $ \ide _ _ _ -> do
       --------- Initialize Shake session --------------------------------------------------------------------
       liftIO $ shakeSessionInit (cmapWithPrio LogShake recorder) ide
 

ghcide/src/Development/IDE/LSP/Server.hs

@@ -22,6 +22,7 @@ import           Ide.Types                    (HasTracing, traceWithSpan)
 import           Language.LSP.Server          (Handlers, LspM)
 import qualified Language.LSP.Server          as LSP
 import           Language.LSP.Types
+import           Language.LSP.VFS
 import           UnliftIO.Chan
 
 data ReactorMessage
@@ -48,14 +49,16 @@ requestHandler m k = LSP.requestHandler m $ \RequestMessage{_method,_id,_params}
 notificationHandler
   :: forall (m :: Method FromClient Notification) c. (HasTracing (MessageParams m)) =>
      SMethod m
-  -> (IdeState -> MessageParams m -> LspM c ())
+  -> (IdeState -> VFS -> MessageParams m -> LspM c ())
   -> Handlers (ServerM c)
 notificationHandler m k = LSP.notificationHandler m $ \NotificationMessage{_params,_method}-> do
   (chan,ide) <- ask
   env <- LSP.getLspEnv
+  -- Take a snapshot of the VFS state on every notification
+  -- We only need to do this here because the VFS state is only updated
+  -- on notifications
+  vfs <- LSP.getVirtualFiles
   let trace x = otTracedHandler "Notification" (show _method) $ \sp -> do
         traceWithSpan sp _params
         x
-  writeChan chan $ ReactorNotification (trace $ LSP.runLspT env $ k ide _params)
-
-
+  writeChan chan $ ReactorNotification (trace $ LSP.runLspT env $ k ide vfs _params)

ghcide/src/Development/IDE/Plugin/HLS.hs

@@ -33,6 +33,7 @@ import           Ide.Plugin.Config
 import           Ide.PluginUtils              (getClientConfig)
 import           Ide.Types                    as HLS
 import qualified Language.LSP.Server          as LSP
+import           Language.LSP.VFS
 import           Language.LSP.Types
 import qualified Language.LSP.Types           as J
 import           Text.Regex.TDFA.Text         ()
@@ -190,7 +191,7 @@ extensibleNotificationPlugins recorder xs = mempty { P.pluginHandlers = handlers
       hs
     handlers = mconcat $ do
       (IdeNotification m :=> IdeNotificationHandler fs') <- DMap.assocs handlers'
-      pure $ notificationHandler m $ \ide params -> do
+      pure $ notificationHandler m $ \ide vfs params -> do
         config <- Ide.PluginUtils.getClientConfig
         let fs = filter (\(pid,_) -> plcGlobalOn $ configForPlugin config pid) fs'
         case nonEmpty fs of
@@ -200,7 +201,7 @@ extensibleNotificationPlugins recorder xs = mempty { P.pluginHandlers = handlers
           Just fs -> do
             -- We run the notifications in order, so the core ghcide provider
             -- (which restarts the shake process) hopefully comes last
-              mapM_ (\(pid,f) -> otTracedProvider pid (fromString $ show m) $ f ide params) fs
+              mapM_ (\(pid,f) -> otTracedProvider pid (fromString $ show m) $ f ide vfs params) fs
 
 -- ---------------------------------------------------------------------
 
@@ -226,7 +227,7 @@ newtype IdeHandler (m :: J.Method FromClient Request)
 
 -- | Combine the 'PluginHandler' for all plugins
 newtype IdeNotificationHandler (m :: J.Method FromClient Notification)
-  = IdeNotificationHandler [(PluginId, IdeState -> MessageParams m -> LSP.LspM Config ())]
+  = IdeNotificationHandler [(PluginId, IdeState -> VFS -> MessageParams m -> LSP.LspM Config ())]
 -- type NotificationHandler (m :: Method FromClient Notification) = MessageParams m -> IO ()`
 
 -- | Combine the 'PluginHandlers' for all plugins

ghcide/test/exe/Main.hs

@@ -421,7 +421,7 @@ diagnosticTests = testGroup "diagnostics"
       let contentA = T.unlines [ "module ModuleA where" ]
       _ <- createDoc "ModuleA.hs" "haskell" contentA
       expectDiagnostics [("ModuleB.hs", [])]
-  , ignoreInWindowsBecause "Broken in windows" $ testSessionWait "add missing module (non workspace)" $ do
+  , ignoreTestBecause "Flaky #2831" $ testSessionWait "add missing module (non workspace)" $ do
       -- need to canonicalize in Mac Os
       tmpDir <- liftIO $ canonicalizePath =<< getTemporaryDirectory
       let contentB = T.unlines
@@ -6417,7 +6417,7 @@ unitTests recorder logger = do
         let plugins = pluginDescToIdePlugins $
                 [ (defaultPluginDescriptor $ fromString $ show i)
                     { pluginNotificationHandlers = mconcat
-                        [ mkPluginNotificationHandler LSP.STextDocumentDidOpen $ \_ _ _ ->
+                        [ mkPluginNotificationHandler LSP.STextDocumentDidOpen $ \_ _ _ _ ->
                             liftIO $ atomicModifyIORef_ orderRef (i:)
                         ]
                     }