Merge pull request swiftlang#36034 from ktoso/wip-fix-withLocal-non-escape

Author: swift-ci

stdlib/public/Concurrency/TaskLocal.swift

@@ -89,48 +89,22 @@ extension Task {
   ///
   /// - Parameters:
   ///   - keyPath: key path to the `TaskLocalKey` to be used for lookup
-  ///   - value:
-  ///   - body:
+  ///   - value: value to bind the task local to for the scope of `operation`
+  ///   - operation: the operation to run with the task local value bound
   /// - Returns: the value returned by the `body` function.
   public static func withLocal<Key, BodyResult>(
     _ keyPath: KeyPath<TaskLocalValues, Key>,
     boundTo value: Key.Value,
-    body: @escaping () async -> BodyResult
-  ) async -> BodyResult where Key: TaskLocalKey {
+    operation: () async throws -> BodyResult
+  ) async rethrows -> BodyResult where Key: TaskLocalKey {
     let task = Builtin.getCurrentAsyncTask()
 
     _taskLocalValuePush(task, keyType: Key.self, value: value)
-  
-    defer {
-      _taskLocalValuePop(task)
-    }
+    defer { _taskLocalValuePop(task) }
 
-    return await body()
+    return try await operation()
   }
 
-  /// Bind the task local key to the given value for the scope of the `body` function.
-  /// Any child tasks spawned within this scope will inherit the binding.
-  ///
-  /// - Parameters:
-  ///   - key:
-  ///   - value:
-  ///   - body:
-  /// - Returns: the value returned by the `body` function, or throws.
-  public static func withLocal<Key, BodyResult>(
-    _ keyPath: KeyPath<TaskLocalValues, Key>,
-    boundTo value: Key.Value,
-    body: @escaping () async throws -> BodyResult
-  ) async throws -> BodyResult where Key: TaskLocalKey {
-    let task = Builtin.getCurrentAsyncTask()
-
-    _taskLocalValuePush(task, keyType: Key.self, value: value)
-
-    defer {
-      _taskLocalValuePop(task)
-    }
-
-    return try! await body()
-  }
 }
 
 // ==== ------------------------------------------------------------------------

test/Concurrency/Runtime/task_locals_basic.swift

@@ -79,6 +79,20 @@ func simple_deinit() async {
   try! await printTaskLocal(\.clazz) // CHECK: ClazzKey: nil {{.*}}
 }
 
+struct Boom: Error {
+  let value: String
+}
+func simple_throw() async {
+  do {
+    try await Task.withLocal(\.clazz, boundTo: ClassTaskLocal()) {
+      throw Boom(value: "oh no!")
+    }
+  } catch {
+    //CHECK: error: Boom(value: "oh no!")
+    print("error: \(error)")
+  }
+}
+
 func nested() async {
   try! await printTaskLocal(\.string) // CHECK: StringKey: <undefined> {{.*}}
   await Task.withLocal(\.string, boundTo: "hello") {
@@ -127,10 +141,18 @@ func nested_3_onlyTopContributes() async {
   try! await printTaskLocal(\.string) // CHECK-NEXT: StringKey: <undefined> {{.*}}
 }
 
+func withLocal_body_mustNotEscape() async {
+  var something = "Nice"
+  await Task.withLocal(\.string, boundTo: "xxx") {
+    something = "very nice"
+  }
+}
+
 @main struct Main {
   static func main() async {
     await simple()
     await simple_deinit()
+    await simple_throw()
     await nested()
     await nested_allContribute()
     await nested_3_onlyTopContributes()