[Concurrency] task local must not copy null values

Author: ktoso

include/swift/ABI/TaskLocal.h

@@ -174,7 +174,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 (!valueType)
+    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,43 @@ 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: "")
+    }
+  }
+  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()
+  }
+}

test/Concurrency/Runtime/async_task_locals_repro.swift

@@ -0,0 +1,63 @@
+// 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 _Concurrency
+
+struct MyError: Error {}
+struct Image {
+}
+@available(SwiftStdlib 5.5, *)
+func downloadImage(from url: String) async throws -> Image {
+  await Task.sleep(10_000)
+  return Image()
+}
+@available(SwiftStdlib 5.5, *)
+actor ImageDownloader {
+  private enum CacheEntry {
+    case inProgress(Task<Image, Error>)
+    case ready(Image)
+  }
+  private var cache: [String: CacheEntry] = [:]
+  func image(from url: String) async throws -> Image? {
+    if let cached = cache[url] {
+      switch cached {
+      case .ready(let image):
+        return image
+      case .inProgress(let handle):
+        return try await handle.value
+      }
+    }
+    let handle = Task {
+      try await downloadImage(from: url)
+    }
+    cache[url] = .inProgress(handle)
+    do {
+      let image = try await handle.value
+      cache[url] = .ready(image)
+      return image
+    } catch {
+      cache[url] = nil
+      throw error
+    }
+  }
+}
+@available(SwiftStdlib 5.5, *)
+@main struct Main {
+  static func main() async {
+    let downloader = ImageDownloader()
+    async let image1 = downloader.image(from: "https://homepages.cae.wisc.edu/~ece533/images/airplane.png")
+    async let image2 = downloader.image(from: "https://homepages.cae.wisc.edu/~ece533/images/airplane.png")
+    async let image3 = downloader.image(from: "https://homepages.cae.wisc.edu/~ece533/images/airplane.png")
+    print("\(try! await image1!)")
+    print("\(try! await image2!)")
+    print("\(try! await image3!)")
+    print("okey") // CHECK: OK
+  }
+}