[Concurrency] Correct withContinuation APIs isolation handling

rdar://125307764

Author: ktoso

stdlib/public/Concurrency/CheckedContinuation.swift

@@ -282,18 +282,32 @@ extension CheckedContinuation {
 /// - SeeAlso: `withUnsafeContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeThrowingContinuation(function:_:)`
 @available(SwiftStdlib 5.1, *)
-@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
 @inlinable
 @_unavailableInEmbedded
 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, *)
+@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
+@inlinable
+@_unavailableInEmbedded
+public func __abi_withCheckedContinuation<T>(
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Never>) -> Void
+) async -> T {
+  return await __abi_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
@@ -323,18 +337,31 @@ public func withCheckedContinuation<T>(
 /// - SeeAlso: `withUnsafeContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeThrowingContinuation(function:_:)`
 @available(SwiftStdlib 5.1, *)
-@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
 @inlinable
 @_unavailableInEmbedded
 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, *)
+@_unsafeInheritExecutor // ABI compatibility with Swift 5.1
+@inlinable
+@_unavailableInEmbedded
+public func __abi_withCheckedThrowingContinuation<T>(
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Error>) -> Void
+) async throws -> T {
+  return try await __abi_withUnsafeThrowingContinuation {
+    body(CheckedContinuation(continuation: $0, function: function))
+  }
+}
+
 #if _runtime(_ObjC)
 
 // Intrinsics used by SILGen to create, resume, or fail checked continuations.

stdlib/public/Concurrency/PartialAsyncTask.swift

@@ -598,9 +598,20 @@ 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 {
+    fn(UnsafeContinuation<T, Never>($0))
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+@_unsafeInheritExecutor
+@_alwaysEmitIntoClient
+public func __abi_withUnsafeContinuation<T>(
   _ fn: (UnsafeContinuation<T, Never>) -> Void
 ) async -> T {
   return await Builtin.withUnsafeContinuation {
@@ -634,9 +645,20 @@ 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 {
+    fn(UnsafeContinuation<T, Error>($0))
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+@_unsafeInheritExecutor
+@_alwaysEmitIntoClient
+public func __abi_withUnsafeThrowingContinuation<T>(
   _ 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/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")