Wait for request to be sent to sourcekitd before cancelling it in SwiftSourceKitPluginTests.testCancellation

Otherwise, the cancellation could get handled before the request actually gets sent to sourcekitd, which is not what we want to test here. There appears to be a second underlying issue that causes unexpected results when we cancel the fast completion request after sending the slow completion request. I’ll look at that in a follow-up PR.

Author: ahoppen

Sources/SourceKitD/DynamicallyLoadedSourceKitD.swift

@@ -108,6 +108,9 @@ package actor DynamicallyLoadedSourceKitD: SourceKitD {
   /// List of notification handlers that will be called for each notification.
   private var notificationHandlers: [WeakSKDNotificationHandler] = []
 
+  /// List of hooks that should be executed for every request sent to sourcekitd.
+  private var requestHandlingHooks: [UUID: (SKDRequestDictionary) -> Void] = [:]
+
   package static func getOrCreate(
     dylibPath: URL,
     pluginPaths: PluginPaths?
@@ -204,6 +207,29 @@ package actor DynamicallyLoadedSourceKitD: SourceKitD {
     notificationHandlers.removeAll(where: { $0.value == nil || $0.value === handler })
   }
 
+  /// Execute `body` and invoke `hook` for every sourcekitd request that is sent during the execution time of `body`.
+  ///
+  /// Note that `hook` will not only be executed for requests sent *by* body but this may also include sourcekitd
+  /// requests that were sent by other clients of the same `DynamicallyLoadedSourceKitD` instance that just happen to
+  /// send a request during that time.
+  ///
+  /// This is intended for testing only.
+  package func withRequestHandlingHook(
+    body: () async throws -> Void,
+    hook: @escaping (SKDRequestDictionary) -> Void
+  ) async rethrows {
+    let id = UUID()
+    requestHandlingHooks[id] = hook
+    defer { requestHandlingHooks[id] = nil }
+    try await body()
+  }
+
+  package func didSend(request: SKDRequestDictionary) {
+    for hook in requestHandlingHooks.values {
+      hook(request)
+    }
+  }
+
   package nonisolated func log(request: SKDRequestDictionary) {
     logger.info(
       """

Sources/SourceKitD/SourceKitD.swift

@@ -83,6 +83,10 @@ package protocol SourceKitD: AnyObject, Sendable {
   /// Removes a previously registered notification handler.
   func removeNotificationHandler(_ handler: SKDNotificationHandler) async
 
+  /// A function that gets called after the request has been sent to sourcekitd but but does not wait for results to be
+  /// received. This can be used by clients to implement hooks that should be executed for every sourcekitd request.
+  func didSend(request: SKDRequestDictionary) async
+
   /// Log the given request.
   ///
   /// This log call is issued during normal operation. It is acceptable for the logger to truncate the log message
@@ -146,6 +150,9 @@ extension SourceKitD {
         self.api.send_request(request.dict, &handle) { response in
           continuation.resume(returning: SKDResponse(response!, sourcekitd: self))
         }
+        Task {
+          await self.didSend(request: request)
+        }
         if let handle {
           return SourceKitDRequestHandle(handle: handle)
         }

Tests/SourceKitDTests/SourceKitDRegistryTests.swift

@@ -74,6 +74,7 @@ final class FakeSourceKitD: SourceKitD {
   var values: sourcekitd_api_values { fatalError() }
   func addNotificationHandler(_ handler: SKDNotificationHandler) { fatalError() }
   func removeNotificationHandler(_ handler: SKDNotificationHandler) { fatalError() }
+  func didSend(request: SKDRequestDictionary) {}
   private init() {
     token = nextToken.fetchAndIncrement()
   }

Tests/SwiftSourceKitPluginTests/SwiftSourceKitPluginTests.swift

@@ -311,22 +311,34 @@ final class SwiftSourceKitPluginTests: XCTestCase {
       compilerArguments: [path]
     )
 
+    let slowCompletionRequestSent = self.expectation(description: "slow completion result sent")
     let slowCompletionResultReceived = self.expectation(description: "slow completion")
-    let slowCompletionTask = Task {
-      do {
-        _ = try await sourcekitd.completeOpen(
-          path: path,
-          position: positions["1️⃣"],
-          filter: ""
-        )
-        XCTFail("Expected completion to be cancelled")
-      } catch {
-        XCTAssert(error is CancellationError, "Expected completion to be cancelled, failed with \(error)")
+    let dynamicallyLoadedSourceKitd = try XCTUnwrap(sourcekitd as? DynamicallyLoadedSourceKitD)
+    try await dynamicallyLoadedSourceKitd.withRequestHandlingHook {
+      let slowCompletionTask = Task {
+        do {
+          _ = try await sourcekitd.completeOpen(
+            path: path,
+            position: positions["1️⃣"],
+            filter: ""
+          )
+          XCTFail("Expected completion to be cancelled")
+        } catch {
+          XCTAssert(error is CancellationError, "Expected completion to be cancelled, failed with \(error)")
+        }
+        slowCompletionResultReceived.fulfill()
       }
-      slowCompletionResultReceived.fulfill()
+      // Wait for the slow completion request to actually be sent to sourcekitd. Otherwise, we might hit a cancellation
+      // check somewhere during request sending and we aren't actually sending the completion request to sourcekitd.
+      try await fulfillmentOfOrThrow([slowCompletionRequestSent])
+
+      slowCompletionTask.cancel()
+      try await fulfillmentOfOrThrow([slowCompletionResultReceived], timeout: 30)
+    } hook: { request in
+      // Check that we aren't matching against a request sent by something else that has handle to the same sourcekitd.
+      XCTAssert(request.description.contains(path), "Received unexpected request: \(request)")
+      slowCompletionRequestSent.fulfill()
     }
-    slowCompletionTask.cancel()
-    try await fulfillmentOfOrThrow([slowCompletionResultReceived], timeout: 30)
 
     let fastCompletionStarted = Date()
     let result = try await sourcekitd.completeOpen(