Merge pull request #1328 from ahoppen/parallel-indexing

Author: ahoppen

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -69,7 +69,11 @@ public final class TestSourceKitLSPClient: MessageHandler {
   ///
   /// Conceptually, this is an array of `RequestHandler<any RequestType>` but
   /// since we can't express this in the Swift type system, we use `[Any]`.
-  private nonisolated(unsafe) var requestHandlers = ThreadSafeBox<[Any]>(initialValue: [])
+  ///
+  /// `isOneShort` if the request handler should only serve a single request and should be removed from
+  /// `requestHandlers` after it has been called.
+  private nonisolated(unsafe) var requestHandlers: ThreadSafeBox<[(requestHandler: Any, isOneShot: Bool)]> =
+    ThreadSafeBox(initialValue: [])
 
   /// A closure that is called when the `TestSourceKitLSPClient` is destructed.
   ///
@@ -147,7 +151,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
         throw ConflictingDiagnosticsError()
       }
       capabilities.textDocument!.diagnostic = .init(dynamicRegistration: true)
-      self.handleNextRequest { (request: RegisterCapabilityRequest) in
+      self.handleSingleRequest { (request: RegisterCapabilityRequest) in
         XCTAssertEqual(request.registrations.only?.method, DocumentDiagnosticsRequest.method)
         return VoidResponse()
       }
@@ -258,22 +262,22 @@ public final class TestSourceKitLSPClient: MessageHandler {
     }
   }
 
-  /// Handle the next request that is sent to the client with the given handler.
-  ///
-  /// By default, `TestSourceKitLSPClient` emits an `XCTFail` if a request is sent
-  /// to the client, since it doesn't know how to handle it. This allows the
-  /// simulation of a single request's handling on the client.
+  /// Handle the next request of the given type that is sent to the client.
   ///
