Implement request cancellation

When receiving a `CancellationNotification`, we cancel the task that handles the request with that ID.

This will cause `cancel_notification` to be sent to sourcekitd or a `CancellationNotification` to be sent to `clangd`, which ultimately cancels the request.

rdar://117492860

Author: ahoppen

Package.swift

@@ -103,6 +103,7 @@ let package = Package(
       name: "LanguageServerProtocol",
       dependencies: [
         "LSPLogging",
+        "SKSupport",
         .product(name: "SwiftToolsSupport-auto", package: "swift-tools-support-core"),
       ],
       exclude: ["CMakeLists.txt"]

Sources/LSPTestSupport/Assertions.swift

@@ -36,7 +36,7 @@ public func assertNoThrow<T>(
   }
 }
 
-/// Same as `XCTAssertThrows` but executes the trailing closure.
+/// Same as `XCTAssertThrows` but allows the expression to be async
 public func assertThrowsError<T>(
   _ expression: @autoclosure () async throws -> T,
   _ message: @autoclosure () -> String = "",

Sources/LanguageServerProtocol/AsyncQueue.swift

@@ -88,7 +88,7 @@ public final class AsyncQueue<TaskMetadata: DependencyTracker> {
     let throwingTask = asyncThrowing(priority: priority, metadata: metadata, operation: operation)
     return Task {
       do {
-        return try await throwingTask.value
+        return try await throwingTask.valuePropagatingCancellation
       } catch {
         // We know this can never happen because `operation` does not throw.
         preconditionFailure("Executing a task threw an error even though the operation did not throw")
@@ -141,7 +141,7 @@ public final class AsyncQueue<TaskMetadata: DependencyTracker> {
 
 /// Convenience overloads for serial queues.
 extension AsyncQueue where TaskMetadata == Serial {
-  /// Same as ``async(priority:operation:)`` but specialized for serial queues 
+  /// Same as ``async(priority:operation:)`` but specialized for serial queues
   /// that don't specify any metadata.
   @discardableResult
   public func async<Success: Sendable>(
@@ -151,7 +151,7 @@ extension AsyncQueue where TaskMetadata == Serial {
     return self.async(priority: priority, metadata: Serial(), operation: operation)
   }
 
-  /// Same as ``asyncThrowing(priority:metadata:operation:)`` but specialized 
+  /// Same as ``asyncThrowing(priority:metadata:operation:)`` but specialized
   /// for serial queues that don't specify any metadata.
   public func asyncThrowing<Success: Sendable>(
     priority: TaskPriority? = nil,

Sources/LanguageServerProtocol/CMakeLists.txt

@@ -138,10 +138,8 @@ add_library(LanguageServerProtocol STATIC
 set_target_properties(LanguageServerProtocol PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})
 target_link_libraries(LanguageServerProtocol PUBLIC
+  LSPLogging
+  SKSupport
   TSCBasic
   $<$<NOT:$<PLATFORM_ID:Darwin>>:swiftDispatch>
   $<$<NOT:$<PLATFORM_ID:Darwin>>:Foundation>)
-
-target_link_libraries(LanguageServerProtocol PUBLIC
-  LSPLogging
-)

Sources/LanguageServerProtocol/Connection.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import Dispatch
+import SKSupport
 
 /// An abstract connection, allow messages to be sent to a (potentially remote) `MessageHandler`.
 public protocol Connection: AnyObject {
@@ -144,3 +145,30 @@ extension LocalConnection: Connection {
     return id
   }
 }
+
+extension Connection {
+  /// Send the given request to the connection and await its result.
+  ///
+  /// This method automatically sends a `CancelRequestNotification` to the
+  /// connection if the task it is executing in is being cancelled.
+  ///
+  /// - Warning: Because this message is `async`, it does not provide any ordering
+  ///   guarantees. If you need to gurantee that messages are sent in-order
+  ///   use the version with a completion handler.
+  public func send<R: RequestType>(_ request: R) async throws -> R.Response {
+    let requestIDWrapper = ThreadSafeBox<RequestID?>(initialValue: nil)
+    try Task.checkCancellation()
+    return try await withTaskCancellationHandler {
+      try await withCheckedThrowingContinuation { continuation in
+        let requestID = self.send(request) { result in
+          continuation.resume(with: result)
+        }
+        requestIDWrapper.value = requestID
+      }
+    } onCancel: {
+      if let requestID = requestIDWrapper.value {
+        self.send(CancelRequestNotification(id: requestID))
+      }
+    }
+  }
+}

Sources/SKSupport/CMakeLists.txt

@@ -2,11 +2,14 @@
 add_library(SKSupport STATIC
   BuildConfiguration.swift
   ByteString.swift
+  dlopen.swift
   FileSystem.swift
   LineTable.swift
   Random.swift
   Result.swift
-  dlopen.swift)
+  Task+ValuePropagatingCancellation.swift
+  ThreadSafeBox.swift
+)
 set_target_properties(SKSupport PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})
 target_link_libraries(SKSupport PRIVATE

Sources/SKSupport/Task+ValuePropagatingCancellation.swift

@@ -0,0 +1,26 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+public extension Task {
+  /// Awaits the value of the result.
+  ///
+  /// If the current task is cancelled, this will cancel the subtask as well.
+  var valuePropagatingCancellation: Success {
+    get async throws {
+      try await withTaskCancellationHandler {
+        return try await self.value
+      } onCancel: {
+        self.cancel()
+      }
+    }
+  }
+}

Sources/SKSupport/ThreadSafeBox.swift

@@ -0,0 +1,47 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Foundation
+
+extension NSLock {
+  /// NOTE: Keep in sync with SwiftPM's 'Sources/Basics/NSLock+Extensions.swift'
+  fileprivate func withLock<T>(_ body: () throws -> T) rethrows -> T {
+    lock()
+    defer { unlock() }
+    return try body()
+  }
+}
+
+/// A thread safe container that contains a value of type `T`.
+public class ThreadSafeBox<T> {
+  /// Lock guarding `_value`.
+  private let lock = NSLock()
+
+  private var _value: T
+
+  public var value: T {
+    get {
+      return lock.withLock {
+        return _value
+      }
+    }
+    set {
+      lock.withLock {
+        _value = newValue
+      }
+    }
+  }
+
+  public init(initialValue: T) {
+    _value = initialValue
+  }
+}

Sources/SourceKitD/SourceKitD.swift

@@ -83,26 +83,30 @@ extension SourceKitD {
   public func send(_ req: SKDRequestDictionary) async throws -> SKDResponseDictionary {
     logRequest(req)
 
-    let sourcekitdResponse: SKDResponse = await withCheckedContinuation { continuation in
-      var handle: sourcekitd_request_handle_t? = nil
-
-      api.send_request(req.dict, &handle) { _resp in
-        continuation.resume(returning: SKDResponse(_resp, sourcekitd: self))
+    let handleWrapper = ThreadSafeBox<sourcekitd_request_handle_t?>(initialValue: nil)
+    let sourcekitdResponse: SKDResponse = try await withTaskCancellationHandler {
+      try Task.checkCancellation()
+      return await withCheckedContinuation { continuation in
+        var handle: sourcekitd_request_handle_t? = nil
+
+        api.send_request(req.dict, &handle) { _resp in
+          continuation.resume(returning: SKDResponse(_resp, sourcekitd: self))
+        }
+        handleWrapper.value = handle
+      }
+    } onCancel: {
+      if let handle = handleWrapper.value {
+        api.cancel_request(handle)
       }
     }
+
     logResponse(sourcekitdResponse)
 
     guard let dict = sourcekitdResponse.value else {
       throw sourcekitdResponse.error!
     }
 
     return dict
-
-    // FIXME: (async) Cancellation
-  }
-
-  public func cancel(_ handle: sourcekitd_request_handle_t) {
-    api.cancel_request(handle)
   }
 }
 

Sources/SourceKitLSP/Clang/ClangLanguageServer.swift

@@ -327,17 +327,7 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
   ///
   /// The response of the request is  returned asynchronously as the return value.
   func forwardRequestToClangd<R: RequestType>(_ request: R) async throws -> R.Response {
-    try await withCheckedThrowingContinuation { continuation in
-      _ = clangd.send(request) { result in
-        switch result {
-        case .success(let response):
-          continuation.resume(returning: response)
-        case .failure(let error):
-          continuation.resume(throwing: error)
-        }
-      }
-    }
-    // FIXME: (async) Cancellation
+    return try await clangd.send(request)
   }
 
   func _crash() {

Sources/SourceKitLSP/SourceKitServer.swift

@@ -262,6 +262,12 @@ public actor SourceKitServer {
   /// The actual semantic handling of the message happens off this queue.
   private let messageHandlingQueue = AsyncQueue<TaskMetadata>()
 
+  /// The queue on which we start and stop keeping track of cancellation.
+  ///
+  /// Having a queue for this ensures that we started keeping track of a
+  /// request's task before handling any cancellation request for it.
+  private let cancellationMessageHandlingQueue = AsyncQueue<Serial>()
+
   /// The connection to the editor.
   public let client: Connection
 
@@ -305,6 +311,17 @@ public actor SourceKitServer {
     }
   }
 
+  /// The requests that we are currently handling.
+  ///
+  /// Used to cancel the tasks if the client requests cancellation.
+  private var inProgressRequests: [RequestID: Task<(), Never>] = [:]
+
+  /// - Note: Needed so we can set an in-progress request from a different
+  ///   isolation context.
+  private func setInProgressRequest(for id: RequestID, task: Task<(), Never>?) {
+    self.inProgressRequests[id] = task
+  }
+
   let fs: FileSystem
 
   var onExit: () -> Void
@@ -400,12 +417,7 @@ public actor SourceKitServer {
 
   /// Send the given request to the editor.
   public func sendRequestToClient<R: RequestType>(_ request: R) async throws -> R.Response {
-    try await withCheckedThrowingContinuation { continuation in
-      _ = client.send(request) { result in
-        continuation.resume(with: result)
-      }
-      // FIXME: (async) Handle cancellation
-    }
+    return try await client.send(request)
   }
 
   func toolchain(for uri: DocumentURI, _ language: Language) -> Toolchain? {
@@ -608,6 +620,13 @@ private func getNextNotificationIDForLogging() -> Int {
 
 extension SourceKitServer: MessageHandler {
   public nonisolated func handle(_ params: some NotificationType, from clientID: ObjectIdentifier) {
+    if let params = params as? CancelRequestNotification {
+      // Request cancellation needs to be able to overtake any other message we
+      // are currently handling. Ordering is not important here. We thus don't
+      // need to execute it on `messageHandlingQueue`.
+      self.cancelRequest(params)
+    }
+
     messageHandlingQueue.async(metadata: TaskMetadata(params)) {
       let notificationID = getNextNotificationIDForLogging()
 
@@ -630,8 +649,6 @@ extension SourceKitServer: MessageHandler {
     switch notification.params {
     case let notification as InitializedNotification:
       self.clientInitialized(notification)
-    case let notification as CancelRequestNotification:
-      self.cancelRequest(notification)
     case let notification as ExitNotification:
       await self.exit(notification)
     case let notification as DidOpenTextDocumentNotification:
@@ -660,10 +677,21 @@ extension SourceKitServer: MessageHandler {
     from clientID: ObjectIdentifier,
     reply: @escaping (LSPResult<R.Response>) -> Void
   ) {
-    messageHandlingQueue.async(metadata: TaskMetadata(params)) {
+    let task = messageHandlingQueue.async(metadata: TaskMetadata(params)) {
       await withLoggingScope("request-\(id)") {
         await self.handleImpl(params, id: id, from: clientID, reply: reply)
       }
+      // We have handled the request and can't cancel it anymore.
+      // Stop keeping track of it to free the memory.
+      self.cancellationMessageHandlingQueue.async(priority: .background) {
+        await self.setInProgressRequest(for: id, task: nil)
+      }
+    }
+    // Keep track of the ID -> Task management with low priority. Once we cancel
+    // a request, the cancellation task runs with a high priority and depends on
+    // this task, which will elevate this task's priority.
+    cancellationMessageHandlingQueue.async(priority: .background) {
+      await self.setInProgressRequest(for: id, task: task)
     }
   }
 
@@ -1077,8 +1105,18 @@ extension SourceKitServer {
     // Nothing to do.
   }
 
-  func cancelRequest(_ notification: CancelRequestNotification) {
-    // TODO: Implement cancellation
+  nonisolated func cancelRequest(_ notification: CancelRequestNotification) {
+    // Since the request is very cheap to execute and stops other requests
+    // from performing more work, we execute it with a high priority.
+    cancellationMessageHandlingQueue.async(priority: .high) {
+      guard let task = await self.inProgressRequests[notification.id] else {
+        logger.error(
+          "Cannot cancel request \(notification.id, privacy: .public) because it hasn't been scheduled for execution yet"
+        )
+        return
+      }
+      task.cancel()
+    }
   }
 
   /// The server is about to exit, and the server should flush any buffered state.

Sources/SourceKitLSP/Swift/CodeCompletionSession.swift

@@ -138,8 +138,7 @@ class CodeCompletionSession {
       return try await session.open(filterText: filterText, position: cursorPosition, in: snapshot, options: options)
     }
 
-    // FIXME: (async) Use valuePropagatingCancellation once we support cancellation
-    return try await task.value
+    return try await task.valuePropagatingCancellation
   }
 
   // MARK: - Implementation

Sources/SourceKitLSP/Swift/SourceKitD+ResponseError.swift

@@ -17,7 +17,7 @@ extension ResponseError {
   public init(_ value: SKDError) {
     switch value {
     case .requestCancelled:
-      self = .serverCancelled
+      self = .cancelled
     case .requestFailed(let desc):
       self = .unknown("sourcekitd request failed: \(desc)")
     case .requestInvalid(let desc):