SessionLoader: batch load

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

@@ -146,10 +146,16 @@
   | LogNewComponentCache !(([FileDiagnostic], Maybe HscEnvEq), DependencyInfo)
   | LogHieBios HieBios.Log
   | LogSessionLoadingChanged
+  | LogSessionNewLoadedFiles ![FilePath]
+  | LogSessionReloadOnError FilePath ![FilePath]
 deriving instance Show Log
 
 instance Pretty Log where
   pretty = \case
+    LogSessionReloadOnError path files ->
+      "Reloading file due to error in" <+> pretty path <+> "with files:" <+> pretty files
+    LogSessionNewLoadedFiles files ->
+      "New loaded files:" <+> pretty files
     LogNoneCradleFound path ->
       "None cradle found for" <+> pretty path <+> ", ignoring the file"
     LogSettingInitialDynFlags ->
@@ -424,7 +430,8 @@
 loadSessionWithOptions :: Recorder (WithPriority Log) -> SessionLoadingOptions -> FilePath -> TQueue (IO ()) -> IO (Action IdeGhcSession)
 loadSessionWithOptions recorder SessionLoadingOptions{..} rootDir que = do
   let toAbsolutePath = toAbsolute rootDir -- see Note [Root Directory]
-  cradle_files <- newIORef []
+  cradle_files <- newIORef (Set.fromList [])
+  error_loading_files <- newIORef (Set.fromList [])
   -- Mapping from hie.yaml file to HscEnv, one per hie.yaml file
   hscEnvs <- newVar Map.empty :: IO (Var HieMap)
   -- Mapping from a Filepath to HscEnv
@@ -434,6 +441,8 @@
   -- they are inconsistent. So, everywhere you modify 'fileToFlags',
   -- you have to modify 'filesMap' as well.
   filesMap <- newVar HM.empty :: IO (Var FilesMap)
+  -- Pending files waiting to be loaded
+  pendingFilesTQueue <- newTQueueIO
   -- Version of the mappings above
   version <- newVar 0
   biosSessionLoadingVar <- newVar Nothing :: IO (Var (Maybe SessionLoadingPreferenceConfig))
@@ -550,7 +559,7 @@
 
 
     let session :: (Maybe FilePath, NormalizedFilePath, ComponentOptions, FilePath)
-                -> IO (IdeResult HscEnvEq,[FilePath])
+                -> IO ((IdeResult HscEnvEq,[FilePath]), HashSet FilePath)
         session args@(hieYaml, _cfp, _opts, _libDir) = do
           (new_deps, old_deps) <- packageSetup args
 
@@ -562,13 +571,15 @@
           let new_cache = newComponentCache recorder optExtensions _cfp hscEnv
           all_target_details <- new_cache old_deps new_deps
 
+          let flags_map' = HM.fromList (concatMap toFlagsMap all_targets')
+              all_targets' = concat all_target_details
+              newLoaded = HM.keys flags_map'
           this_dep_info <- getDependencyInfo $ maybeToList hieYaml
           let (all_targets, this_flags_map, this_options)
                 = case HM.lookup _cfp flags_map' of
                     Just this -> (all_targets', flags_map', this)
                     Nothing -> (this_target_details : all_targets', HM.insert _cfp this_flags flags_map', this_flags)
-                  where all_targets' = concat all_target_details
-                        flags_map' = HM.fromList (concatMap toFlagsMap all_targets')
+                  where
                         this_target_details = TargetDetails (TargetFile _cfp) this_error_env this_dep_info [_cfp]
                         this_flags = (this_error_env, this_dep_info)
                         this_error_env = ([this_error], Nothing)
@@ -580,27 +591,25 @@
 
           void $ modifyVar' fileToFlags $ Map.insert hieYaml this_flags_map
           void $ modifyVar' filesMap $ flip HM.union (HM.fromList (map ((,hieYaml) . fst) $ concatMap toFlagsMap all_targets))
+          -- Typecheck all files in the project on startup
+          checkProject <- getCheckProject
           -- The VFS doesn't change on cradle edits, re-use the old one.
           -- Invalidate all the existing GhcSession build nodes by restarting the Shake session
-          keys2 <- invalidateShakeCache
           restartShakeSession VFSUnmodified "new component" [] $ do
+            keys2 <- invalidateShakeCache
             keys1 <- extendKnownTargets all_targets
