[Concurrency] Correct handling of mixed nested values

Author: ktoso

include/swift/ABI/TaskLocal.h

@@ -47,6 +47,10 @@ class TaskLocal {
     /// The task local binding was created inside the body of a `withTaskGroup`,
     /// and therefore must either copy it, or crash when a child task is created
     /// using 'group.addTask' and it would refer to this task local.
+    ///
+    /// Items of this kind must be copied by a group child task for access
+    /// safety reasons, as otherwise the pop would happen before the child task
+    /// has completed.
     IsNextCreatedInTaskGroupBody = 0b10,
   };
 
@@ -126,6 +130,8 @@ class TaskLocal {
     }
 
     void relinkNext(Item* nextOverride) {
+      fprintf(stderr, "[%s:%d](%s) try relink item:%p\n", __FILE_NAME__, __LINE__, __FUNCTION__, this);
+      fprintf(stderr, "[%s:%d](%s) try relink to target:%p\n", __FILE_NAME__, __LINE__, __FUNCTION__, nextOverride);
       assert(!getNext() &&
                "Can only relink task local item that was not pointing at anything yet");
       assert(nextOverride->isNextLinkPointer() ||
@@ -149,7 +155,7 @@ class TaskLocal {
              NextLinkType::IsNext;
     }
 
-    bool isNextCreatedInTaskGroupBody() const {
+    bool isNextLinkPointerCreatedInTaskGroupBody() const {
       return static_cast<NextLinkType>(next & statusMask) ==
              NextLinkType::IsNextCreatedInTaskGroupBody;
     }

stdlib/public/Concurrency/TaskLocal.cpp

@@ -436,7 +436,8 @@ void TaskLocal::Storage::pushValue(AsyncTask *task,
 
   valueType->vw_initializeWithTake(item->getStoragePtr(), value);
   head = item;
-  SWIFT_TASK_LOCAL_DEBUG_LOG(item->key, "Created link item:%p", item);
+  SWIFT_TASK_LOCAL_DEBUG_LOG(item->key, "Created link item:%p, in group body:%d",
+                             item, inTaskGroupBody);
 }
 
 bool TaskLocal::Storage::popValue(AsyncTask *task) {
@@ -513,29 +514,35 @@ void TaskLocal::Storage::copyToOnlyOnlyFromCurrent(AsyncTask *target) {
   std::set<const HeapObject*> copied = {};
 
   auto item = head;
-  TaskLocal::Item *lastCopiedItem = nullptr;
+  TaskLocal::Item *copiedHead = nullptr;
   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) {
 
-      if (!item->isNextCreatedInTaskGroupBody() && lastCopiedItem) {
-        SWIFT_TASK_LOCAL_DEBUG_LOG(item->key, "break out, next item is not within body, item value [%p]", item->getStoragePtr());
+      if (!item->isNextLinkPointerCreatedInTaskGroupBody() && copiedHead) {
+        SWIFT_TASK_LOCAL_DEBUG_LOG(item->key, "break out, next item is not within body, item:%p", item);
         // The next item is not the "risky one" so we can directly link to it,
         // as we would have within normal child task relationships. E.g. this is
         // a parent or next pointer to a "safe" (withValue { withTaskGroup { ... } })
         // binding, so we re-link our current head to point at this item.
-        lastCopiedItem->relinkNext(item);
+        copiedHead->relinkNext(item);
         break;
       }
 
-      lastCopiedItem = item->copyTo(target);
+      auto copy = item->copyTo(target);
+      if (!copiedHead) {
+        SWIFT_TASK_LOCAL_DEBUG_LOG(item->key, "store copied head item:%p",
+                                   copiedHead);
+        copiedHead = copy;
+      }
+
       SWIFT_TASK_LOCAL_DEBUG_LOG(item->key, "copy value [%p] to target:%p, item:%p, copied:%p",
-                                 item->getStoragePtr(), target, item, lastCopiedItem);
+                                 item->getStoragePtr(), target, item, copiedHead);
 
       // If we didn't copy an item, e.g. because it was a pointer to parent,
       // break out of the loop and keep pointing at parent still.
-      if (lastCopiedItem == nullptr) {
+      if (!copy) {
         SWIFT_TASK_LOCAL_DEBUG_LOG(item->key, "break out, next is %p", 0);
         break;
       }

test/Concurrency/Runtime/async_task_locals_in_task_group_may_need_to_copy.swift

@@ -1,12 +1,11 @@
-// RUN: %target-run-simple-swift( -plugin-path %swift-plugin-dir -Xfrontend -disable-availability-checking -parse-as-library %import-libdispatch) | %FileCheck %s --dump-input=always
+// RUN: %target-run-simple-swift( -plugin-path %swift-plugin-dir -Xfrontend -strict-concurrency=complete -Xfrontend -disable-availability-checking -parse-as-library %import-libdispatch) | %FileCheck %s --dump-input=always
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency
 // REQUIRES: libdispatch
 // REQUIRES: concurrency_runtime
 // UNSUPPORTED: back_deployment_runtime
 
-@available(SwiftStdlib 5.1, *)
 enum TL {
   @TaskLocal
   static var one: Int = 1
@@ -22,22 +21,43 @@ func test() async {
   await TL.$one.withValue(11) {
     await TL.$one.withValue(1111) {
       await withTaskGroup(of: Void.self) { group in
+
         TL.$two.withValue(2222) {
           group.addTask { // will have to copy the `2222`
             print("Survived, one: \(TL.one) @ \(#fileID):\(#line)") // CHECK: Survived, one: 1111
             print("Survived, two: \(TL.two) @ \(#fileID):\(#line)") // CHECK: Survived, two: 2222
           }
         }
+        await group.next()
+        print("--")
 
+        TL.$two.withValue(2) {
+          TL.$two.withValue(22) {
+              TL.$two.withValue(2222) {
+                group.addTask { // will have to copy the `2222`
+                  print("Survived, one: \(TL.one) @ \(#fileID):\(#line)") // CHECK: Survived, one: 1111
+                  print("Survived, two: \(TL.two) @ \(#fileID):\(#line)") // CHECK: Survived, two: 2222
+                }
+              }
+          }
+        }
         await group.next()
         print("--")
 
         TL.$two.withValue(2) {
           TL.$two.withValue(22) {
-            TL.$two.withValue(2222) {
-              group.addTask { // will have to copy the `2222`
-                print("Survived, one: \(TL.one) @ \(#fileID):\(#line)") // CHECK: Survived, one: 1111
-                print("Survived, two: \(TL.two) @ \(#fileID):\(#line)") // CHECK: Survived, two: 2222
+            TL.$three.withValue(33) {
+              TL.$two.withValue(2222) {
+                group.addTask { // will have to copy the `2222`
+                  print("Survived, one: \(TL.one) @ \(#fileID):\(#line)") // CHECK: Survived, one: 1111
+                  print("Survived, two: \(TL.two) @ \(#fileID):\(#line)") // CHECK: Survived, two: 2222
+
+                  TL.$three.withValue(3333) {
+                    print("Survived, one: \(TL.one) @ \(#fileID):\(#line)") // CHECK: Survived, one: 1111
+                    print("Survived, two: \(TL.two) @ \(#fileID):\(#line)") // CHECK: Survived, two: 2222
+                    print("Survived, three: \(TL.three) @ \(#fileID):\(#line)") // CHECK: Survived, three: 3333
+                  }
+                }
               }
             }
           }