[Concurrency] Correct enqueue behavior in with... functions by adopting nonisolated(nonsending)

stdlib/public/Concurrency/CheckedContinuation.swift

@@ -290,12 +290,31 @@ extension CheckedContinuation {
 /// - SeeAlso: `withCheckedThrowingContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeThrowingContinuation(function:_:)`
+@available(SwiftStdlib 5.1, *)
+@export(implementation)
+@_silgen_name("$ss25withCheckedContinuation_X8function_xSS_yScCyxs5NeverOGXEtYalF") // The `_X` suffix is only here to avoid silgen_name clash with original _unsafeInheritExecutor ABI
+public nonisolated(nonsending) func withCheckedContinuation<T>(
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Never>) -> Void
+) async -> sending T {
+  return await Builtin.withUnsafeContinuation {
+    let unsafeContinuation = unsafe UnsafeContinuation<T, Never>($0)
+    return body(unsafe CheckedContinuation(continuation: unsafeContinuation,
+                                           function: function))
+  }
+}
+
 @inlinable
 @available(SwiftStdlib 5.1, *)
 #if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
 #endif
-public func withCheckedContinuation<T>(
+@abi(func withCheckedContinuation<T>(
+  isolation: isolated (any Actor)?,
+  function: String,
+  _ body: (CheckedContinuation<T, Never>) -> Void
+) async -> sending T)
+public func _isolatedParam_withCheckedContinuation<T>(
   isolation: isolated (any Actor)? = #isolation,
   function: String = #function,
   _ body: (CheckedContinuation<T, Never>) -> Void
@@ -325,7 +344,6 @@ public func _unsafeInheritExecutor_withCheckedContinuation<T>(
   }
 }
 
-
 /// 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
@@ -354,12 +372,36 @@ public func _unsafeInheritExecutor_withCheckedContinuation<T>(
 /// - SeeAlso: `withCheckedContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeContinuation(function:_:)`
 /// - SeeAlso: `withUnsafeThrowingContinuation(function:_:)`
+@available(SwiftStdlib 5.1, *)
+@export(implementation)
+@_silgen_name("$ss33withCheckedThrowingContinuation_X8function_xSS_yScCyxs5Error_pGXEtYaKlF") // The `_X` suffix is only here to avoid silgen_name clash with original _unsafeInheritExecutor ABI
+public nonisolated(nonsending) func withCheckedThrowingContinuation<T>(
+  function: String = #function,
+  _ body: (CheckedContinuation<T, Error>) -> Void
+) async throws -> sending T {
+  // ABI NOTE: Interestingly enough the 'nonisolated(nonsending)' version of this func has the same ABI
+  // as the initial implementation, that was '@_unsafeInheritExecutor'. So we can remove the "ABI shim"
+  // as the current correct impl and the previous impl are both one and the same.
+  //
+  // We need to keep around the isolated parameter ABI for compatibility though.
+  return try await Builtin.withUnsafeThrowingContinuation {
+    let unsafeContinuation = unsafe UnsafeContinuation<T, Error>($0)
+    return body(unsafe CheckedContinuation(continuation: unsafeContinuation,
+                                           function: function))
+  }
+}
+
 @inlinable
 @available(SwiftStdlib 5.1, *)
 #if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
 #endif
-public func withCheckedThrowingContinuation<T>(
+@abi(func withCheckedThrowingContinuation<T>(
+  isolation: isolated (any Actor)?,
+  function: String,
+  _ body: (CheckedContinuation<T, Error>) -> Void
+) async throws -> sending T)
+func _isolatedParam_withCheckedThrowingContinuation<T>(
   isolation: isolated (any Actor)? = #isolation,
   function: String = #function,
   _ body: (CheckedContinuation<T, Error>) -> Void

stdlib/public/Concurrency/DiscardingTaskGroup.swift

@@ -69,12 +69,38 @@ import Swift
 ///
 /// - SeeAlso: ``TaskGroup``
 /// - SeeAlso: ``withThrowingDiscardingTaskGroup(returning:body:)``
+@available(SwiftStdlib 5.9, *)
+@export(implementation)
+public nonisolated(nonsending) func withDiscardingTaskGroup<GroupResult>(
+  returning returnType: GroupResult.Type = GroupResult.self,
+  body: nonisolated(nonsending) (inout DiscardingTaskGroup) async -> GroupResult
+) async -> GroupResult {
+  let flags = taskGroupCreateFlags(
+    discardResults: true
+  )
+
+  let _group = Builtin.createTaskGroupWithFlags(flags, Void.self)
+  var group = DiscardingTaskGroup(group: _group)
+  defer { Builtin.destroyTaskGroup(_group) }
+
+  let result = await body(&group)
+
+  try! await group.awaitAllRemainingTasksNonsending() // try!-safe, cannot throw since this is a non throwing group
+
+  return result
+}
+
 @available(SwiftStdlib 5.9, *)
 #if !hasFeature(Embedded)
 @backDeployed(before: SwiftStdlib 6.0)
 #endif
 @inlinable
-public func withDiscardingTaskGroup<GroupResult>(
+@abi(func withDiscardingTaskGroup<GroupResult>(
+  returning returnType: GroupResult.Type,
+  isolation: isolated (any Actor)?,
+  body: (inout DiscardingTaskGroup) async -> GroupResult
+) async -> GroupResult)
+func _isolatedParameter_withDiscardingTaskGroup<GroupResult>(
   returning returnType: GroupResult.Type = GroupResult.self,
   isolation: isolated (any Actor)? = #isolation,
   body: (inout DiscardingTaskGroup) async -> GroupResult
@@ -182,6 +208,11 @@ public struct DiscardingTaskGroup {
   /// Await all the remaining tasks on this group.
   ///
   /// - Throws: The first error that was encountered by this group.
+  @export(implementation)
+  internal nonisolated(nonsending) mutating func awaitAllRemainingTasksNonsending() async throws {
+    let _: Void? = try await _taskGroupWaitAll(group: _group, bodyError: nil)
+  }
+
   @usableFromInline
   internal mutating func awaitAllRemainingTasks() async throws {
     let _: Void? = try await _taskGroupWaitAll(group: _group, bodyError: nil)
@@ -336,15 +367,50 @@ extension DiscardingTaskGroup: Sendable { }
 ///   }
 /// }
 /// ```
+@available(SwiftStdlib 5.9, *)
+@export(implementation)
+public nonisolated(nonsending) func withThrowingDiscardingTaskGroup<GroupResult>(
+  returning returnType: GroupResult.Type = GroupResult.self,
+  body: nonisolated(nonsending) (inout ThrowingDiscardingTaskGroup<Error>) async throws -> GroupResult
+) async throws -> GroupResult {
+  let flags = taskGroupCreateFlags(
+      discardResults: true
+  )
+
+  let _group = Builtin.createTaskGroupWithFlags(flags, Void.self)
+  var group = ThrowingDiscardingTaskGroup<Error>(group: _group)
+  defer { Builtin.destroyTaskGroup(_group) }
+
+  let result: GroupResult
+  do {
+    result = try await body(&group)
+  } catch {
+    group.cancelAll()
+
+    try await group.awaitAllRemainingTasksNonisolated(bodyError: error)
+
+    throw error
+  }
+
+  try await group.awaitAllRemainingTasksNonisolated(bodyError: nil)
+
+  return result
+}
+
 @available(SwiftStdlib 5.9, *)
 #if !hasFeature(Embedded)
 @backDeployed(before: SwiftStdlib 6.0)
 #endif
 @inlinable
-public func withThrowingDiscardingTaskGroup<GroupResult>(
-    returning returnType: GroupResult.Type = GroupResult.self,
-    isolation: isolated (any Actor)? = #isolation,
-    body: (inout ThrowingDiscardingTaskGroup<Error>) async throws -> GroupResult
+@abi(func withThrowingDiscardingTaskGroup<GroupResult>(
+  returning returnType: GroupResult.Type,
+  isolation: isolated (any Actor)?,
+  body: (inout ThrowingDiscardingTaskGroup<Error>) async throws -> GroupResult
+) async throws -> GroupResult)
+public func _isolatedParam_withThrowingDiscardingTaskGroup<GroupResult>(
+  returning returnType: GroupResult.Type = GroupResult.self,
+  isolation: isolated (any Actor)? = #isolation,
+  body: (inout ThrowingDiscardingTaskGroup<Error>) async throws -> GroupResult
 ) async throws -> GroupResult {
   let flags = taskGroupCreateFlags(
       discardResults: true
@@ -474,6 +540,11 @@ public struct ThrowingDiscardingTaskGroup<Failure: Error> {
   }
 
   /// Await all the remaining tasks on this group.
+  @export(implementation)
+  internal nonisolated(nonsending) mutating func awaitAllRemainingTasksNonisolated(bodyError: Error?) async throws {
+    let _: Void? = try await _taskGroupWaitAll(group: _group, bodyError: bodyError)
+  }
+
   @usableFromInline
   internal mutating func awaitAllRemainingTasks(bodyError: Error?) async throws {
     let _: Void? = try await _taskGroupWaitAll(group: _group, bodyError: bodyError)

stdlib/public/Concurrency/PartialAsyncTask.swift

@@ -887,10 +887,9 @@ internal func _resumeUnsafeThrowingContinuationWithError<T>(
 /// - SeeAlso: `withCheckedContinuation(function:_:)`
 /// - SeeAlso: `withCheckedThrowingContinuation(function:_:)`
 @available(SwiftStdlib 5.1, *)
-@_alwaysEmitIntoClient
+@export(implementation)
 @unsafe
-public func withUnsafeContinuation<T>(
-  isolation: isolated (any Actor)? = #isolation,
+public nonisolated(nonsending) func withUnsafeContinuation<T>(
   _ fn: (UnsafeContinuation<T, Never>) -> Void
 ) async -> sending T {
   return await Builtin.withUnsafeContinuation {
@@ -924,10 +923,9 @@ public func withUnsafeContinuation<T>(
 /// - SeeAlso: `withCheckedContinuation(function:_:)`
 /// - SeeAlso: `withCheckedThrowingContinuation(function:_:)`
 @available(SwiftStdlib 5.1, *)
-@_alwaysEmitIntoClient
+@export(implementation)
 @unsafe
-public func withUnsafeThrowingContinuation<T>(
-  isolation: isolated (any Actor)? = #isolation,
+public nonisolated(nonsending) func withUnsafeThrowingContinuation<T>(
   _ fn: (UnsafeContinuation<T, Error>) -> Void
 ) async throws -> sending T {
   return try await Builtin.withUnsafeThrowingContinuation {

stdlib/public/Concurrency/Task+TaskExecutor.swift

@@ -208,18 +208,22 @@ extension UnsafeCurrentTask {
 // ==== Runtime ---------------------------------------------------------------
 
 @available(SwiftStdlib 6.0, *)
+@usableFromInline
 @_silgen_name("swift_task_getPreferredTaskExecutor")
 internal func _getPreferredUnownedTaskExecutor() -> Builtin.Executor
 
+@usableFromInline
 typealias TaskExecutorPreferenceStatusRecord = UnsafeRawPointer
 
 @available(SwiftStdlib 6.0, *)
+@usableFromInline
 @_silgen_name("swift_task_pushTaskExecutorPreference")
 internal func _pushTaskExecutorPreference(_ executor: Builtin.Executor)
   -> TaskExecutorPreferenceStatusRecord
 
 @available(SwiftStdlib 6.0, *)
 @_silgen_name("swift_task_popTaskExecutorPreference")
+@usableFromInline
 internal func _popTaskExecutorPreference(
   record: TaskExecutorPreferenceStatusRecord
 )

stdlib/public/Concurrency/TaskCancellation.swift

@@ -74,10 +74,25 @@ import Swift
 /// Therefore, if a cancellation handler must acquire a lock, other code should
 /// not cancel tasks or resume continuations while holding that lock.
 @available(SwiftStdlib 5.1, *)
-#if !$Embedded
-@backDeployed(before: SwiftStdlib 6.0)
-#endif
-public func withTaskCancellationHandler<T>(
+@export(implementation)
+public nonisolated(nonsending) func withTaskCancellationHandler<T>(
+  operation: nonisolated(nonsending) () async throws -> T,
+  onCancel handler: @Sendable () -> Void
+) async rethrows -> T {
+  // unconditionally add the cancellation record to the task.
+  // if the task was already cancelled, it will be executed right away.
+  let record = unsafe _taskAddCancellationHandler(handler: handler)
+  defer { unsafe _taskRemoveCancellationHandler(record: record) }
+
+  return try await operation()
+}
+
+// Note: Deprecated version which would still hop if we did not close over an `isolated` parameter
+// with the operation closure. Instead, we should do what the docs of this method promise - and not hop at all,
+// by using the new nonisolated(nonsending)
+@available(SwiftStdlib 5.1, *)
+@_silgen_name("$ss27withTaskCancellationHandler9operation8onCancel9isolationxxyYaKXE_yyYbXEScA_pSgYitYaKlF")
+public func _isolatedParam_withTaskCancellationHandler<T>(
   operation: () async throws -> T,
   onCancel handler: @Sendable () -> Void,
   isolation: isolated (any Actor)? = #isolation