Avoid a potential race condition in which a sourcekitd/clangd request wouldn't get cancelled

Author: ahoppen

Sources/LanguageServerProtocol/Connection.swift

@@ -157,18 +157,32 @@ extension Connection {
   ///   use the version with a completion handler.
   public func send<R: RequestType>(_ request: R) async throws -> R.Response {
     let requestIDWrapper = ThreadSafeBox<RequestID?>(initialValue: nil)
-    return try await withTaskCancellationHandler {
+
+    @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 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: {
-      if let requestID = requestIDWrapper.value {
-        self.send(CancelRequestNotification(id: requestID))
-      }
-    }
+    }, onCancel: sendCancelNotification)
   }
 }

Sources/SKSupport/ThreadSafeBox.swift

@@ -44,4 +44,14 @@ public class ThreadSafeBox<T> {
   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,22 +83,35 @@ extension SourceKitD {
   public func send(_ req: SKDRequestDictionary) async throws -> SKDResponseDictionary {
     logRequest(req)
 
+    @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)
+      }
+    }
+
     let handleWrapper = ThreadSafeBox<sourcekitd_request_handle_t?>(initialValue: nil)
-    let sourcekitdResponse: SKDResponse = try await withTaskCancellationHandler {
+    let sourcekitdResponse: SKDResponse = try await withTaskCancellationHandler(operation: {
       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
+        // 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()
+        }
       }
-    } onCancel: {
-      if let handle = handleWrapper.value {
-        api.cancel_request(handle)
-      }
-    }
+    }, onCancel: cancelSourcekitdRequest)
 
     logResponse(sourcekitdResponse)