[Concurrency] waitForAll and next of TaskGroups must inherit isolation

ktoso · ktoso · commit 27a8fe8a7bc7 · 2024-04-04T12:39:41.000+09:00
Otherwise we get warnings when task groups are used within an actor.
Add global actor tests to cover the specific situation.

resolves rdar://122846553
diff --git a/stdlib/public/Concurrency/TaskGroup.swift b/stdlib/public/Concurrency/TaskGroup.swift
@@ -80,6 +80,7 @@ public func withTaskGroup<ChildTaskResult, GroupResult>(
   // Run the withTaskGroup body.
   let result = await body(&group)
 
+  // TODO(concurrency): should get isolation from param from withThrowingTaskGroup
   await group.awaitAllRemainingTasks()
 
   Builtin.destroyTaskGroup(_group)
@@ -183,13 +184,15 @@ public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
     // Run the withTaskGroup body.
     let result = try await body(&group)
 
+    // TODO(concurrency): should get isolation from param from withThrowingTaskGroup
     await group.awaitAllRemainingTasks()
     Builtin.destroyTaskGroup(_group)
 
     return result
   } catch {
     group.cancelAll()
 
+    // TODO(concurrency): should get isolation from param from withThrowingTaskGroup
     await group.awaitAllRemainingTasks()
     Builtin.destroyTaskGroup(_group)
 
@@ -563,22 +566,35 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
   /// that method can't be called from a concurrent execution context like a child task.
   ///
   /// - Returns: The value returned by the next child task that completes.
-  public mutating func next() async -> ChildTaskResult? {
+  public mutating func next(isolation: isolated (any Actor)? = #isolation) async -> ChildTaskResult? {
+    // try!-safe because this function only exists for Failure == Never,
+    // and as such, it is impossible to spawn a throwing child task.
+    return try! await _taskGroupWaitNext(group: _group) // !-safe cannot throw, we're a non-throwing TaskGroup
+  }
+
+  @usableFromInline
+  @_silgen_name("$sScG4nextxSgyYaF")
+  internal mutating func __abi_next() async -> ChildTaskResult? {
     // try!-safe because this function only exists for Failure == Never,
     // and as such, it is impossible to spawn a throwing child task.
     return try! await _taskGroupWaitNext(group: _group) // !-safe cannot throw, we're a non-throwing TaskGroup
   }
 
   /// Await all of the pending tasks added this group.
   @usableFromInline
-  internal mutating func awaitAllRemainingTasks() async {
-    while let _ = await next() {}
+  internal mutating func awaitAllRemainingTasks(isolation: isolated (any Actor)? = #isolation) async {
+    while let _ = await next(isolation: isolation) {}
+  }
+  @usableFromInline
+  @_silgen_name("$sScG22awaitAllRemainingTasksyyYaF")
+  internal mutating func __abi_awaitAllRemainingTasks() async {
+    while let _ = await next(isolation: nil) {}
   }
 
   /// Wait for all of the group's remaining tasks to complete.
   @_alwaysEmitIntoClient
-  public mutating func waitForAll() async {
-    await awaitAllRemainingTasks()
+  public mutating func waitForAll(isolation: isolated (any Actor)? = #isolation) async {
+    await awaitAllRemainingTasks(isolation: isolation)
   }
 
   /// A Boolean value that indicates whether the group has any remaining tasks.
@@ -703,15 +719,19 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
 
   /// Await all the remaining tasks on this group.
   @usableFromInline
-  internal mutating func awaitAllRemainingTasks() async {
+  internal mutating func awaitAllRemainingTasks(isolation: isolated (any Actor)? = #isolation) async {
     while true {
       do {
-        guard let _ = try await next() else {
+        guard let _ = try await next(isolation: isolation) else {
           return
         }
       } catch {}
     }
   }
+  @usableFromInline
+  internal mutating func awaitAllRemainingTasks() async {
+    await awaitAllRemainingTasks(isolation: nil)
+  }
 
   @usableFromInline
   internal mutating func _waitForAll() async throws {
@@ -750,7 +770,7 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   /// - Throws: The *first* error that was thrown by a child task during draining all the tasks.
   ///           This first error is stored until all other tasks have completed, and is re-thrown afterwards.
   @_alwaysEmitIntoClient
-  public mutating func waitForAll() async throws {
+  public mutating func waitForAll(isolation: isolated (any Actor)? = #isolation) async throws {
     var firstError: Error? = nil
 
     // Make sure we loop until all child tasks have completed
@@ -999,22 +1019,14 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   /// - Throws: The error thrown by the next child task that completes.
   ///
   /// - SeeAlso: `nextResult()`
-  public mutating func next() async throws -> ChildTaskResult? {
+  public mutating func next(isolation: isolated (any Actor)? = #isolation) async throws -> ChildTaskResult? {
     return try await _taskGroupWaitNext(group: _group)
   }
 
-  @_silgen_name("$sScg10nextResults0B0Oyxq_GSgyYaKF")
   @usableFromInline
-  mutating func nextResultForABI() async throws -> Result<ChildTaskResult, Failure>? {
-    do {
-      guard let success: ChildTaskResult = try await _taskGroupWaitNext(group: _group) else {
-        return nil
-      }
-
-      return .success(success)
-    } catch {
-      return .failure(error as! Failure) // as!-safe, because we are only allowed to throw Failure (Error)
-    }
+  @_silgen_name("$sScg4nextxSgyYaKF")
+  internal mutating func __abi_next() async throws -> ChildTaskResult? {
+    return try await _taskGroupWaitNext(group: _group)
   }
 
   /// Wait for the next child task to complete,
@@ -1052,10 +1064,23 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   ///
   /// - SeeAlso: `next()`
   @_alwaysEmitIntoClient
-  public mutating func nextResult() async -> Result<ChildTaskResult, Failure>? {
-    return try! await nextResultForABI()
+  public mutating func nextResult(isolation: isolated (any Actor)? = #isolation) async -> Result<ChildTaskResult, Failure>? {
+    return try! await __abi_nextResult()
   }
 
+  @_silgen_name("$sScg10nextResults0B0Oyxq_GSgyYaKF")
+  @usableFromInline
+  mutating func __abi_nextResult() async throws -> Result<ChildTaskResult, Failure>? {
+    do {
+      guard let success: ChildTaskResult = try await _taskGroupWaitNext(group: _group) else {
+        return nil
+      }
+
+      return .success(success)
+    } catch {
+      return .failure(error as! Failure) // as!-safe, because we are only allowed to throw Failure (Error)
+    }
+  }
   /// A Boolean value that indicates whether the group has any remaining tasks.
   ///
   /// At the start of the body of a `withThrowingTaskGroup(of:returning:body:)` call,
diff --git a/test/Concurrency/async_task_groups_and_actors.swift b/test/Concurrency/async_task_groups_and_actors.swift
@@ -0,0 +1,36 @@
+// RUN: %target-swift-frontend  -disable-availability-checking %s -emit-sil -o /dev/null -verify -strict-concurrency=complete -enable-upcoming-feature RegionBasedIsolation
+
+// REQUIRES: concurrency
+// REQUIRES: asserts
+// REQUIRES: libdispatch
+
+@MainActor
+class MyActor {
+  func check() async throws {
+    await withTaskGroup(of: Int.self) { group in
+      group.addTask {
+        2
+      }
+      await group.waitForAll()
+    }
+
+    try await withThrowingTaskGroup(of: Int.self) { throwingGroup in
+      throwingGroup.addTask {
+        2
+      }
+      try await throwingGroup.waitForAll()
+    }
+
+    await withDiscardingTaskGroup { discardingGroup in
+      discardingGroup.addTask {
+        ()
+      }
+    }
+
+    try await withThrowingDiscardingTaskGroup { throwingDiscardingGroup in
+      throwingDiscardingGroup.addTask {
+        ()
+      }
+    }
+  }
+}
diff --git a/test/abi/macOS/arm64/concurrency.swift b/test/abi/macOS/arm64/concurrency.swift
@@ -276,6 +276,19 @@ Added: _$ss26withTaskExecutorPreference_9isolation9operationxSch_pSg_ScA_pSgYixy
 // async function pointer to 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_lFTu
 
+// === Add #isolation to next() and waitForAll() in task groups
+// Swift.TaskGroup.awaitAllRemainingTasks(isolation: isolated Swift.Actor?) async -> ()
+Added: _$sScG22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaF
+Added: _$sScG22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaFTu
+// Swift.TaskGroup.next(isolation: isolated Swift.Actor?) async -> A?
+Added: _$sScG4next9isolationxSgScA_pSgYi_tYaF
+Added: _$sScG4next9isolationxSgScA_pSgYi_tYaFTu
+// Swift.ThrowingTaskGroup.next(isolation: isolated Swift.Actor?) async throws -> A?
+Added: _$sScg4next9isolationxSgScA_pSgYi_tYaKF
+Added: _$sScg4next9isolationxSgScA_pSgYi_tYaKFTu
+// Swift.ThrowingTaskGroup.awaitAllRemainingTasks(isolation: isolated Swift.Actor?) async -> ()
+Added: _$sScg22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaF
+Added: _$sScg22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaFTu
 
 // next() default implementation in terms of next(isolation:)
 Added: _$sScIsE4next7ElementQzSgyYa7FailureQzYKF
diff --git a/test/abi/macOS/x86_64/concurrency.swift b/test/abi/macOS/x86_64/concurrency.swift
@@ -276,6 +276,19 @@ Added: _$ss26withTaskExecutorPreference_9isolation9operationxSch_pSg_ScA_pSgYixy
 // async function pointer to 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_lFTu
 
+// === Add #isolation to next() and waitForAll() in task groups
+// Swift.TaskGroup.awaitAllRemainingTasks(isolation: isolated Swift.Actor?) async -> ()
+Added: _$sScG22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaF
+Added: _$sScG22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaFTu
+// Swift.TaskGroup.next(isolation: isolated Swift.Actor?) async -> A?
+Added: _$sScG4next9isolationxSgScA_pSgYi_tYaF
+Added: _$sScG4next9isolationxSgScA_pSgYi_tYaFTu
+// Swift.ThrowingTaskGroup.next(isolation: isolated Swift.Actor?) async throws -> A?
+Added: _$sScg4next9isolationxSgScA_pSgYi_tYaKF
+Added: _$sScg4next9isolationxSgScA_pSgYi_tYaKFTu
+// Swift.ThrowingTaskGroup.awaitAllRemainingTasks(isolation: isolated Swift.Actor?) async -> ()
+Added: _$sScg22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaF
+Added: _$sScg22awaitAllRemainingTasks9isolationyScA_pSgYi_tYaFTu
 
 // next() default implementation in terms of next(isolation:)
 Added: _$sScIsE4next7ElementQzSgyYa7FailureQzYKF
