[Concurrency] tasklocal withValue inside actor is safe

Author: ktoso

stdlib/public/BackDeployConcurrency/TaskLocal.swift

@@ -134,8 +134,9 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
   ///
   /// If the value is a reference type, it will be retained for the duration of
   /// the operation closure.
+  @Sendable
   @discardableResult
-  public func withValue<R>(_ valueDuringOperation: Value, operation: () async throws -> R,
+  public func withValue<R>(_ valueDuringOperation: Value, operation: @Sendable () async throws -> R,
                            file: String = #file, line: UInt = #line) async rethrows -> R {
     // check if we're not trying to bind a value from an illegal context; this may crash
     _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)

stdlib/public/Concurrency/TaskLocal.swift

@@ -135,7 +135,7 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
   /// If the value is a reference type, it will be retained for the duration of
   /// the operation closure.
   @discardableResult
-  public func withValue<R>(_ valueDuringOperation: Value, operation: () async throws -> R,
+  public func withValue<R>(_ valueDuringOperation: Value, operation: @Sendable () async throws -> R,
                            file: String = #fileID, line: UInt = #line) async rethrows -> R {
     // check if we're not trying to bind a value from an illegal context; this may crash
     _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)

test/Concurrency/Runtime/async_task_locals_basic.swift

@@ -199,6 +199,24 @@ func withLocal_body_mustNotEscape() async {
   _ = something // silence not used warning
 }
 
+@available(SwiftStdlib 5.1, *)
+actor Worker {
+  @TaskLocal
+  static var declaredInActor: String = "<default-value>"
+
+  func setAndRead() async {
+    print("setAndRead") // CHECK: setAndRead
+    await Worker.$declaredInActor.withValue("value-1") {
+      await printTaskLocalAsync(Worker.$declaredInActor) // CHECK-NEXT: TaskLocal<String>(defaultValue: <default-value>) (value-1)
+    }
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+func inside_actor() async {
+  await Worker().setAndRead()
+}
+
 @available(SwiftStdlib 5.1, *)
 @main struct Main {
   static func main() async {
@@ -210,5 +228,6 @@ func withLocal_body_mustNotEscape() async {
     await nested_3_onlyTopContributes()
     await nested_3_onlyTopContributesAsync()
     await nested_3_onlyTopContributesMixed()
+    await inside_actor()
   }
 }

test/Concurrency/Runtime/async_taskgroup_void_neverConsumingTasks.swift

@@ -0,0 +1,76 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -parse-as-library) | %FileCheck %s --dump-input=always
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+// UNSUPPORTED: OS=linux-gnu
+import Darwin
+
+actor Waiter {
+  let until: Int
+  var count: Int
+
+  var cc: CheckedContinuation<Int, Never>?
+
+  init(until: Int) {
+    self.until = until
+    self.count = 0
+  }
+
+  func increment() {
+    self.count += 1
+    fputs("> increment (\(self.count)/\(self.until))\n", stderr);
+    if self.until <= self.count {
+      if let cc = self.cc {
+        cc.resume(returning: self.count)
+      }
+    }
+  }
+
+  func wait() async -> Int {
+    if self.until <= self.count {
+      fputs("> RETURN in Waiter\n", stderr);
+      return self.count
+    }
+
+    return await withCheckedContinuation { cc in
+      fputs("> WAIT in Waiter\n", stderr);
+      self.cc = cc
+    }
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+func test_taskGroup_void_neverConsume() async {
+  let until = 100_000_000
+  let waiter = Waiter(until: until)
+
+  let allTasks = await withTaskGroup(of: Void.self, returning: Int.self) { group in
+    for n in 1...until {
+    fputs("> enqueue: \(n)\n", stderr);
+      group.addTask {
+        fputs("> run: \(n)\n", stderr);
+        try? await Task.sleep(until: .now + .milliseconds(100), clock: .continuous)
+        await waiter.increment()
+      }
+    }
+
+    let void = await group.next()
+
+    // wait a little bit, so some tasks complete before we hit the implicit "wait at end of task group scope"
+    try? await Task.sleep(until: .now + .milliseconds(500), clock: .continuous)
+
+    return until
+  }
+
+  // CHECK: all tasks: 100
+  print("all tasks: \(allTasks)")
+  print("actor: \(allTasks)")
+}
+
+@available(SwiftStdlib 5.1, *)
+@main struct Main {
+  static func main() async {
+    await test_taskGroup_void_neverConsume()
+  }
+}

test/Concurrency/async_task_locals_basic_warnings.swift

@@ -0,0 +1,23 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -emit-module-path %t/OtherActors.swiftmodule -module-name OtherActors %S/Inputs/OtherActors.swift -disable-availability-checking
+// RUN: %target-typecheck-verify-swift -I %t  -disable-availability-checking -warn-concurrency -parse-as-library
+// REQUIRES: concurrency
+
+// REQUIRES: concurrency
+
+@available(SwiftStdlib 5.1, *)
+actor Test {
+
+  @TaskLocal static var local: Int?
+
+  func run() async {
+    // This should NOT produce any warnings, the closure withValue uses is @Sendable:
+    await Test.$local.withValue(42) {
+      await work()
+    }
+  }
+
+  func work() async {
+    print("Hello \(Test.local ?? 0)")
+  }
+}