Merge pull request #38218 from ktoso/wip-tasklocal

Author: swift-ci

include/swift/ABI/TaskLocal.h

@@ -69,6 +69,15 @@ class TaskLocal {
     // Trailing storage for the value itself. The storage will be
     // uninitialized or contain an instance of \c valueType.
 
+    /// Returns true if this item is a 'parent pointer'.
+    ///
+    /// A parent pointer is special kind of `Item` is created when pointing at
+    /// the parent storage, forming a chain of task local items spanning multiple
+    /// tasks.
+    bool isParentPointer() const {
+      return !valueType;
+    }
+
   protected:
     explicit Item()
       : next(0),
@@ -174,7 +183,7 @@ class TaskLocal {
     /// returns, as such, any child tasks potentially accessing the value stack
     /// are guaranteed to be completed by the time we pop values off the stack
     /// (after the body has completed).
-    TaskLocal::Item *head;
+    TaskLocal::Item *head = nullptr;
 
   public:
 

stdlib/public/Concurrency/TaskLocal.cpp

@@ -164,11 +164,9 @@ void TaskLocal::Storage::initializeLinkParent(AsyncTask* task,
 TaskLocal::Item*
 TaskLocal::Item::createParentLink(AsyncTask *task, AsyncTask *parent) {
   size_t amountToAllocate = Item::itemSize(/*valueType*/nullptr);
-  // assert(amountToAllocate % MaximumAlignment == 0); // TODO: do we need this?
   void *allocation = _swift_task_alloc_specific(task, amountToAllocate);
   Item *item = new(allocation) Item();
 
-  // FIXME: parent pointer must to be the parent STORAGE not just the next item.
   auto parentHead = parent->_private().Local.head;
   if (parentHead) {
     if (parentHead->isEmpty()) {
@@ -223,8 +221,14 @@ TaskLocal::Item::createLink(AsyncTask *task,
 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_initializeWithCopy(item->getStoragePtr(), this->getStoragePtr());
+  // 'parent' pointers are signified by null valueType.
+  // We must not copy parent pointers, but rather perform a deep copy of all values,
+  // as such, we skip parent pointers here entirely.
+  if (isParentPointer())
+    return;
+
+  auto item = Item::createLink(target, key, valueType);
+  valueType->vw_initializeWithCopy(item->getStoragePtr(), 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

test/Concurrency/Runtime/async_task_locals_copy_to_async.swift

@@ -95,7 +95,7 @@ class CustomClass {
 }
 
 @available(SwiftStdlib 5.5, *)
-func test_async_retains() async {
+func test_unstructured_retains() async {
   let instance = CustomClass()
   CustomClass.$current.withValue(instance) {
     print("BEFORE send: \(String(reflecting: CustomClass.current))")
@@ -119,12 +119,46 @@ func test_async_retains() async {
   await Task.sleep(2 * 1_000_000_000)
 }
 
+@available(SwiftStdlib 5.5, *)
+func test_unstructured_noValues() async {
+  await Task {
+    // no values to copy
+  }.value
+}
+
+@available(SwiftStdlib 5.5, *)
+func downloadImage(from url: String) async throws -> String {
+  await Task.sleep(10_000)
+  return ""
+}
+
+@available(SwiftStdlib 5.5, *)
+func test_unstructured_noValues_childTasks() async {
+  @Sendable func work() async throws {
+    let handle = Task {
+      try await downloadImage(from: "")
+    }
+  }
+
+  // these child tasks have a parent pointer in their task local storage.
+  // we must not copy it when performing the copyTo for a new unstructured task.
+  async let one = work()
+  async let two = work()
+  async let three = work()
+
+  try! await one
+  try! await two
+  try! await three
+
+}
+
 @available(SwiftStdlib 5.5, *)
 @main struct Main {
   static func main() async {
-    await copyTo_async()
     await copyTo_async()
     await copyTo_async_noWait()
-    await test_async_retains()
+    await test_unstructured_retains()
+    await test_unstructured_noValues()
+    await test_unstructured_noValues_childTasks()
   }
 }

test/Concurrency/Runtime/async_task_locals_copy_to_sync.swift

@@ -0,0 +1,83 @@
+// 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
+
+import Dispatch
+
+// For sleep
+#if canImport(Darwin)
+import Darwin
+#elseif canImport(Glibc)
+import Glibc
+#endif
+
+@available(SwiftStdlib 5.5, *)
+enum TL {
+  @TaskLocal
+  static var number: Int = 0
+  @TaskLocal
+  static var other: Int = 0
+}
+
+@available(SwiftStdlib 5.5, *)
+@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(SwiftStdlib 5.5, *)
+func copyTo_sync_noWait() {
+  print(#function)
+  TL.$number.withValue(1111) {
+    TL.$number.withValue(2222) {
+      TL.$other.withValue(9999) {
+        Task {
+          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)
+          }
+        }
+      }
+    }
+  }
+
+  sleep(1)
+}
+
+@available(SwiftStdlib 5.5, *)
+func copyTo_sync_noValues() {
+  Task {
+    printTaskLocal(TL.$number) // CHECK: TaskLocal<Int>(defaultValue: 0) (0)
+  }
+
+  sleep(1)
+}
+
+/// Similar to tests in `async_task_locals_copy_to_async_ but without any task involved at the top level.
+@available(SwiftStdlib 5.5, *)
+@main struct Main {
+  static func main() {
+    copyTo_sync_noWait()
+    copyTo_sync_noValues()
+  }
+}