Factor withCancellableCheckedThrowingContinuation into a separate function

Author: ahoppen

Sources/LanguageServerProtocol/Connection.swift

@@ -156,33 +156,12 @@ extension Connection {
   ///   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)
-
-    @Sendable
-    func sendCancelNotification() {
-      /// 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 requestID = requestIDWrapper.takeValue() {
-        self.send(CancelRequestNotification(id: requestID))
+    return try await withCancellableCheckedThrowingContinuation { continuation in
+      return self.send(request) { result in
+        continuation.resume(with: result)
       }
+    } cancel: { requestID in
+      self.send(CancelRequestNotification(id: requestID))
     }
-
-    return try await withTaskCancellationHandler(operation: {
-      try Task.checkCancellation()
-      return try await withCheckedThrowingContinuation { continuation in
-        let requestID = self.send(request) { result in
-          continuation.resume(with: result)
-        }
-        requestIDWrapper.value = requestID
-
-        // 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 {
-          sendCancelNotification()
-        }
-      }
-    }, onCancel: sendCancelNotification)
   }
 }

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,13 +1,13 @@
 
 add_library(SKSupport STATIC
+  AsyncUtils.swift
   BuildConfiguration.swift
   ByteString.swift
   dlopen.swift
   FileSystem.swift
   LineTable.swift
   Random.swift
   Result.swift
-  Task+ValuePropagatingCancellation.swift
   ThreadSafeBox.swift
 )
 set_target_properties(SKSupport PROPERTIES

Sources/SKSupport/Task+ValuePropagatingCancellation.swift

@@ -83,45 +83,26 @@ extension SourceKitD {
   public func send(_ req: SKDRequestDictionary) async throws -> SKDResponseDictionary {
     logRequest(req)
 
-    let handleWrapper = ThreadSafeBox<sourcekitd_request_handle_t?>(initialValue: nil)
-
-    @Sendable
-    func cancelSourcekitdRequest() {
-      /// 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() {
-        api.cancel_request(handle)
-      }
-    }
-
-    return try await withTaskCancellationHandler(operation: {
-      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)
+    return try await withCancellableCheckedThrowingContinuation { continuation in
+      var handle: sourcekitd_request_handle_t?
+      api.send_request(req.dict, &handle) { [weak self] _resp in
+        guard let self = self else { return }
 
-          logResponse(resp)
+        let resp = SKDResponse(_resp, sourcekitd: self)
 
-          guard let dict = resp.value else {
-            continuation.resume(throwing: resp.error!)
-            return
-          }
+        logResponse(resp)
 
-          continuation.resume(returning: dict)
-        }
-        handleWrapper.value = handle
-        // 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 request handle.
-        if Task.isCancelled {
-          cancelSourcekitdRequest()
+        guard let dict = resp.value else {
+          continuation.resume(throwing: resp.error!)
+          return
         }
+
+        continuation.resume(returning: dict)
       }
-    }, onCancel: cancelSourcekitdRequest)
+      return handle
+    } cancel: { handle in
+      api.cancel_request(handle)
+    }
   }
 }