Wait for cancellation to propagate in testDontReturnEmptyDiagnosticsIfDiagnosticRequestIsCancelled

I saw a few non-deterministic test failures. I think the issue was that handling of the `CancelRequestNotification` is done asynchronously, which left a short window in which the sourcekitd diagnostics request could run and return results instead of being cancelled. Wait for the diagnostic request to actually be cancelled before running it for real.

While doing this, also introduce proper sourcekitd test hooks instead of relying on the preparation test hooks, which just got run as a side effect.

Author: ahoppen

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])