Merge pull request #1355 from ahoppen/wait-for-preparation-before-diags

Wait for target to be prepared before returning diagnostics

Author: ahoppen

Package.swift

@@ -367,6 +367,7 @@ let package = Package(
       name: "SourceKitLSPTests",
       dependencies: [
         "BuildServerProtocol",
+        "CAtomics",
         "LSPLogging",
         "LSPTestSupport",
         "LanguageServerProtocol",

Sources/SKTestSupport/MultiFileTestProject.swift

@@ -82,6 +82,7 @@ public class MultiFileTestProject {
     workspaces: (URL) async throws -> [WorkspaceFolder] = { [WorkspaceFolder(uri: DocumentURI($0))] },
     capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
+    enableBackgroundIndexing: Bool = false,
     usePullDiagnostics: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     cleanUp: (() -> Void)? = nil,
@@ -117,6 +118,7 @@ public class MultiFileTestProject {
       serverOptions: serverOptions,
       capabilities: capabilities,
       usePullDiagnostics: usePullDiagnostics,
+      enableBackgroundIndexing: enableBackgroundIndexing,
       workspaceFolders: workspaces(scratchDirectory),
       preInitialization: preInitialization,
       cleanUp: { [scratchDirectory] in

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -44,9 +44,10 @@ public class SwiftPMTestProject: MultiFileTestProject {
     allowBuildFailure: Bool = false,
     capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
+    enableBackgroundIndexing: Bool = false,
+    usePullDiagnostics: Bool = true,
     pollIndex: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
-    usePullDiagnostics: Bool = true,
     cleanUp: (() -> Void)? = nil,
     testName: String = #function
   ) async throws {
@@ -71,6 +72,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
       workspaces: workspaces,
       capabilities: capabilities,
       serverOptions: serverOptions,
+      enableBackgroundIndexing: enableBackgroundIndexing,
       usePullDiagnostics: usePullDiagnostics,
       preInitialization: preInitialization,
       cleanUp: cleanUp,

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -88,7 +88,8 @@ public final class TestSourceKitLSPClient: MessageHandler {
   ///     `true` by default
   ///   - initializationOptions: Initialization options to pass to the SourceKit-LSP server.
   ///   - capabilities: The test client's capabilities.
-  ///   - usePullDiagnostics: Whether to use push diagnostics or use push-based diagnostics
+  ///   - usePullDiagnostics: Whether to use push diagnostics or use push-based diagnostics.
+  ///   - enableBackgroundIndexing: Whether background indexing should be enabled in the project.
   ///   - workspaceFolders: Workspace folders to open.
   ///   - preInitialization: A closure that is called after the test client is created but before SourceKit-LSP is
   ///     initialized. This can be used to eg. register request handlers.
@@ -102,6 +103,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
     initializationOptions: LSPAny? = nil,
     capabilities: ClientCapabilities = ClientCapabilities(),
     usePullDiagnostics: Bool = true,
+    enableBackgroundIndexing: Bool = false,
     workspaceFolders: [WorkspaceFolder]? = nil,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     cleanUp: @Sendable @escaping () -> Void = {}
@@ -115,6 +117,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
     if let moduleCache {
       serverOptions.buildSetup.flags.swiftCompilerFlags += ["-module-cache-path", moduleCache.path]
     }
+    serverOptions.indexOptions.enableBackgroundIndexing = enableBackgroundIndexing
 
     var notificationYielder: AsyncStream<any NotificationType>.Continuation!
     self.notifications = AsyncStream { continuation in
@@ -155,8 +158,8 @@ public final class TestSourceKitLSPClient: MessageHandler {
         XCTAssertEqual(request.registrations.only?.method, DocumentDiagnosticsRequest.method)
         return VoidResponse()
       }
-      preInitialization?(self)
     }
+    preInitialization?(self)
     if initialize {
       _ = try await self.send(
         InitializeRequest(
@@ -193,12 +196,22 @@ public final class TestSourceKitLSPClient: MessageHandler {
   /// Send the request to `server` and return the request result.
   public func send<R: RequestType>(_ request: R) async throws -> R.Response {
     return try await withCheckedThrowingContinuation { continuation in
-      server.handle(request, id: .number(Int(nextRequestID.fetchAndIncrement()))) { result in
+      self.send(request) { result in
         continuation.resume(with: result)
       }
     }
   }
 
+  /// Send the request to `server` and return the result via a completion handler.
+  ///
+  /// This version of the `send` function should only be used if some action needs to be performed after the request is
+  /// sent but before it returns a result.
+  public func send<R: RequestType>(_ request: R, completionHandler: @escaping (LSPResult<R.Response>) -> Void) {
+    server.handle(request, id: .number(Int(nextRequestID.fetchAndIncrement()))) { result in
+      completionHandler(result)
+    }
+  }
+
   /// Send the notification to `server`.
   public func send(_ notification: some NotificationType) {
     server.handle(notification)

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -357,10 +357,7 @@ public final actor SemanticIndexManager {
   /// Schedule preparation of the target that contains the given URI, building all modules that the file depends on.
   ///
   /// This is intended to be called when the user is interacting with the document at the given URI.
-  public func schedulePreparationForEditorFunctionality(
-    of uri: DocumentURI,
-    priority: TaskPriority? = nil
-  ) {
+  public func schedulePreparationForEditorFunctionality(of uri: DocumentURI, priority: TaskPriority? = nil) {
     if inProgressPrepareForEditorTask?.document == uri {
       // We are already preparing this document, so nothing to do. This is necessary to avoid the following scenario:
       // Determining the canonical configured target for a document takes 1s and we get a new document request for the
@@ -371,13 +368,7 @@ public final actor SemanticIndexManager {
     let id = UUID()
     let task = Task(priority: priority) {
       await withLoggingScope("preparation") {
-        guard let target = await buildSystemManager.canonicalConfiguredTarget(for: uri) else {
-          return
-        }
-        if Task.isCancelled {
-          return
-        }
-        await self.prepare(targets: [target], priority: priority)
+        await self.prepareFileForEditorFunctionality(uri)
         if inProgressPrepareForEditorTask?.id == id {
           inProgressPrepareForEditorTask = nil
           self.indexProgressStatusDidChange()
@@ -389,6 +380,20 @@ public final actor SemanticIndexManager {
     self.indexProgressStatusDidChange()
   }
 
+  /// Prepare the target that the given file is in, building all modules that the file depends on. Returns when
+  /// preparation has finished.
+  ///
+  /// If file's target is known to be up-to-date, this returns almost immediately.
+  public func prepareFileForEditorFunctionality(_ uri: DocumentURI) async {
+    guard let target = await buildSystemManager.canonicalConfiguredTarget(for: uri) else {
+      return
+    }
+    if Task.isCancelled {
+      return
+    }
+    await self.prepare(targets: [target], priority: nil)
+  }
+
   // MARK: - Helper functions
 
   /// Prepare the given targets for indexing.

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -17,6 +17,7 @@ import LSPLogging
 import LanguageServerProtocol
 import SKCore
 import SKSupport
+import SemanticIndex
 import SourceKitD
 import SwiftParser
 import SwiftParserDiagnostics
@@ -123,6 +124,9 @@ public actor SwiftLanguageService: LanguageService, Sendable {
 
   let syntaxTreeManager = SyntaxTreeManager()
 
+  /// The `semanticIndexManager` of the workspace this language service was created for.
+  private let semanticIndexManager: SemanticIndexManager?
+
   nonisolated var keys: sourcekitd_api_keys { return sourcekitd.keys }
   nonisolated var requests: sourcekitd_api_requests { return sourcekitd.requests }
   nonisolated var values: sourcekitd_api_values { return sourcekitd.values }
@@ -192,6 +196,7 @@ public actor SwiftLanguageService: LanguageService, Sendable {
     self.swiftFormat = toolchain.swiftFormat
     self.sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitd)
     self.capabilityRegistry = workspace.capabilityRegistry
+    self.semanticIndexManager = workspace.semanticIndexManager
     self.serverOptions = options
     self.documentManager = DocumentManager()
     self.state = .connected
@@ -875,6 +880,7 @@ extension SwiftLanguageService {
 
   public func documentDiagnostic(_ req: DocumentDiagnosticsRequest) async throws -> DocumentDiagnosticReport {
     do {
+      await semanticIndexManager?.prepareFileForEditorFunctionality(req.textDocument.uri)
       let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
       let buildSettings = await self.buildSettings(for: req.textDocument.uri)
       let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -18,10 +18,6 @@ import SemanticIndex
 import SourceKitLSP
 import XCTest
 
-fileprivate let backgroundIndexingOptions = SourceKitLSPServer.Options(
-  indexOptions: IndexOptions(enableBackgroundIndexing: true)
-)
-
 final class BackgroundIndexingTests: XCTestCase {
   func testBackgroundIndexingOfSingleFile() async throws {
     let project = try await SwiftPMTestProject(
@@ -33,7 +29,7 @@ final class BackgroundIndexingTests: XCTestCase {
         }
         """
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -76,7 +72,7 @@ final class BackgroundIndexingTests: XCTestCase {
         }
         """,
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -134,7 +130,7 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -166,7 +162,7 @@ final class BackgroundIndexingTests: XCTestCase {
   }
 
   func testBackgroundIndexingHappensWithLowPriority() async throws {
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     serverOptions.indexTestHooks.preparationTaskDidFinish = { taskDescription in
       XCTAssert(Task.currentPriority == .low, "\(taskDescription) ran with priority \(Task.currentPriority)")
     }
@@ -199,6 +195,7 @@ final class BackgroundIndexingTests: XCTestCase {
         )
         """,
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       pollIndex: false
     )
 
@@ -249,7 +246,7 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let dependencyUrl = try XCTUnwrap(
@@ -300,7 +297,7 @@ final class BackgroundIndexingTests: XCTestCase {
         }
         """,
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.c")
@@ -339,7 +336,7 @@ final class BackgroundIndexingTests: XCTestCase {
     let receivedReportProgressNotification = self.expectation(
       description: "Received work done progress saying indexing"
     )
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     serverOptions.indexTestHooks = IndexTestHooks(
       buildGraphGenerationDidFinish: {
         await self.fulfillment(of: [receivedBeginProgressNotification], timeout: defaultTimeout)
@@ -359,6 +356,7 @@ final class BackgroundIndexingTests: XCTestCase {
       ],
       capabilities: ClientCapabilities(window: WindowClientCapabilities(workDoneProgress: true)),
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       pollIndex: false,
       preInitialization: { testClient in
         testClient.handleMultipleRequests { (request: CreateWorkDoneProgressRequest) in
@@ -419,7 +417,7 @@ final class BackgroundIndexingTests: XCTestCase {
         """,
         "MyOtherFile.swift": "",
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -486,7 +484,7 @@ final class BackgroundIndexingTests: XCTestCase {
         #include "Header.h"
         """,
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("Header.h", language: .c)
@@ -544,7 +542,7 @@ final class BackgroundIndexingTests: XCTestCase {
 
   func testPrepareTargetAfterEditToDependency() async throws {
     try await SkipUnless.swiftpmStoresModulesInSubdirectory()
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     let expectedPreparationTracker = ExpectedIndexTaskTracker(expectedPreparations: [
       [
         ExpectedPreparation(targetID: "LibA", runDestinationID: "dummy"),
@@ -581,6 +579,7 @@ final class BackgroundIndexingTests: XCTestCase {
         """,
       capabilities: ClientCapabilities(window: WindowClientCapabilities(workDoneProgress: true)),
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       cleanUp: { expectedPreparationTracker.keepAlive() }
     )
 
@@ -653,7 +652,7 @@ final class BackgroundIndexingTests: XCTestCase {
     let libDPreparedForEditing = self.expectation(description: "LibD prepared for editing")
 
     try await SkipUnless.swiftpmStoresModulesInSubdirectory()
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     let expectedPreparationTracker = ExpectedIndexTaskTracker(expectedPreparations: [
       // Preparation of targets during the initial of the target
       [
@@ -705,6 +704,7 @@ final class BackgroundIndexingTests: XCTestCase {
         )
         """,
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       cleanUp: { expectedPreparationTracker.keepAlive() }
     )
 
@@ -734,7 +734,7 @@ final class BackgroundIndexingTests: XCTestCase {
       files: [
         "MyFile.swift": ""
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
     let targetPrepareNotification = try await project.testClient.nextNotification(ofType: LogMessageNotification.self)
     XCTAssert(
@@ -748,13 +748,13 @@ final class BackgroundIndexingTests: XCTestCase {
     )
   }
 
-  func testPreparationHappensInParallel() async throws {
+  func testIndexingHappensInParallel() async throws {
     try await SkipUnless.swiftpmStoresModulesInSubdirectory()
 
     let fileAIndexingStarted = self.expectation(description: "FileA indexing started")
     let fileBIndexingStarted = self.expectation(description: "FileB indexing started")
 
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     let expectedIndexTaskTracker = ExpectedIndexTaskTracker(
       expectedIndexStoreUpdates: [
         [
@@ -787,6 +787,7 @@ final class BackgroundIndexingTests: XCTestCase {
         "FileB.swift": "",
       ],
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       cleanUp: { expectedIndexTaskTracker.keepAlive() }
     )
   }
@@ -817,10 +818,10 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
-    var otherClientOptions = backgroundIndexingOptions
+    var otherClientOptions = SourceKitLSPServer.Options.testDefault
     otherClientOptions.indexTestHooks = IndexTestHooks(
       preparationTaskDidStart: { taskDescription in
         XCTFail("Did not expect any target preparation, got \(taskDescription.targetsToPrepare)")
@@ -831,6 +832,7 @@ final class BackgroundIndexingTests: XCTestCase {
     )
     let otherClient = try await TestSourceKitLSPClient(
       serverOptions: otherClientOptions,
+      enableBackgroundIndexing: true,
       workspaceFolders: [
         WorkspaceFolder(uri: DocumentURI(project.scratchDirectory))
       ]