Implement request cancellation

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/LSPTestSupport/TestJSONRPCConnection.swift

@@ -129,9 +129,7 @@ public final class TestMessageHandler: MessageHandler {
     from clientID: ObjectIdentifier,
     reply: @escaping (LSPResult<R.Response>) -> Void
   ) {
-    let cancellationToken = CancellationToken()
-
-    let request = Request(params, id: id, clientID: clientID, cancellation: cancellationToken, reply: reply)
+    let request = Request(params, id: id, clientID: clientID, reply: reply)
 
     guard !oneShotRequestHandlers.isEmpty else {
       fatalError("unexpected request \(request)")
@@ -179,14 +177,11 @@ public final class TestServer: MessageHandler {
     from clientID: ObjectIdentifier,
     reply: @escaping (LSPResult<R.Response>) -> Void
   ) {
-    let cancellationToken = CancellationToken()
-
     if let params = params as? EchoRequest {
       let req = Request(
         params,
         id: id,
         clientID: clientID,
-        cancellation: cancellationToken,
         reply: { result in
           reply(result.map({ $0 as! R.Response }))
         }
@@ -197,7 +192,6 @@ public final class TestServer: MessageHandler {
         params,
         id: id,
         clientID: clientID,
-        cancellation: cancellationToken,
         reply: { result in
           reply(result.map({ $0 as! R.Response }))
         }

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

@@ -1,6 +1,5 @@
 add_library(LanguageServerProtocol STATIC
   AsyncQueue.swift
-  Cancellation.swift
   Connection.swift
   CustomCodable.swift
   Error.swift
@@ -138,10 +137,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,23 @@ 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 {
+    return try await withCancellableCheckedThrowingContinuation { continuation in
+      return self.send(request) { result in
+        continuation.resume(with: result)
+      }
+    } cancel: { requestID in
+      self.send(CancelRequestNotification(id: requestID))
+    }
+  }
+}

Sources/LanguageServerProtocol/Request.swift

@@ -38,20 +38,15 @@ public final class Request<R: RequestType> {
     }
   }
 
-  /// The request's cancellation state.
-  public let cancellationToken: CancellationToken
-
   public init(
     _ request: Params,
     id: RequestID,
     clientID: ObjectIdentifier,
-    cancellation: CancellationToken,
     reply: @escaping (LSPResult<Response>) -> Void
   ) {
     self.id = id
     self.clientID = clientID
     self.params = request
-    self.cancellationToken = cancellation
     self.replyBlock = reply
   }
 
@@ -71,9 +66,6 @@ public final class Request<R: RequestType> {
   public func reply(_ result: Response) {
     reply(.success(result))
   }
-
-  /// Whether the result has been cancelled.
-  public var isCancelled: Bool { return cancellationToken.isCancelled }
 }
 
 /// A request object, wrapping the parameters of a `NotificationType`.

Sources/SKSupport/AsyncUtils.swift

@@ -0,0 +1,64 @@
+//===----------------------------------------------------------------------===//
+//
+// 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()
+      }
+    }
+  }
+}
+
+/// Allows the execution of a cancellable operation that returns the results
+/// via a completion handler.
+///
+/// `operation` must invoke the continuation's `resume` method exactly once.
+///
+/// If the task executing `withCancellableCheckedThrowingContinuation` gets
+/// cancelled, `cancel` is invoked with the handle that `operation` provided.
+public func withCancellableCheckedThrowingContinuation<Handle, Result>(
+  _ operation: (_ continuation: CheckedContinuation<Result, any Error>) -> Handle,
+  cancel: (Handle) -> Void
+) async throws -> Result {
+  let handleWrapper = ThreadSafeBox<Handle?>(initialValue: nil)
+
+  @Sendable
+  func callCancel() {
+    /// Take the request ID out of the box. This ensures that we only send the
+    /// cancel notification once in case the `Task.isCancelled` and the
+    /// `onCancel` check race.
+    if let handle = handleWrapper.takeValue() {
+      cancel(handle)
+    }
+  }
+
+  return try await withTaskCancellationHandler(operation: {
+    try Task.checkCancellation()
+    return try await withCheckedThrowingContinuation { continuation in
+      handleWrapper.value = operation(continuation)
+
+      // Check if the task was cancelled. This ensures we send a
+      // CancelNotification even if the task gets cancelled after we register
+      // the cancellation handler but before we set the `requestID`.
+      if Task.isCancelled {
+        callCancel()
+      }
+    }
+  }, onCancel: callCancel)
+}

Sources/SKSupport/CMakeLists.txt

@@ -1,12 +1,15 @@
 
 add_library(SKSupport STATIC
+  AsyncUtils.swift
   BuildConfiguration.swift
   ByteString.swift
+  dlopen.swift
   FileSystem.swift
   LineTable.swift
   Random.swift
   Result.swift
-  dlopen.swift)
+  ThreadSafeBox.swift
+)
 set_target_properties(SKSupport PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})
 target_link_libraries(SKSupport PRIVATE

Sources/SKSupport/ThreadSafeBox.swift

@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+  }
+
+  /// If the value in the box is an optional, return it and reset it to `nil`
+  /// in an atomic operation.
+  public func takeValue<U>() -> T where U? == T {
+    lock.withLock {
+      guard let value = self._value else { return nil }
+      self._value = nil
+      return value
+    }
+  }
+}

Sources/SourceKitD/SourceKitD.swift

@@ -83,26 +83,23 @@ extension SourceKitD {
   public func send(_ req: SKDRequestDictionary) async throws -> SKDResponseDictionary {
     logRequest(req)
 
-    let sourcekitdResponse: SKDResponse = await withCheckedContinuation { continuation in
+    let sourcekitdResponse: SKDResponse = try await withCancellableCheckedThrowingContinuation { continuation in
       var handle: sourcekitd_request_handle_t? = nil
-
       api.send_request(req.dict, &handle) { _resp in
         continuation.resume(returning: SKDResponse(_resp, sourcekitd: self))
       }
+      return handle
+    } cancel: { handle in
+      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

@@ -298,7 +298,6 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
         params,
         id: id,
         clientID: clientID,
-        cancellation: CancellationToken(),
         reply: { result in
           reply(result)
         }
@@ -327,17 +326,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() {