Merge pull request #2311 from ahoppen/long-build-server-init

Do not block SourceKit-LSP functionality when a build server takes long to initialize

Author: ahoppen

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -343,7 +343,26 @@ package actor BuildServerManager: QueueBasedMessageHandler {
   }
 
   /// Provider of file to main file mappings.
-  private var mainFilesProvider: MainFilesProvider?
+  ///
+  /// Force-unwrapped optional because initializing it requires access to `self`.
+  private var mainFilesProvider: Task<MainFilesProvider?, Never>! {
+    didSet {
+      // Must only be set once
+      precondition(oldValue == nil)
+      precondition(mainFilesProvider != nil)
+    }
+  }
+
+  package func mainFilesProvider<T: MainFilesProvider>(as: T.Type) async -> T? {
+    guard let mainFilesProvider = mainFilesProvider else {
+      return nil
+    }
+    guard let index = await mainFilesProvider.value as? T else {
+      logger.fault("Expected the main files provider of the build server manager to be an `\(T.self)`")
+      return nil
+    }
+    return index
+  }
 
   /// Build server delegate that will receive notifications about setting changes, etc.
   private weak var delegate: BuildServerManagerDelegate?
@@ -465,7 +484,11 @@ package actor BuildServerManager: QueueBasedMessageHandler {
     toolchainRegistry: ToolchainRegistry,
     options: SourceKitLSPOptions,
     connectionToClient: BuildServerManagerConnectionToClient,
-    buildServerHooks: BuildServerHooks
+    buildServerHooks: BuildServerHooks,
+    createMainFilesProvider:
+      @escaping @Sendable (
+        SourceKitInitializeBuildResponseData?, _ mainFilesChangedCallback: @escaping @Sendable () async -> Void
+      ) async -> MainFilesProvider?
   ) async {
     self.toolchainRegistry = toolchainRegistry
     self.options = options
@@ -578,6 +601,11 @@ package actor BuildServerManager: QueueBasedMessageHandler {
       await buildServerAdapter.send(OnBuildInitializedNotification())
       return initializeResponse
     }
+    self.mainFilesProvider = Task {
+      await createMainFilesProvider(initializationData) { [weak self] in
+        await self?.mainFilesChanged()
+      }
+    }
   }
 
   /// Explicitly shut down the build server.
@@ -630,12 +658,6 @@ package actor BuildServerManager: QueueBasedMessageHandler {
     self.delegate = delegate
   }
 
