Wait for cancellation to propagate in testDontReturnEmptyDiagnosticsIfDiagnosticRequestIsCancelled

Sources/Diagnose/RunSourcekitdRequestCommand.swift

@@ -55,7 +55,8 @@ public struct RunSourceKitdRequestCommand: AsyncParsableCommand {
         throw ExitCode(1)
       }
     let sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(
-      dylibPath: try! AbsolutePath(validating: sourcekitdPath)
+      dylibPath: try! AbsolutePath(validating: sourcekitdPath),
+      testHooks: SourceKitDTestHooks()
     )
 
     if let lineColumn = position?.split(separator: ":", maxSplits: 2).map(Int.init),

Sources/SourceKitD/DynamicallyLoadedSourceKitD.swift

@@ -25,17 +25,27 @@ extension sourcekitd_api_keys: @unchecked Sendable {}
 extension sourcekitd_api_requests: @unchecked Sendable {}
 extension sourcekitd_api_values: @unchecked Sendable {}
 
+public struct SourceKitDTestHooks: Sendable {
+  public var sourcekitdRequestDidStart: (@Sendable (SKDRequestDictionary) -> Void)?
+
+  public init(sourcekitdRequestDidStart: (@Sendable (SKDRequestDictionary) -> Void)? = nil) {
+    self.sourcekitdRequestDidStart = sourcekitdRequestDidStart
+  }
+}
+
 /// Wrapper for sourcekitd, taking care of initialization, shutdown, and notification handler
 /// multiplexing.
 ///
 /// Users of this class should not call the api functions `initialize`, `shutdown`, or
 /// `set_notification_handler`, which are global state managed internally by this class.
 public actor DynamicallyLoadedSourceKitD: SourceKitD {
   /// The path to the sourcekitd dylib.
-  public let path: AbsolutePath
+  private let path: AbsolutePath
 
   /// The handle to the dylib.
-  let dylib: DLHandle
+  private let dylib: DLHandle
+
+  public let testHooks: SourceKitDTestHooks
 
   /// The sourcekitd API functions.
   public let api: sourcekitd_api_functions_t
@@ -54,18 +64,23 @@ public actor DynamicallyLoadedSourceKitD: SourceKitD {
   /// List of notification handlers that will be called for each notification.
   private var notificationHandlers: [WeakSKDNotificationHandler] = []
 
-  public static func getOrCreate(dylibPath: AbsolutePath) async throws -> SourceKitD {
+  /// If there is already a `sourcekitd` instance from the given return it, otherwise create a new one.
+  ///
+  /// `testHooks` are only considered when an instance is being created. If a sourcekitd instance at the given path
+  /// already exists, its test hooks will be used.
+  public static func getOrCreate(dylibPath: AbsolutePath, testHooks: SourceKitDTestHooks) async throws -> SourceKitD {
     try await SourceKitDRegistry.shared
-      .getOrAdd(dylibPath, create: { try DynamicallyLoadedSourceKitD(dylib: dylibPath) })
+      .getOrAdd(dylibPath, create: { try DynamicallyLoadedSourceKitD(dylib: dylibPath, testHooks: testHooks) })
   }
 
-  init(dylib path: AbsolutePath) throws {
+  init(dylib path: AbsolutePath, testHooks: SourceKitDTestHooks) throws {
     self.path = path
     #if os(Windows)
     self.dylib = try dlopen(path.pathString, mode: [])
     #else
     self.dylib = try dlopen(path.pathString, mode: [.lazy, .local, .first])
     #endif
+    self.testHooks = testHooks
     self.api = try sourcekitd_api_functions_t(self.dylib)
     self.keys = sourcekitd_api_keys(api: self.api)
     self.requests = sourcekitd_api_requests(api: self.api)

Sources/SourceKitD/SourceKitD.swift

@@ -30,6 +30,8 @@ extension sourcekitd_api_request_handle_t: @unchecked Sendable {}
 /// *Implementors* are expected to handle initialization and shutdown, e.g. during `init` and
 /// `deinit` or by wrapping an existing sourcekitd session that outlives this object.
 public protocol SourceKitD: AnyObject, Sendable {
+  var testHooks: SourceKitDTestHooks { get }
+
   /// The sourcekitd API functions.
   var api: sourcekitd_api_functions_t { get }
 
@@ -95,18 +97,20 @@ extension SourceKitD {
   ///   - req: The request to send to sourcekitd.
   ///   - fileContents: The contents of the file that the request operates on. If sourcekitd crashes, the file contents
   ///     will be logged.
-  public func send(_ req: SKDRequestDictionary, fileContents: String?) async throws -> SKDResponseDictionary {
-    log(request: req)
+  public func send(_ request: SKDRequestDictionary, fileContents: String?) async throws -> SKDResponseDictionary {
+    log(request: request)
+
+    testHooks.sourcekitdRequestDidStart?(request)
 
     let sourcekitdResponse: SKDResponse = try await withCancellableCheckedThrowingContinuation { continuation in
       var handle: sourcekitd_api_request_handle_t? = nil
-      api.send_request(req.dict, &handle) { response in
+      api.send_request(request.dict, &handle) { response in
         continuation.resume(returning: SKDResponse(response!, sourcekitd: self))
       }
       return handle
     } cancel: { handle in
       if let handle {
-        logRequestCancellation(request: req)
+        logRequestCancellation(request: request)
         api.cancel_request(handle)
       }
     }
@@ -115,7 +119,7 @@ extension SourceKitD {
 
     guard let dict = sourcekitdResponse.value else {
       if sourcekitdResponse.error == .connectionInterrupted {
-        log(crashedRequest: req, fileContents: fileContents)
+        log(crashedRequest: request, fileContents: fileContents)
       }
       throw sourcekitdResponse.error!
     }

Sources/SourceKitD/SourceKitDRegistry.swift

@@ -35,10 +35,10 @@ extension NSLock {
 public actor SourceKitDRegistry {
 
   /// Mapping from path to active SourceKitD instance.
-  var active: [AbsolutePath: SourceKitD] = [:]
+  private var active: [AbsolutePath: SourceKitD] = [:]
 
   /// Instances that have been unregistered, but may be resurrected if accessed before destruction.
-  var cemetary: [AbsolutePath: WeakSourceKitD] = [:]
+  private var cemetary: [AbsolutePath: WeakSourceKitD] = [:]
 
   /// Initialize an empty registry.
   public init() {}
@@ -79,14 +79,8 @@ public actor SourceKitDRegistry {
     }
     return existing
   }
-
-  /// Remove all SourceKitD instances, including weak ones.
-  public func clear() {
-    active.removeAll()
-    cemetary.removeAll()
-  }
 }
 
-struct WeakSourceKitD {
+fileprivate struct WeakSourceKitD {
   weak var value: SourceKitD?
 }

Sources/SourceKitLSP/SourceKitLSPServer+Options.swift

@@ -15,12 +15,12 @@ import LanguageServerProtocol
 import SKCore
 import SKSupport
 import SemanticIndex
+import SourceKitD
 
 import struct TSCBasic.AbsolutePath
 import struct TSCBasic.RelativePath
 
 extension SourceKitLSPServer {
-
   /// Configuration options for the SourceKitServer.
   public struct Options: Sendable {
     /// Additional compiler flags (e.g. `-Xswiftc` for SwiftPM projects) and other build-related
@@ -52,6 +52,8 @@ extension SourceKitLSPServer {
     /// Experimental features that are enabled.
     public var experimentalFeatures: Set<ExperimentalFeature>
 
+    public var sourcekitdTestHooks: SourceKitDTestHooks
+
     public var indexTestHooks: IndexTestHooks
 
     public init(
@@ -63,6 +65,7 @@ extension SourceKitLSPServer {
       generatedInterfacesPath: AbsolutePath = defaultDirectoryForGeneratedInterfaces,
       swiftPublishDiagnosticsDebounceDuration: TimeInterval = 2, /* 2s */
       experimentalFeatures: Set<ExperimentalFeature> = [],
+      sourcekitdTestHooks: SourceKitDTestHooks = SourceKitDTestHooks(),
       indexTestHooks: IndexTestHooks = IndexTestHooks()
     ) {
       self.buildSetup = buildSetup
@@ -73,6 +76,7 @@ extension SourceKitLSPServer {
       self.generatedInterfacesPath = generatedInterfacesPath
       self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
       self.experimentalFeatures = experimentalFeatures
+      self.sourcekitdTestHooks = sourcekitdTestHooks
       self.indexTestHooks = indexTestHooks
     }
   }

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -194,7 +194,10 @@ public actor SwiftLanguageService: LanguageService, Sendable {
     guard let sourcekitd = toolchain.sourcekitd else { return nil }
     self.sourceKitLSPServer = sourceKitLSPServer
     self.swiftFormat = toolchain.swiftFormat
-    self.sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitd)
+    self.sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(
+      dylibPath: sourcekitd,
+      testHooks: options.sourcekitdTestHooks
+    )
     self.capabilityRegistry = workspace.capabilityRegistry
     self.semanticIndexManager = workspace.semanticIndexManager
     self.serverOptions = options

Tests/DiagnoseTests/DiagnoseTests.swift

@@ -318,7 +318,8 @@ private class InProcessSourceKitRequestExecutor: SourceKitRequestExecutor {
     logger.info("Sending request: \(requestString)")
 
     let sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(
-      dylibPath: try! AbsolutePath(validating: sourcekitd.path)
+      dylibPath: try! AbsolutePath(validating: sourcekitd.path),
+      testHooks: SourceKitDTestHooks()
     )
     let response = try await sourcekitd.run(requestYaml: requestString)
 

Tests/SourceKitDTests/SourceKitDRegistryTests.swift

@@ -63,6 +63,7 @@ private nonisolated(unsafe) var nextToken = AtomicUInt32(initialValue: 0)
 
 final class FakeSourceKitD: SourceKitD {
   let token: UInt32
+  var testHooks: SourceKitDTestHooks { SourceKitDTestHooks() }
   var api: sourcekitd_api_functions_t { fatalError() }
   var keys: sourcekitd_api_keys { fatalError() }
   var requests: sourcekitd_api_requests { fatalError() }

Tests/SourceKitDTests/SourceKitDTests.swift

@@ -28,7 +28,10 @@ import class TSCBasic.Process
 final class SourceKitDTests: XCTestCase {
   func testMultipleNotificationHandlers() async throws {
     let sourcekitdPath = await ToolchainRegistry.forTesting.default!.sourcekitd!
-    let sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitdPath)
+    let sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(
+      dylibPath: sourcekitdPath,
+      testHooks: SourceKitDTestHooks()
+    )
     let keys = sourcekitd.keys
     let path = DocumentURI(for: .swift).pseudoPath
 

Tests/SourceKitLSPTests/PullDiagnosticsTests.swift

@@ -310,18 +310,28 @@ final class PullDiagnosticsTests: XCTestCase {
   }
 
   func testDontReturnEmptyDiagnosticsIfDiagnosticRequestIsCancelled() async throws {
+    let diagnosticSourcekitdRequestDidStart = self.expectation(description: "diagnostic sourcekitd request did start")
     let diagnosticRequestCancelled = self.expectation(description: "diagnostic request cancelled")
     var serverOptions = SourceKitLSPServer.Options.testDefault
-    serverOptions.indexTestHooks.preparationTaskDidStart = { _ in
-      await self.fulfillment(of: [diagnosticRequestCancelled], timeout: defaultTimeout)
+    serverOptions.sourcekitdTestHooks.sourcekitdRequestDidStart = { request in
+      guard request.description.contains("source.request.diagnostics") else {
+        return
+      }
+      diagnosticSourcekitdRequestDidStart.fulfill()
+      self.wait(for: [diagnosticRequestCancelled], timeout: defaultTimeout)
+      // Poll until the `CancelRequestNotification` has been propagated to the request handling.
+      for _ in 0..<Int(defaultTimeout * 100) {
+        if Task.isCancelled {
+          break
+        }
+        usleep(10_000)
+      }
     }
     let project = try await SwiftPMTestProject(
       files: [
         "Lib.swift": "let x: String = 1"
       ],
-      serverOptions: serverOptions,
-      enableBackgroundIndexing: true,
-      pollIndex: false
+      serverOptions: serverOptions
     )
     let (uri, _) = try project.openDocument("Lib.swift")
 
@@ -332,6 +342,7 @@ final class PullDiagnosticsTests: XCTestCase {
       XCTAssertEqual(result, .failure(ResponseError.cancelled))
       diagnosticResponseReceived.fulfill()
     }
+    try await fulfillmentOfOrThrow([diagnosticSourcekitdRequestDidStart])
     project.testClient.send(CancelRequestNotification(id: requestID))
     diagnosticRequestCancelled.fulfill()
     try await fulfillmentOfOrThrow([diagnosticResponseReceived])