[TaskLocals] Propagate task-locals through async{}

Author: ktoso

include/swift/ABI/TaskLocal.h

@@ -118,6 +118,8 @@ class TaskLocal {
         reinterpret_cast<char *>(this) + storageOffset(valueType));
     }
 
+    void copyTo(AsyncTask *task);
+
     /// Compute the offset of the storage from the base of the item.
     static size_t storageOffset(const Metadata *valueType) {
       size_t offset = sizeof(Item);
@@ -189,6 +191,18 @@ class TaskLocal {
     /// can be safely disposed of.
     bool popValue(AsyncTask *task);
 
+    /// Copy all task-local bindings to the target task.
+    ///
+    /// The new bindings allocate their own items and can out-live the current task.
+    ///
+    /// ### Optimizations
+    /// Only the most recent binding of a value is copied over, i.e. given
+    /// a key bound to `A` and then `B`, only the `B` binding will be copied.
+    /// This is safe and correct because the new task would never have a chance
+    /// to observe the `A` value, because it semantically will never observe a
+    /// "pop" of the `B` value - it was spawned from a scope where only B was observable.
+    void copyTo(AsyncTask *target);
+
     /// Destroy and deallocate all items stored by this specific task.
     ///
     /// Items owned by a parent task are left untouched, since we do not own them.

include/swift/Runtime/Concurrency.h

@@ -480,6 +480,16 @@ void swift_task_localValuePush(const HeapObject *key,
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_localValuePop();
 
+/// Copy all task locals from the current context to the target task.
+///
+/// Its Swift signature is
+///
+/// \code
+/// func _taskLocalValueGet<Key>(AsyncTask* task)
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_localsCopyTo(AsyncTask* target);
+
 /// This should have the same representation as an enum like this:
 ///    enum NearestTaskDeadline {
 ///      case none

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -247,6 +247,12 @@ OVERRIDE_TASK_LOCAL(task_localValuePop, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                     swift::, ,)
 
+OVERRIDE_TASK_LOCAL(task_localsCopyTo, void,
+                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
+                    swift::,
+                    (AsyncTask *target),
+                    (target))
+
 OVERRIDE_TASK_STATUS(task_addStatusRecord, bool,
                      SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                      swift::, (TaskStatusRecord *newRecord), (newRecord))

stdlib/public/Concurrency/Task.cpp

@@ -195,6 +195,10 @@ AsyncTask::~AsyncTask() {
   }
 
   // Release any objects potentially held as task local values.
+  //
+  // This must be called last when destroying a task - to keep stack discipline of the allocator.
+  // because it may have created some task-locals immediately upon creation,
+  // e.g. if the task is spawned with async{} and inherited some task-locals.
   Local.destroy(this);
 }
 

stdlib/public/Concurrency/Task.swift

@@ -535,6 +535,15 @@ public func async<T>(
   // Create the asynchronous task future.
   let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, operation)
 
+  // Copy all task locals to the newly created task.
+  // We must copy them rather than point to the current task since the new task
+  // is not structured and may out-live the current task.
+  //
+  // WARNING: This MUST be done BEFORE we enqueue the task,
+  // because it acts as-if it was running inside the task and thus does not
+  // take any extra steps to synchronize the task-local operations.
+  _taskLocalsCopy(to: task)
+
   // Enqueue the resulting job.
   _enqueueJobGlobal(Builtin.convertTaskToJob(task))
 

stdlib/public/Concurrency/TaskLocal.cpp

@@ -24,6 +24,7 @@
 #include "swift/ABI/Metadata.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "TaskPrivate.h"
+#include <set>
 
 #if defined(__APPLE__)
 #include <asl.h>
@@ -134,6 +135,22 @@ static void swift_task_localValuePopImpl() {
   assert(false && "Attempted to pop value but no task or thread-local storage available!");
 }
 
+SWIFT_CC(swift)
+static void swift_task_localsCopyToImpl(AsyncTask *task) {
+  TaskLocal::Storage *Local = nullptr;
+
+  if (AsyncTask *task = swift_task_getCurrent()) {
+    Local = &task->Local;
+  } else if (auto *storage = FallbackTaskLocalStorage::get()) {
+    Local = storage;
+  } else {
+    // bail out, there are no values to copy
+    return;
+  }
+
+  Local->copyTo(task);
+}
+
 // =============================================================================
 // ==== Initialization ---------------------------------------------------------
 
@@ -185,6 +202,38 @@ TaskLocal::Item::createParentLink(AsyncTask *task, AsyncTask *parent) {
   return item;
 }
 
+TaskLocal::Item*
+TaskLocal::Item::createLink(AsyncTask *task,
+                            const HeapObject *key,
+                            const Metadata *valueType) {
+  size_t amountToAllocate = Item::itemSize(valueType);
+  void *allocation = task ? _swift_task_alloc_specific(task, amountToAllocate)
+                          : malloc(amountToAllocate);
+  Item *item = new (allocation) Item(key, valueType);
+
+  auto next = task ? task->Local.head : FallbackTaskLocalStorage::get()->head;
+  item->next = reinterpret_cast<uintptr_t>(next) |
+      static_cast<uintptr_t>(NextLinkType::IsNext);
+
+  return item;
+}
+
+
+void TaskLocal::Item::copyTo(AsyncTask *target) {
+  assert(target && "TaskLocal item attempt to copy to null target task!");
+
+  auto item = Item::createLink(target, this->key, this->valueType);
+  valueType->vw_initializeWithTake(item->getStoragePtr(), this->getStoragePtr());
+
+  /// A `copyTo` may ONLY be invoked BEFORE the task is actually scheduled,
+  /// so right now we can safely copy the value into the task without additional
+  /// synchronization.
+  target->Local.head = item;
+}
+
+// =============================================================================
+// ==== checks -----------------------------------------------------------------
+
 SWIFT_CC(swift)
 static void swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroupImpl(
     const unsigned char *file, uintptr_t fileLength,
@@ -257,22 +306,6 @@ static void swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroupImpl(
   abort();
 }
 
-TaskLocal::Item*
-TaskLocal::Item::createLink(AsyncTask *task,
-                            const HeapObject *key,
-                            const Metadata *valueType) {
-  size_t amountToAllocate = Item::itemSize(valueType);
-  void *allocation = task ? _swift_task_alloc_specific(task, amountToAllocate)
-                          : malloc(amountToAllocate);
-  Item *item = new (allocation) Item(key, valueType);
-
-  auto next = task ? task->Local.head : FallbackTaskLocalStorage::get()->head;
-  item->next = reinterpret_cast<uintptr_t>(next) |
-      static_cast<uintptr_t>(NextLinkType::IsNext);
-
-  return item;
-}
-
 // =============================================================================
 // ==== destroy ----------------------------------------------------------------
 
@@ -311,7 +344,7 @@ void TaskLocal::Storage::destroy(AsyncTask *task) {
 }
 
 // =============================================================================
-// ==== push / pop / get -------------------------------------------------------
+// ==== Task Local Storage: operations -----------------------------------------
 
 void TaskLocal::Storage::pushValue(AsyncTask *task,
                                    const HeapObject *key,
@@ -350,5 +383,32 @@ OpaqueValue* TaskLocal::Storage::getValue(AsyncTask *task,
   return nullptr;
 }
 
+
+void TaskLocal::Storage::copyTo(AsyncTask *target) {
+  assert(target && "task must not be null when copying values into it");
+  assert(!(target->Local.head) &&
+      "Task must not have any task-local values bound before copying into it");
+
+  AsyncTask *task = swift_task_getCurrent();
+
+  // Set of keys for which we already have copied to the new task.
+  // We only ever need to copy the *first* encounter of any given key,
+  // because it is the most "specific"/"recent" binding and any other binding
+  // of a key does not matter for the target task as it will never be able to
+  // observe it.
+  std::set<const HeapObject*> copied = {};
+
+  auto item = head;
+  while (item) {
+    // we only have to copy an item if it is the most recent binding of a key.
+    // i.e. if we've already seen an item for this key, we can skip it.
+    if (copied.emplace(item->key).second) {
+      item->copyTo(target);
+    }
+
+    item = item->getNext();
+  }
+}
+
 #define OVERRIDE_TASK_LOCAL COMPATIBILITY_OVERRIDE
 #include COMPATIBILITY_OVERRIDE_INCLUDE_PATH

stdlib/public/Concurrency/TaskLocal.swift

@@ -212,21 +212,27 @@ public final class TaskLocal<Value>: UnsafeSendable, CustomStringConvertible {
 
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_task_localValuePush")
-public func _taskLocalValuePush<Value>(
+func _taskLocalValuePush<Value>(
   key: Builtin.RawPointer/*: Key*/,
   value: __owned Value
 ) // where Key: TaskLocal
 
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_task_localValuePop")
-public func _taskLocalValuePop()
+func _taskLocalValuePop()
 
 @available(SwiftStdlib 5.5, *)
 @_silgen_name("swift_task_localValueGet")
-public func _taskLocalValueGet(
+func _taskLocalValueGet(
   key: Builtin.RawPointer/*Key*/
 ) -> UnsafeMutableRawPointer? // where Key: TaskLocal
 
+@available(SwiftStdlib 5.5, *)
+@_silgen_name("swift_task_localsCopyTo")
+func _taskLocalsCopy(
+  to target: Builtin.NativeObject
+)
+
 // ==== Checks -----------------------------------------------------------------
 
 @available(SwiftStdlib 5.5, *)

test/Concurrency/Runtime/async_task_locals_copy_to_async.swift

@@ -0,0 +1,90 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library %import-libdispatch) | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: libdispatch
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+enum TL {
+  @TaskLocal
+  static var number: Int = 0
+  @TaskLocal
+  static var other: Int = 0
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+@discardableResult
+func printTaskLocal<V>(
+    _ key: TaskLocal<V>,
+    _ expected: V? = nil,
+    file: String = #file, line: UInt = #line
+) -> V? {
+  let value = key.get()
+  print("\(key) (\(value)) at \(file):\(line)")
+  if let expected = expected {
+    assert("\(expected)" == "\(value)",
+        "Expected [\(expected)] but found: \(value), at \(file):\(line)")
+  }
+  return expected
+}
+
+// ==== ------------------------------------------------------------------------
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+func copyTo_async() async {
+  await TL.$number.withValue(1111) {
+    printTaskLocal(TL.$number) // CHECK: TaskLocal<Int>(defaultValue: 0) (1111)
+
+    await TL.$number.withValue(2222) {
+      await TL.$other.withValue(9999) {
+        printTaskLocal(TL.$number) // CHECK: TaskLocal<Int>(defaultValue: 0) (2222)
+        printTaskLocal(TL.$other) // CHECK: TaskLocal<Int>(defaultValue: 0) (9999)
+        let handle = async {
+          printTaskLocal(TL.$number) // CHECK: TaskLocal<Int>(defaultValue: 0) (2222)
+          printTaskLocal(TL.$other) // CHECK: TaskLocal<Int>(defaultValue: 0) (9999)
+          TL.$number.withValue(3333) {
+            printTaskLocal(TL.$number) // CHECK: TaskLocal<Int>(defaultValue: 0) (3333)
+            printTaskLocal(TL.$other) // CHECK: TaskLocal<Int>(defaultValue: 0) (9999)
+          }
+        }
+
+        _ = await handle.get()
+      }
+    }
+  }
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+func copyTo_async_noWait() async {
+  print(#function)
+  TL.$number.withValue(1111) {
+    TL.$number.withValue(2222) {
+      TL.$other.withValue(9999) {
+        async {
+          printTaskLocal(TL.$number) // CHECK: TaskLocal<Int>(defaultValue: 0) (2222)
+          printTaskLocal(TL.$other) // CHECK: TaskLocal<Int>(defaultValue: 0) (9999)
+          TL.$number.withValue(3333) {
+            printTaskLocal(TL.$number) // CHECK: TaskLocal<Int>(defaultValue: 0) (3333)
+            printTaskLocal(TL.$other) // CHECK: TaskLocal<Int>(defaultValue: 0) (9999)
+          }
+        }
+      }
+    }
+  }
+
+  let second = UInt64(100_000_000) // ns
+  await Task.sleep(2 * second)
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+@main struct Main {
+  static func main() async {
+    await copyTo_async()
+    await copyTo_async()
+    await copyTo_async_noWait()
+  }
+}

unittests/runtime/CompatibilityOverrideConcurrency.cpp

@@ -236,6 +236,10 @@ TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_localValuePop) {
   swift_task_localValuePop();
 }
 
+TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_localsCopyTo) {
+  swift_task_localsCopyTo(nullptr);
+}
+
 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_addStatusRecord) {
   swift_task_addStatusRecord(nullptr);
 }