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.

Author: ahoppen

Package.swift

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

Sources/LanguageServerProtocol/AsyncQueue.swift

@@ -87,7 +87,7 @@ public final class AsyncQueue {
     let throwingTask = asyncThrowing(priority: priority, barrier: isBarrier, 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")

Sources/LanguageServerProtocol/CMakeLists.txt

@@ -141,3 +141,7 @@ target_link_libraries(LanguageServerProtocol PUBLIC
   TSCBasic
   $<$<NOT:$<PLATFORM_ID:Darwin>>:swiftDispatch>
   $<$<NOT:$<PLATFORM_ID:Darwin>>:Foundation>)
+
+target_link_libraries(LanguageServerProtocol PRIVATE
+  SKSupport
+)

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 {
@@ -136,3 +137,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/SKTestSupport/INPUTS/Cancellation/main.swift

@@ -0,0 +1,13 @@
+class Foo {
+  func slow(x: Invalid1, y: Invalid2) {
+    x / y / x / y / x / y / x / y . /*slowLoc*/
+  }
+
+  struct Foo {
+    let fooMember: String
+  }
+
+  func fast(a: Foo) {
+    a . /*fastLoc*/
+  }
+}

Sources/SKTestSupport/INPUTS/Cancellation/project.json

@@ -0,0 +1 @@
+{"sources": ["main.swift"]}

Sources/SourceKitD/SourceKitD.swift

@@ -84,31 +84,34 @@ extension SourceKitD {
   public func send(_ req: SKDRequestDictionary) async throws -> SKDResponseDictionary {
     logRequest(req)
 
-    return try await withCheckedThrowingContinuation { continuation in
-      var handle: sourcekitd_request_handle_t? = nil
+    let handleWrapper = ThreadSafeBox<sourcekitd_request_handle_t?>(initialValue: nil)
 
-      api.send_request(req.dict, &handle) { [weak self] _resp in
-        guard let self = self else { return }
+    return try await withTaskCancellationHandler {
+      try Task.checkCancellation()
+      return try await withCheckedThrowingContinuation { continuation in
+        var handle: sourcekitd_request_handle_t?
+        api.send_request(req.dict, &handle) { [weak self] _resp in
+          guard let self = self else { return }
 
-        let resp = SKDResponse(_resp, sourcekitd: self)
+          let resp = SKDResponse(_resp, sourcekitd: self)
 
-        logResponse(resp)
+          logResponse(resp)
 
-        guard let dict = resp.value else {
-          continuation.resume(throwing: resp.error!)
-          return
-        }
+          guard let dict = resp.value else {
+            continuation.resume(throwing: resp.error!)
+            return
+          }
 
-        continuation.resume(returning: dict)
+          continuation.resume(returning: dict)
+        }
+        handleWrapper.value = handle
+      }
+    } onCancel: {
+      if let handle = handleWrapper.value {
+        api.cancel_request(handle)
       }
-
-      // FIXME: (async) Cancellation
     }
   }
-
-  public func cancel(_ handle: sourcekitd_request_handle_t) {
-    api.cancel_request(handle)
-  }
 }
 
 private func logRequest(_ request: SKDRequestDictionary) {

Sources/SourceKitLSP/Clang/ClangLanguageServer.swift

@@ -312,17 +312,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

@@ -144,6 +144,12 @@ public actor SourceKitServer {
   /// The actual semantic handling of the message happens off this queue.
   private let messageHandlingQueue = AsyncQueue(.concurrent)
 
+  /// 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
 
@@ -184,6 +190,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
@@ -275,12 +292,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? {
@@ -450,28 +462,32 @@ public actor SourceKitServer {
 
 extension SourceKitServer: MessageHandler {
   public nonisolated func handle(_ params: some NotificationType, from clientID: ObjectIdentifier) {
-    // All of the notifications sourcekit-lsp currently handles might modify the
+    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)
+    }
+
+    // All of the other notifications sourcekit-lsp currently handles might modify the
     // global state (eg. whether a document is open or its contents) in a way
     // that changes the results of requsts before and after.
     // We thus need to ensure that we handle the notifications in order, so they
     // need to be dispatch barriers.
     //
     // Technically, we could optimize this further by having an `AsyncQueue` for
     // each file, because edits on one file should not block requests on another
-    // file from executing but, at least in Swift, this would get us any real 
+    // file from executing but, at least in Swift, this would get us any real
     // benefits at the moment because sourcekitd only has a single, global queue,
     // instead of a queue per file.
     // Additionally, usually you are editing one file in a source editor, which
     // means that concurrent requests to multiple files tend to be rare.
     messageHandlingQueue.async(barrier: true) {
-      let notification = Notification(params, clientID: clientID)
-      await self._logNotification(notification)
+      await self._logNotification(Notification(params, clientID: clientID))
 
-      switch notification.params {
+      switch params {
       case let notification as InitializedNotification:
         await self.clientInitialized(notification)
-      case let notification as CancelRequestNotification:
-        await self.cancelRequest(notification)
       case let notification as ExitNotification:
         await self.exit(notification)
       case let notification as DidOpenTextDocumentNotification:
@@ -503,7 +519,7 @@ extension SourceKitServer: MessageHandler {
     //    completion session but it only makes sense for the client to request
     //    more results for this completion session after it has received the
     //    initial results.
-    messageHandlingQueue.async(barrier: false) {
+    let task = messageHandlingQueue.async(barrier: false) {
       let cancellationToken = CancellationToken()
 
       let request = Request(params, id: id, clientID: clientID, cancellation: cancellationToken, reply: { [weak self] result in
@@ -581,6 +597,17 @@ extension SourceKitServer: MessageHandler {
       default:
         reply(.failure(ResponseError.methodNotFound(R.method)))
       }
+      // 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)
     }
   }
 
@@ -910,8 +937,16 @@ 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 {
+        log("Cannot cancel task because it hasn't been scheduled for execution yet", level: .error)
+        return
+      }
+      task.cancel()
+    }
   }
 
   /// The server is about to exit, and the server should flush any buffered state.

Sources/SourceKitLSP/Swift/CodeCompletionSession.swift

@@ -81,7 +81,7 @@ actor CodeCompletionSession {
         return try await self.updateImpl(filterText: filterText, position: position, in: snapshot, options: options)
       }
     }
-    return try await task.value
+    return try await task.valuePropagatingCancellation
   }
 
   private func open(

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