-  /// If the next request that is sent to the client is of a different kind than
-  /// the given handler, `TestSourceKitLSPClient` will emit an `XCTFail`.
-  public func handleNextRequest<R: RequestType>(_ requestHandler: @escaping RequestHandler<R>) {
-    requestHandlers.value.append(requestHandler)
+  /// The request handler will only handle a single request. If the request is called again, the request handler won't
+  /// call again
+  public func handleSingleRequest<R: RequestType>(_ requestHandler: @escaping RequestHandler<R>) {
+    requestHandlers.value.append((requestHandler: requestHandler, isOneShot: true))
+  }
+
+  /// Handle all requests of the given type that are sent to the client.
+  public func handleMultipleRequests<R: RequestType>(_ requestHandler: @escaping RequestHandler<R>) {
+    requestHandlers.value.append((requestHandler: requestHandler, isOneShot: false))
   }
 
   // MARK: - Conformance to MessageHandler
 
-  /// - Important: Implementation detail of `TestSourceKitLSPServer`. Do not call
-  ///   from tests.
+  /// - Important: Implementation detail of `TestSourceKitLSPServer`. Do not call from tests.
   public func handle(_ params: some NotificationType) {
     notificationYielder.yield(params)
   }
@@ -285,19 +289,21 @@ public final class TestSourceKitLSPClient: MessageHandler {
     reply: @escaping (LSPResult<Request.Response>) -> Void
   ) {
     requestHandlers.withLock { requestHandlers in
-      let requestHandlerAndIndex = requestHandlers.enumerated().compactMap {
-        (index, handler) -> (RequestHandler<Request>, Int)? in
-        guard let handler = handler as? RequestHandler<Request> else {
+      let requestHandlerIndexAndIsOneShot = requestHandlers.enumerated().compactMap {
+        (index, handlerAndIsOneShot) -> (RequestHandler<Request>, Int, Bool)? in
+        guard let handler = handlerAndIsOneShot.requestHandler as? RequestHandler<Request> else {
           return nil
         }
-        return (handler, index)
+        return (handler, index, handlerAndIsOneShot.isOneShot)
       }.first
-      guard let (requestHandler, index) = requestHandlerAndIndex else {
+      guard let (requestHandler, index, isOneShot) = requestHandlerIndexAndIsOneShot else {
         reply(.failure(.methodNotFound(Request.method)))
         return
       }
       reply(.success(requestHandler(params)))
-      requestHandlers.remove(at: index)
+      if isOneShot {
+        requestHandlers.remove(at: index)
+      }
     }
   }
 

Sources/SemanticIndex/PreparationTaskDescription.swift

@@ -78,7 +78,6 @@ public struct PreparationTaskDescription: IndexTaskDescription {
       subsystem: "org.swift.sourcekit-lsp.indexing",
       scope: "preparation-\(id % 100)"
     ) {
-      await testHooks.preparationTaskDidStart?(self)
       let targetsToPrepare = await targetsToPrepare.asyncFilter {
         await !preparationUpToDateStatus.isUpToDate($0)
       }.sorted(by: {
@@ -87,6 +86,7 @@ public struct PreparationTaskDescription: IndexTaskDescription {
       if targetsToPrepare.isEmpty {
         return
       }
+      await testHooks.preparationTaskDidStart?(self)
 
       let targetsToPrepareDescription =
         targetsToPrepare

Sources/SemanticIndex/TestHooks.swift

@@ -16,16 +16,20 @@ public struct IndexTestHooks: Sendable {
 
   public var preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)?
 
+  public var updateIndexStoreTaskDidStart: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)?
+
   /// A callback that is called when an index task finishes.
   public var updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)?
 
   public init(
     preparationTaskDidStart: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
     preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
+    updateIndexStoreTaskDidStart: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)? = nil,
     updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)? = nil
   ) {
     self.preparationTaskDidStart = preparationTaskDidStart
     self.preparationTaskDidFinish = preparationTaskDidFinish
+    self.updateIndexStoreTaskDidStart = updateIndexStoreTaskDidStart
     self.updateIndexStoreTaskDidFinish = updateIndexStoreTaskDidFinish
   }
 }

Sources/SemanticIndex/UpdateIndexStoreTaskDescription.swift

@@ -22,7 +22,7 @@ import class TSCBasic.Process
 
 private nonisolated(unsafe) var updateIndexStoreIDForLogging = AtomicUInt32(initialValue: 1)
 
-enum FileToIndex: CustomLogStringConvertible {
+public enum FileToIndex: CustomLogStringConvertible {
   /// A non-header file
   case indexableFile(DocumentURI)
 
@@ -33,7 +33,7 @@ enum FileToIndex: CustomLogStringConvertible {
   ///
   /// This file might be a header file that doesn't have build settings associated with it. For the actual compiler
   /// invocation that updates the index store, the `mainFile` should be used.
-  var sourceFile: DocumentURI {
+  public var sourceFile: DocumentURI {
     switch self {
     case .indexableFile(let uri): return uri
     case .headerFile(header: let header, mainFile: _): return header
@@ -51,7 +51,7 @@ enum FileToIndex: CustomLogStringConvertible {
     }
   }
 
-  var description: String {
+  public var description: String {
     switch self {
     case .indexableFile(let uri):
       return uri.description
@@ -60,7 +60,7 @@ enum FileToIndex: CustomLogStringConvertible {
     }
   }
 
-  var redactedDescription: String {
+  public var redactedDescription: String {
     switch self {
     case .indexableFile(let uri):
       return uri.redactedDescription
@@ -71,9 +71,9 @@ enum FileToIndex: CustomLogStringConvertible {
 }
 
 /// A file to index and the target in which the file should be indexed.
-struct FileAndTarget {
-  let file: FileToIndex
-  let target: ConfiguredTarget
+public struct FileAndTarget: Sendable {
+  public let file: FileToIndex
+  public let target: ConfiguredTarget
 }
 
 /// Describes a task to index a set of source files.
@@ -84,7 +84,7 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
   public let id = updateIndexStoreIDForLogging.fetchAndIncrement()
 
   /// The files that should be indexed.
-  private let filesToIndex: [FileAndTarget]
+  public let filesToIndex: [FileAndTarget]
 
   /// The build system manager that is used to get the toolchain and build settings for the files to index.
   private let buildSystemManager: BuildSystemManager
@@ -140,6 +140,8 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
     ) {
       let startDate = Date()
 
+      await testHooks.updateIndexStoreTaskDidStart?(self)
+
       let filesToIndexDescription = filesToIndex.map {
         $0.file.sourceFile.fileURL?.lastPathComponent ?? $0.file.sourceFile.stringValue
       }
@@ -166,9 +168,7 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
   ) -> [TaskDependencyAction<UpdateIndexStoreTaskDescription>] {
     let selfMainFiles = Set(filesToIndex.map(\.file.mainFile))
     return currentlyExecutingTasks.compactMap { (other) -> TaskDependencyAction<UpdateIndexStoreTaskDescription>? in
-      guard
-        !other.filesToIndex.lazy.map(\.file.mainFile).contains(where: { selfMainFiles.contains($0) })
-      else {
+      if !other.filesToIndex.lazy.map(\.file.mainFile).contains(where: { selfMainFiles.contains($0) }) {
         // Disjoint sets of files can be indexed concurrently.
         return nil
       }

Sources/sourcekit-lsp/CMakeLists.txt

@@ -5,6 +5,7 @@ target_link_libraries(sourcekit-lsp PRIVATE
   Diagnose
   LanguageServerProtocol
   LanguageServerProtocolJSONRPC
+  SemanticIndex
   SKCore
   SKSupport
   SourceKitLSP

Tests/SourceKitDTests/CrashRecoveryTests.swift

@@ -75,7 +75,7 @@ final class CrashRecoveryTests: XCTestCase {
       "Sanity check failed. The Hover response did not contain foo(), even before crashing sourcekitd. Received response: \(String(describing: preCrashHoverResponse))"
     )
 
-    testClient.handleNextRequest { (request: CreateWorkDoneProgressRequest) -> VoidResponse in
+    testClient.handleSingleRequest { (request: CreateWorkDoneProgressRequest) -> VoidResponse in
       return VoidResponse()
     }