Added helper types for dealing with continuations that can be cancelled

Author: dimitribouniol

Sources/WebAuthn/Helpers/CancellableContinuationTask.swift

@@ -0,0 +1,80 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the WebAuthn Swift open source project
+//
+// Copyright (c) 2023 the WebAuthn Swift project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of WebAuthn Swift project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+/// An internal type to assist kicking off work without primarily awaiting it, instead allowing that work to call into a continuation as needed.
+/// Use ``withCancellableFirstSuccessfulContinuation()`` instead of invoking this directly.
+actor CancellableContinuation<T: Sendable>: Sendable {
+    private var task: Task<Void, Error>?
+    private var continuation: CheckedContinuation<T, Error>?
+    private var isCancelled = false
+    
+    private func cancelMainTask() {
+        continuation?.resume(throwing: CancellationError())
+        continuation = nil
+        task?.cancel()
+        isCancelled = true
+    }
+    
+    private func isolatedResume(returning value: T) {
+        continuation?.resume(returning: value)
+        continuation = nil
+        cancelMainTask()
+    }
+    
+    nonisolated func cancel() {
+        Task { await cancelMainTask() }
+    }
+    
+    nonisolated func resume(returning value: T) {
+        Task { await isolatedResume(returning: value) }
+    }
+    
+    /// Wrap an asynchronous closure providing a continuation for when results are ready that can be called any number of times, but also allowing the closure to be cancelled at any time, including once the first successful value is provided.
+    fileprivate func wrap(_ body: @Sendable (_ continuation: CancellableContinuation) async throws -> ()) async throws -> T {
+        assert(task == nil, "A CancellableContinuationTask should only be used once.")
+        /// Register a cancellation callback that will: a) immediately cancel the continuation if we have one, b) unset it so it doesn't get called a second time, and c) cancel the main task.
+        return try await withTaskCancellationHandler {
+            try await withoutActuallyEscaping(body) { escapingBody in
+                try await withCheckedThrowingContinuation { localContinuation in
+                    /// Synchronously a) check if we've been cancelled, stopping early, b) save the contnuation, and c) assign the task, which runs immediately.
+                    /// This works since we are guaranteed to hear back from the cancellation handler either immediately, since Task.isCancelled is already set, or after task is set, since we are executing on the actor's executor.
+                    guard !Task.isCancelled else {
+                        localContinuation.resume(throwing: CancellationError())
+                        return
+                    }
+                    
+                    self.continuation = localContinuation
+                    self.task = Task { [unowned self] in
+                        /// If the continuation doesn't exist at this point, it's because we've already been cancelled. This is guaranteed to run after the task has been set and potentially cancelled since it also runs on the task executor.
+                        guard let continuation = self.continuation else { return }
+                        do {
+                            try await escapingBody(self)
+                        } catch {
+                            /// If the main body fails for any reason, pass along the error. This will be a no-op if the continuation was already resumed or cancelled.
+                            continuation.resume(throwing: error)
+                            self.continuation = nil
+                        }
+                    }
+                }
+            }
+        } onCancel: {
+            cancel()
+        }
+    }
+}
+
+/// Execute an operation providing it a continuation for when results are ready that can be called any number of times, but also allowing the operation to be cancelled at any time, including once the first successful value is provided.
+func withCancellableFirstSuccessfulContinuation<T: Sendable>(_ body: (_ continuation: CancellableContinuation<T>) async throws -> ()) async throws -> T {
+    try await CancellableContinuation().wrap { try await body($0) }
+}