[Concurrency] Correct withContinuation APIs isolation handling

rdar://125307764

abi cleanup

Author: ktoso

stdlib/public/Concurrency/CheckedContinuation.swift

@@ -281,19 +281,37 @@ extension CheckedContinuation {
 /// - SeeAlso: `withCheckedThrowingContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeThrowingContinuation(function:_:)`
-@available(SwiftStdlib 5.1, *)
-@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
 @inlinable
 @_unavailableInEmbedded
+@available(SwiftStdlib 5.1, *)
+#if !$Embedded
+@backDeployed(before: SwiftStdlib 6.0)
+#endif
 public func withCheckedContinuation<T>(
-    function: String = #function,
-    _ body: (CheckedContinuation<T, Never>) -> Void
+  isolation: isolated (any Actor)? = #isolation,
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Never>) -> Void
+) async -> T {
+  return await withUnsafeContinuation {
+    body(CheckedContinuation(continuation: $0, function: function))
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+@usableFromInline
+@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
+@_unavailableInEmbedded
+@_silgen_name("$ss23withCheckedContinuation8function_xSS_yScCyxs5NeverOGXEtYalF")
+internal func __abi_withCheckedContinuation<T>(
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Never>) -> Void
 ) async -> T {
   return await withUnsafeContinuation {
     body(CheckedContinuation(continuation: $0, function: function))
   }
 }
 
+
 /// Invokes the passed in closure with a checked continuation for the current task.
 ///
 /// The body of the closure executes synchronously on the calling task, and once it returns
@@ -322,13 +340,30 @@ public func withCheckedContinuation<T>(
 /// - SeeAlso: `withCheckedContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeThrowingContinuation(function:_:)`
-@available(SwiftStdlib 5.1, *)
-@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
 @inlinable
 @_unavailableInEmbedded
+@available(SwiftStdlib 5.1, *)
+#if !$Embedded
+@backDeployed(before: SwiftStdlib 6.0)
+#endif
 public func withCheckedThrowingContinuation<T>(
-    function: String = #function,
-    _ body: (CheckedContinuation<T, Error>) -> Void
+  isolation: isolated (any Actor)? = #isolation,
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Error>) -> Void
+) async throws -> T {
+  return try await withUnsafeThrowingContinuation {
+    body(CheckedContinuation(continuation: $0, function: function))
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+@usableFromInline
+@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
+@_unavailableInEmbedded
+@_silgen_name("$ss31withCheckedThrowingContinuation8function_xSS_yScCyxs5Error_pGXEtYaKlF")
+internal func __abi_withCheckedThrowingContinuation<T>(
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Error>) -> Void
 ) async throws -> T {
   return try await withUnsafeThrowingContinuation {
     body(CheckedContinuation(continuation: $0, function: function))

stdlib/public/Concurrency/PartialAsyncTask.swift

@@ -598,9 +598,9 @@ internal func _resumeUnsafeThrowingContinuationWithError<T>(
 /// - SeeAlso: `withCheckedContinuation(function:_:)`
 /// - SeeAlso: `withCheckedThrowingContinuation(function:_:)`
 @available(SwiftStdlib 5.1, *)
-@_unsafeInheritExecutor
 @_alwaysEmitIntoClient
 public func withUnsafeContinuation<T>(
+  isolation: isolated (any Actor)? = #isolation,
   _ fn: (UnsafeContinuation<T, Never>) -> Void
 ) async -> T {
   return await Builtin.withUnsafeContinuation {
@@ -634,9 +634,9 @@ public func withUnsafeContinuation<T>(
 /// - SeeAlso: `withCheckedContinuation(function:_:)`
 /// - SeeAlso: `withCheckedThrowingContinuation(function:_:)`
 @available(SwiftStdlib 5.1, *)
-@_unsafeInheritExecutor
 @_alwaysEmitIntoClient
 public func withUnsafeThrowingContinuation<T>(
+  isolation: isolated (any Actor)? = #isolation,
   _ fn: (UnsafeContinuation<T, Error>) -> Void
 ) async throws -> T {
   return try await Builtin.withUnsafeThrowingContinuation {

stdlib/public/Concurrency/Task.cpp

@@ -346,9 +346,9 @@ static SerialExecutorRef executorForEnqueuedJob(Job *job) {
   return SerialExecutorRef::generic();
 #else
   void *jobQueue = job->SchedulerPrivate[Job::DispatchQueueIndex];
-  if (jobQueue == DISPATCH_QUEUE_GLOBAL_EXECUTOR)
+  if (jobQueue == DISPATCH_QUEUE_GLOBAL_EXECUTOR) {
     return SerialExecutorRef::generic();
-  else
+  } else
     return SerialExecutorRef::forOrdinary(reinterpret_cast<HeapObject*>(jobQueue),
                     _swift_task_getDispatchQueueSerialExecutorWitnessTable());
 #endif

stdlib/public/Concurrency/TaskGroup.swift

@@ -593,7 +593,6 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
 
   @usableFromInline
   @available(SwiftStdlib 5.1, *)
-  @_silgen_name("$sScG22awaitAllRemainingTasksyyYaF")
   internal mutating func awaitAllRemainingTasks() async {
     while let _ = await next(isolation: nil) {}
   }

stdlib/toolchain/Compatibility56/Concurrency/Actor.cpp

@@ -67,7 +67,7 @@ class ExecutorTrackingInfo {
 
   /// Unconditionally initialize a fresh tracking state on the
   /// current state, shadowing any previous tracking state.
-  /// leave() must be called beforet the object goes out of scope.
+  /// leave() must be called before the object goes out of scope.
   void enterAndShadow(ExecutorRef currentExecutor) {
     ActiveExecutor = currentExecutor;
     SavedInfo = ActiveInfoInThread.get();

test/Concurrency/Runtime/continuation_validation.swift

@@ -16,11 +16,72 @@
 
 import StdlibUnittest
 
+@MainActor
+@available(SwiftStdlib 5.1, *)
+func test_isolation_withUnsafeContinuation() async {
+  // This test specifically should have only one suspension point,
+  // as it would trigger a problem with the previous @_unsafeInheritExecutor
+  // implementation, where we optimize away a switch accidentally, causing
+  // wrong isolation.
+  await withUnsafeContinuation { continuation in
+    MainActor.shared.assertIsolated() // OK
+    continuation.resume(returning: ())
+  }
+}
+@MainActor
+@available(SwiftStdlib 5.1, *)
+func test_isolation_withUnsafeThrowingContinuation() async {
+  // See comment in `test_isolation_withUnsafeContinuation` about exact test case shape
+  try! await withUnsafeThrowingContinuation { continuation in
+    MainActor.shared.assertIsolated() // OK
+    continuation.resume(returning: ())
+  }
+}
+@MainActor
+@available(SwiftStdlib 5.1, *)
+func test_isolation_withCheckedContinuation() async {
+  // See comment in `test_isolation_withUnsafeContinuation` about exact test case shape
+  await withCheckedContinuation { continuation in
+    MainActor.shared.assertIsolated() // OK
+    continuation.resume(returning: ())
+  }
+}
+@MainActor
+@available(SwiftStdlib 5.1, *)
+func test_isolation_withCheckedThrowingContinuation() async {
+  // See comment in `test_isolation_withUnsafeContinuation` about exact test case shape
+  try! await withCheckedThrowingContinuation { continuation in
+    MainActor.shared.assertIsolated() // OK
+    continuation.resume(returning: ())
+  }
+}
+
 @main struct Main {
   static func main() async {
     let tests = TestSuite("ContinuationValidation")
 
     if #available(SwiftStdlib 5.1, *) {
+      tests.test("withUnsafeThrowingContinuation: continuation should be on calling isolation") {
+        await Task.detached {
+          await test_isolation_withUnsafeThrowingContinuation()
+        }.value
+      }
+      tests.test("withUnsafeContinuation: continuation should be on calling isolation") {
+        await Task.detached {
+          await test_isolation_withUnsafeContinuation()
+        }.value
+      }
+      tests.test("withCheckedContinuation: continuation should be on calling isolation") {
+        await Task.detached {
+          await test_isolation_withCheckedContinuation()
+        }.value
+      }
+      tests.test("withCheckedThrowingContinuation: continuation should be on calling isolation") {
+        await Task.detached {
+          await test_isolation_withCheckedThrowingContinuation()
+        }.value
+      }
+
       tests.test("trap on double resume of unchecked continuation") {
         expectCrashLater(withMessage: "may have already been resumed")
 

test/abi/macOS/arm64/concurrency.swift

@@ -270,6 +270,14 @@ Added: _swift_task_getPreferredTaskExecutor
 Added: _swift_task_popTaskExecutorPreference
 Added: _swift_task_pushTaskExecutorPreference
 
+// Adopt #isolation in with...Continuation APIs
+// Swift.withCheckedThrowingContinuation<A>(isolation: isolated Swift.Actor?, function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Error>) -> ()) async throws -> A
+Added: _$ss31withCheckedThrowingContinuation9isolation8function_xScA_pSgYi_SSyScCyxs5Error_pGXEtYaKlF
+Added: _$ss31withCheckedThrowingContinuation9isolation8function_xScA_pSgYi_SSyScCyxs5Error_pGXEtYaKlFTu
+// Swift.withCheckedContinuation<A>(isolation: isolated Swift.Actor?, function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Never>) -> ()) async -> A
+Added: _$ss23withCheckedContinuation9isolation8function_xScA_pSgYi_SSyScCyxs5NeverOGXEtYalF
+Added: _$ss23withCheckedContinuation9isolation8function_xScA_pSgYi_SSyScCyxs5NeverOGXEtYalFTu
+
 // Updated signature for withTaskExecutorPreference to used typed throws and #isolation
 // Swift.withTaskExecutorPreference<A, B where B: Swift.Error>(_: Swift.TaskExecutor?, isolation: isolated Swift.Actor?, operation: () async throws(B) -> A) async throws(B) -> A
 Added: _$ss26withTaskExecutorPreference_9isolation9operationxSch_pSg_ScA_pSgYixyYaq_YKXEtYaq_YKs5ErrorR_r0_lF

test/api-digester/stability-concurrency-abi.test

@@ -78,6 +78,19 @@ Protocol SerialExecutor has added inherited protocol Copyable
 Protocol SerialExecutor has added inherited protocol Escapable
 Protocol AsyncSequence is now without @rethrows
 
+// #isolated adoption in with...Continuation
+// api-digester is not aware of silgen_name trickery we do to keep this ABI compatible
+Func withCheckedContinuation(function:_:) has parameter 0 type change from Swift.String to (any _Concurrency.Actor)?
+Func withCheckedContinuation(function:_:) has parameter 1 type change from (_Concurrency.CheckedContinuation<τ_0_0, Swift.Never>) -> () to Swift.String
+Func withCheckedThrowingContinuation(function:_:) has been renamed to Func withCheckedThrowingContinuation(isolation:function:_:)
+Func withCheckedThrowingContinuation(function:_:) has mangled name changing from '_Concurrency.withCheckedThrowingContinuation<A>(function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Error>) -> ()) async throws -> A' to '_Concurrency.withCheckedThrowingContinuation<A>(isolation: isolated Swift.Optional<Swift.Actor>, function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Error>) -> ()) async throws -> A'
+Func withCheckedThrowingContinuation(function:_:) has parameter 0 type change from Swift.String to (any _Concurrency.Actor)?
+Func withCheckedThrowingContinuation(function:_:) has parameter 1 type change from (_Concurrency.CheckedContinuation<τ_0_0, any Swift.Error>) -> () to Swift.String
+
+// #isolated was adopted and the old methods kept: $ss31withCheckedThrowingContinuation8function_xSS_yScCyxs5Error_pGXEtYaKlF
+Func withCheckedContinuation(function:_:) has been renamed to Func withCheckedContinuation(isolation:function:_:)
+Func withCheckedContinuation(function:_:) has mangled name changing from '_Concurrency.withCheckedContinuation<A>(function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Never>) -> ()) async -> A' to '_Concurrency.withCheckedContinuation<A>(isolation: isolated Swift.Optional<Swift.Actor>, function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Never>) -> ()) async -> A'
+
 // SerialExecutor gained `enqueue(_: __owned Job)`, protocol requirements got default implementations
 Func SerialExecutor.enqueue(_:) has been added as a protocol requirement