+            unless (null new_deps || not checkProject) $ do
+                  cfps' <- liftIO $ filterM (IO.doesFileExist . fromNormalizedFilePath) (concatMap targetLocations all_targets)
+                  void $ shakeEnqueue extras $ mkDelayedAction "InitialLoad" Debug $ void $ do
+                      mmt <- uses GetModificationTime cfps'
+                      let cs_exist = catMaybes (zipWith (<$) cfps' mmt)
+                      modIfaces <- uses GetModIface cs_exist
+                      -- update exports map
+                      shakeExtras <- getShakeExtras
+                      let !exportsMap' = createExportsMap $ mapMaybe (fmap hirModIface) modIfaces
+                      liftIO $ atomically $ modifyTVar' (exportsMap shakeExtras) (exportsMap' <>)
             return [keys1, keys2]
-
-          -- Typecheck all files in the project on startup
-          checkProject <- getCheckProject
-          unless (null new_deps || not checkProject) $ do
-                cfps' <- liftIO $ filterM (IO.doesFileExist . fromNormalizedFilePath) (concatMap targetLocations all_targets)
-                void $ shakeEnqueue extras $ mkDelayedAction "InitialLoad" Debug $ void $ do
-                    mmt <- uses GetModificationTime cfps'
-                    let cs_exist = catMaybes (zipWith (<$) cfps' mmt)
-                    modIfaces <- uses GetModIface cs_exist
-                    -- update exports map
-                    shakeExtras <- getShakeExtras
-                    let !exportsMap' = createExportsMap $ mapMaybe (fmap hirModIface) modIfaces
-                    liftIO $ atomically $ modifyTVar' (exportsMap shakeExtras) (exportsMap' <>)
-
-          return $ second Map.keys this_options
+          return $ (second Map.keys this_options, Set.fromList $ fromNormalizedFilePath <$> newLoaded)
 
     let consultCradle :: Maybe FilePath -> FilePath -> IO (IdeResult HscEnvEq, [FilePath])
         consultCradle hieYaml cfp = do
@@ -615,11 +624,20 @@
            -- Display a user friendly progress message here: They probably don't know what a cradle is
            let progMsg = "Setting up " <> T.pack (takeBaseName (cradleRootDir cradle))
                          <> " (for " <> T.pack lfpLog <> ")"
+
+           pendingFiles <- Set.insert cfp . Set.fromList <$> (atomically $ flushTQueue pendingFilesTQueue)
+           errorFiles <- readIORef error_loading_files
+           old_files <- readIORef cradle_files
+           -- if the file is in error loading files, we fall back to single loading mode
+           let extraToLoads = if cfp `Set.member` errorFiles
+                                then Set.empty
+                                -- remove error files from pending files since error loading need to load one by one
+                                else Set.delete cfp $ pendingFiles `Set.difference` errorFiles
+
            eopts <- mRunLspTCallback lspEnv (\act -> withIndefiniteProgress progMsg Nothing NotCancellable (const act)) $
               withTrace "Load cradle" $ \addTag -> do
                   addTag "file" lfpLog
-                  old_files <- readIORef cradle_files
-                  res <- cradleToOptsAndLibDir recorder (sessionLoading clientConfig) cradle cfp old_files
+                  res <- cradleToOptsAndLibDir recorder (sessionLoading clientConfig) cradle cfp (Set.toList $ Set.delete cfp $ extraToLoads <> old_files)
                   addTag "result" (show res)
                   return res
 
@@ -633,18 +651,43 @@
                  [] -> error $ "GHC version could not be parsed: " <> version
                  ((runTime, _):_)
                    | compileTime == runTime -> do