-  /// - Note: Needed because we need the `indexStorePath` and `indexDatabasePath` from the build server to create an
-  ///   IndexStoreDB, which serves as the `MainFilesProvider`. And thus this can't be set during initialization.
-  package func setMainFilesProvider(_ mainFilesProvider: MainFilesProvider?) {
-    self.mainFilesProvider = mainFilesProvider
-  }
-
   // MARK: Handling messages from the build server
 
   package func handle(notification: some NotificationType) async {
@@ -1274,13 +1296,12 @@ package actor BuildServerManager: QueueBasedMessageHandler {
   }
 
   private func buildTargets() async throws -> [BuildTargetIdentifier: BuildTargetInfo] {
-    guard let buildServerAdapter = try await buildServerAdapterAfterInitialized else {
-      return [:]
-    }
-
     let request = WorkspaceBuildTargetsRequest()
     let result = try await cachedBuildTargets.get(request, isolation: self) { request in
       let result = try await withTimeout(self.options.buildServerWorkspaceRequestsTimeoutOrDefault) {
+        guard let buildServerAdapter = try await self.buildServerAdapterAfterInitialized else {
+          return [:]
+        }
         let buildTargets = try await buildServerAdapter.send(request).targets
         let (depths, dependents) = await self.targetDepthsAndDependents(for: buildTargets)
         var result: [BuildTargetIdentifier: BuildTargetInfo] = [:]
@@ -1325,7 +1346,7 @@ package actor BuildServerManager: QueueBasedMessageHandler {
   }
 
   package func sourceFiles(in targets: Set<BuildTargetIdentifier>) async throws -> [SourcesItem] {
-    guard let buildServerAdapter = try await buildServerAdapterAfterInitialized, !targets.isEmpty else {
+    guard !targets.isEmpty else {
       return []
     }
 
@@ -1346,6 +1367,9 @@ package actor BuildServerManager: QueueBasedMessageHandler {
 
     let response = try await cachedTargetSources.get(request, isolation: self) { request in
       try await withTimeout(self.options.buildServerWorkspaceRequestsTimeoutOrDefault) {
+        guard let buildServerAdapter = try await self.buildServerAdapterAfterInitialized else {
+          return BuildTargetSourcesResponse(items: [])
+        }
         return try await buildServerAdapter.send(request)
       } resultReceivedAfterTimeout: { newResult in
         await self.buildTargetsDidChange(.sourceFilesReceivedResultAfterTimeout(request: request, newResult: newResult))
@@ -1393,7 +1417,6 @@ package actor BuildServerManager: QueueBasedMessageHandler {
   /// - Important: This method returns both buildable and non-buildable source files. Callers need to check
   /// `SourceFileInfo.isBuildable` if they are only interested in buildable source files.
   private func sourceFilesAndDirectories() async throws -> SourceFilesAndDirectories {
-    let supportsOutputPaths = await initializationData?.outputPathsProvider ?? false
 
     return try await cachedSourceFilesAndDirectories.get(
       SourceFilesAndDirectoriesKey(),
@@ -1411,7 +1434,7 @@ package actor BuildServerManager: QueueBasedMessageHandler {
         for sourceItem in sourcesItem.sources {
           let sourceKitData = sourceItem.sourceKitData
           let outputPath: OutputPath? =
-            if !supportsOutputPaths {
+            if !(await self.initializationData?.outputPathsProvider ?? false) {
               .notSupported
             } else if let outputPath = sourceKitData?.outputPath {
               .path(outputPath)
@@ -1499,7 +1522,7 @@ package actor BuildServerManager: QueueBasedMessageHandler {
   /// path is `/tmp`). If the realpath that indexstore-db returns could not be found in the build server's source files
   /// but the standardized path is part of the source files, return the standardized path instead.
   package func mainFiles(containing uri: DocumentURI) async -> [DocumentURI] {
-    guard let mainFilesProvider else {
+    guard let mainFilesProvider = await mainFilesProvider.value else {
       return [uri]
     }
     let mainFiles = Array(await mainFilesProvider.mainFiles(containing: uri, crossLanguage: false))

Sources/DocumentationLanguageService/DoccDocumentationHandler.swift

@@ -101,7 +101,7 @@ extension DocumentationLanguageService {
     }
     guard let moduleName, symbolName == moduleName else {
       // This is a symbol extension page. Find the symbol so that we can include it in the request.
-      guard let index = workspace.index(checkedFor: .deletedFiles) else {
+      guard let index = await workspace.index(checkedFor: .deletedFiles) else {
         throw ResponseError.requestFailed(doccDocumentationError: .indexNotAvailable)
       }
       return try await sourceKitLSPServer.withOnDiskDocumentManager { onDiskDocumentManager in
@@ -206,7 +206,7 @@ extension DocumentationLanguageService {
       return nil
     }
     let catalogIndex = try await documentationManager.catalogIndex(for: catalogURL)
-    guard let index = workspace.index(checkedFor: .deletedFiles) else {
+    guard let index = await workspace.index(checkedFor: .deletedFiles) else {
       return nil
     }
     let symbolInformation = try await index.doccSymbolInformation(

Sources/SourceKitLSP/IndexProgressManager.swift

@@ -83,7 +83,7 @@ actor IndexProgressManager {
       return
     }
     var status = IndexProgressStatus.upToDate
-    for indexManager in await sourceKitLSPServer.workspaces.compactMap({ $0.semanticIndexManager }) {
+    for indexManager in await sourceKitLSPServer.workspaces.asyncCompactMap({ await $0.semanticIndexManager }) {
       status = status.merging(with: await indexManager.progressStatus)
     }
 

Sources/SourceKitLSP/Rename.swift

@@ -290,7 +290,7 @@ extension SourceKitLSPServer {
     // location. This way we are still able to rename occurrences in files where eg. only one line has been modified but
     // all the line:column locations of occurrences are still up-to-date.
     // This should match the check level in prepareRename.
-    guard let usr = renameResult.usr, let index = workspace.index(checkedFor: .deletedFiles) else {
+    guard let usr = renameResult.usr, let index = await workspace.index(checkedFor: .deletedFiles) else {
       // We don't have enough information to perform a cross-file rename.
       return renameResult.edits
     }
@@ -442,7 +442,7 @@ extension SourceKitLSPServer {
     var prepareRenameResult = languageServicePrepareRename.prepareRename
 
     guard
-      let index = workspace.index(checkedFor: .deletedFiles),
+      let index = await workspace.index(checkedFor: .deletedFiles),
       let usr = languageServicePrepareRename.usr,
       let oldName = try await self.getCrossLanguageName(forUsr: usr, workspace: workspace, index: index),
       var definitionName = oldName.definitionName

Sources/SourceKitLSP/SourceKitIndexDelegate.swift

@@ -19,14 +19,16 @@ import SwiftExtensions
 /// `IndexDelegate` for the SourceKit workspace.
 actor SourceKitIndexDelegate: IndexDelegate {
   /// Registered `MainFilesDelegate`s to notify when main files change.
-  var mainFilesChangedCallbacks: [@Sendable () async -> Void] = []
+  var mainFilesChangedCallback: @Sendable () async -> Void
 
   /// The count of pending unit events. Whenever this transitions to 0, it represents a time where
   /// the index finished processing known events. Of course, that may have already changed by the
   /// time we are notified.
   let pendingUnitCount = AtomicInt32(initialValue: 0)
 
-  package init() {}
+  package init(mainFilesChangedCallback: @escaping @Sendable () async -> Void) {
+    self.mainFilesChangedCallback = mainFilesChangedCallback
+  }
 
   nonisolated package func processingAddedPending(_ count: Int) {
     pendingUnitCount.value += Int32(count)
@@ -53,13 +55,6 @@ actor SourceKitIndexDelegate: IndexDelegate {
 
   private func indexChanged() async {
     logger.debug("IndexStoreDB changed")
-    for callback in mainFilesChangedCallbacks {
-      await callback()
-    }
-  }
-
-  /// Register a delegate to receive notifications when main files change.
-  package func addMainFileChangedCallback(_ callback: @escaping @Sendable () async -> Void) {
-    mainFilesChangedCallbacks.append(callback)
+    await mainFilesChangedCallback()
   }
 }

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1257,8 +1257,6 @@ extension SourceKitLSPServer {
         taskGroup.addTask {
           await orLog("Shutting down build server") {
             await workspace.buildServerManager.shutdown()
-            await workspace.buildServerManager.setMainFilesProvider(nil)
-            workspace.closeIndex()
           }
         }
       }
@@ -1306,7 +1304,7 @@ extension SourceKitLSPServer {
     }
     let isIndexing =
       await workspaces
-      .compactMap(\.semanticIndexManager)
+      .asyncCompactMap { await $0.semanticIndexManager }
       .asyncContains { await $0.progressStatus != .upToDate }
     return IsIndexingResponse(indexing: isIndexing)
   }
@@ -1568,7 +1566,7 @@ extension SourceKitLSPServer {
       target = nil
     }
     let didPrepareTarget: Bool?
-    if request.prepareTarget, let target, let semanticIndexManager = workspace.semanticIndexManager {
+    if request.prepareTarget, let target, let semanticIndexManager = await workspace.semanticIndexManager {
       didPrepareTarget = await semanticIndexManager.prepareTargetsForSourceKitOptions(target: target)
     } else {
       didPrepareTarget = nil
@@ -1669,7 +1667,7 @@ extension SourceKitLSPServer {
     }
     var symbolsAndIndex: [(symbol: SymbolOccurrence, index: CheckedIndex)] = []
     for workspace in workspaces {
-      guard let index = workspace.index(checkedFor: .deletedFiles) else {
+      guard let index = await workspace.index(checkedFor: .deletedFiles) else {
         continue
       }
       var symbolOccurrences: [SymbolOccurrence] = []
@@ -2084,7 +2082,7 @@ extension SourceKitLSPServer {
       // overrides.
       if let location = locations.only,
         let usr = symbol.usr,
-        let index = workspace.index(checkedFor: .deletedFiles),
+        let index = await workspace.index(checkedFor: .deletedFiles),
         await isAtCanonicalOriginatorLocation(location)
       {
         let baseUSRs = index.occurrences(ofUSR: usr, roles: .overrideOf).flatMap {

Sources/SourceKitLSP/TestDiscovery.swift

@@ -207,7 +207,7 @@ extension SourceKitLSPServer {
     //  - For all files that don't have any in-memory modifications, include swift-testing tests from the syntactic test
     //    index.
     //  - All files that have in-memory modifications are syntactically scanned for tests here.
-    let index = workspace.index(checkedFor: .inMemoryModifiedFiles(documentManager))
+    let index = await workspace.index(checkedFor: .inMemoryModifiedFiles(documentManager))
 
     // TODO: Remove this workaround once https://github.com/swiftlang/swift/issues/75600 is fixed
     func documentManagerHasInMemoryModifications(_ uri: DocumentURI) -> Bool {
@@ -252,7 +252,7 @@ extension SourceKitLSPServer {
     )
     let filesWithTestsFromSemanticIndex = Set(testsFromSemanticIndex.map(\.testItem.location.uri))
 
-    let indexOnlyDiscardingDeletedFiles = workspace.index(checkedFor: .deletedFiles)
+    let indexOnlyDiscardingDeletedFiles = await workspace.index(checkedFor: .deletedFiles)
 
     let syntacticTestsToInclude =
       testsFromSyntacticIndex
@@ -334,7 +334,7 @@ extension SourceKitLSPServer {
     let indexCheckLevel: IndexCheckLevel =
       syntacticTests == nil ? .deletedFiles : .inMemoryModifiedFiles(documentManager)
 
-    if let index = workspace.index(checkedFor: indexCheckLevel) {
+    if let index = await workspace.index(checkedFor: indexCheckLevel) {
       var syntacticSwiftTestingTests: [AnnotatedTestItem] {
         syntacticTests?.filter { $0.testItem.style == TestStyle.swiftTesting } ?? []
       }