-                     atomicModifyIORef' cradle_files (\xs -> (cfp:xs,()))
-                     session (hieYaml, toNormalizedFilePath' cfp, opts, libDir)
+                     (results, allNewLoaded) <- session (hieYaml, toNormalizedFilePath' cfp, opts, libDir)
+                     -- put back to pending que if not listed in the results
+                     -- delete cfp even if we report No cradle target found for the cfp
+                     let remainPendingFiles = Set.delete cfp $ pendingFiles `Set.difference` allNewLoaded
+                     let newLoaded = pendingFiles `Set.intersection` allNewLoaded
+                     atomically $ forM_ remainPendingFiles (writeTQueue pendingFilesTQueue)
+                     -- log new loaded files
+                     logWith recorder Info $ LogSessionNewLoadedFiles $ Set.toList newLoaded
+                     -- remove all new loaded file from error loading files
+                     atomicModifyIORef' error_loading_files (\old -> (old `Set.difference` allNewLoaded, ()))
+                     atomicModifyIORef' cradle_files (\xs -> (newLoaded <> xs,()))
+                     return results
                    | otherwise -> return (([renderPackageSetupException cfp GhcVersionMismatch{..}], Nothing),[])
              -- Failure case, either a cradle error or the none cradle
              Left err -> do
-               dep_info <- getDependencyInfo (maybeToList hieYaml)
-               let ncfp = toNormalizedFilePath' cfp
-               let res = (map (\err' -> renderCradleError err' cradle ncfp) err, Nothing)
-               void $ modifyVar' fileToFlags $
-                    Map.insertWith HM.union hieYaml (HM.singleton ncfp (res, dep_info))
-               void $ modifyVar' filesMap $ HM.insert ncfp hieYaml
-               return (res, maybe [] pure hieYaml ++ concatMap cradleErrorDependencies err)
+                if (not $ null extraToLoads)
+                then do
+                    -- mark as less loaded files as failedLoadingFiles as possible
+                    -- limitation is that when we are loading files, and the dependencies of old_files
+                    -- are changed, and old_files are not valid anymore.
+                    -- but they will still be in the old_files, and will not move to error_loading_files.
+                    -- And make other files failed to load in batch mode.
+                    let failedLoadingFiles = (Set.insert cfp extraToLoads) `Set.difference` old_files
+                    atomicModifyIORef' error_loading_files (\xs -> (failedLoadingFiles <> xs,()))
+                    -- retry without other files
+                    atomically $ forM_ pendingFiles (writeTQueue pendingFilesTQueue)
+                    logWith recorder Info $ LogSessionReloadOnError cfp (Set.toList pendingFiles)
+                    consultCradle hieYaml cfp
+                else do
+                    dep_info <- getDependencyInfo (maybeToList hieYaml)
+                    let ncfp = toNormalizedFilePath' cfp
+                    let res = (map (\err' -> renderCradleError err' cradle ncfp) err, Nothing)
+                    void $ modifyVar' fileToFlags $
+                            Map.insertWith HM.union hieYaml (HM.singleton ncfp (res, dep_info))
+                    void $ modifyVar' filesMap $ HM.insert ncfp hieYaml
+                    atomicModifyIORef' error_loading_files (\xs -> (Set.insert cfp xs,()))
+                    return (res, maybe [] pure hieYaml ++ concatMap cradleErrorDependencies err)
 
     let
         -- | We allow users to specify a loading strategy.
@@ -677,39 +720,46 @@
             modifyVar_ filesMap (const (return HM.empty))
             -- Don't even keep the name cache, we start from scratch here!
             modifyVar_ hscEnvs (const (return Map.empty))
+            -- cleanup error loading files and cradle files
+            atomicModifyIORef' error_loading_files (\_ -> (Set.empty,()))
+            atomicModifyIORef' cradle_files (\_ -> (Set.empty,()))
 
           v <- Map.findWithDefault HM.empty hieYaml <$> readVar fileToFlags
-          let cfp = toAbsolutePath file
-          case HM.lookup (toNormalizedFilePath' cfp) v of
+          case HM.lookup (toNormalizedFilePath' file) v of
             Just (opts, old_di) -> do
               deps_ok <- checkDependencyInfo old_di
               if not deps_ok
                 then do
+                  -- if deps are old, we can try to load the error files again
+                  atomicModifyIORef' error_loading_files (\xs -> (Set.delete file xs,()))
+                  atomicModifyIORef' cradle_files (\xs -> (Set.delete file xs,()))
                   -- If the dependencies are out of date then clear both caches and start
                   -- again.
                   modifyVar_ fileToFlags (const (return Map.empty))
                   modifyVar_ filesMap (const (return HM.empty))
                   -- Keep the same name cache
                   modifyVar_ hscEnvs (return . Map.adjust (const []) hieYaml )
-                  consultCradle hieYaml cfp
+                  consultCradle hieYaml file
                 else return (opts, Map.keys old_di)
-            Nothing -> consultCradle hieYaml cfp
+            Nothing -> consultCradle hieYaml file
 
     -- The main function which gets options for a file. We only want one of these running
     -- at a time. Therefore the IORef contains the currently running cradle, if we try
     -- to get some more options then we wait for the currently running action to finish
     -- before attempting to do so.
     let getOptions :: FilePath -> IO (IdeResult HscEnvEq, [FilePath])
         getOptions file = do
-            let ncfp = toNormalizedFilePath' (toAbsolutePath file)
+            let ncfp = toNormalizedFilePath' file
             cachedHieYamlLocation <- HM.lookup ncfp <$> readVar filesMap
             hieYaml <- cradleLoc file
             sessionOpts (join cachedHieYamlLocation <|> hieYaml, file) `Safe.catch` \e ->
                 return (([renderPackageSetupException file e], Nothing), maybe [] pure hieYaml)
 
     returnWithVersion $ \file -> do
+      let absFile = toAbsolutePath file
+      atomically $ writeTQueue pendingFilesTQueue absFile
       -- see Note [Serializing runs in separate thread]
-      awaitRunInThread que $ getOptions file
+      awaitRunInThread que $ getOptions absFile
 
 -- | Run the specific cradle on a specific FilePath via hie-bios.
 -- This then builds dependencies or whatever based on the cradle, gets the