[Concurrency] Don't add new task locals copy runtime func

Author: ktoso

include/swift/ABI/TaskLocal.h

@@ -129,11 +129,11 @@ class TaskLocal {
       return reinterpret_cast<Item *>(next & ~statusMask);
     }
 
-    void relinkNext(Item* nextOverride) {
+    void relinkTaskGroupLocalHeadToSafeNext(Item* nextOverride) {
       assert(!getNext() &&
                "Can only relink task local item that was not pointing at anything yet");
-      assert(nextOverride->isNextLinkPointer() ||
-             nextOverride->isParentPointer() &&
+      assert((nextOverride->isNextLinkPointer() ||
+              nextOverride->isParentPointer()) &&
                  "Currently relinking is only done within a task group to "
                  "avoid within-taskgroup next pointers; attempted to point at "
                  "task local declared within task group body though!");
@@ -230,6 +230,10 @@ class TaskLocal {
 
   public:
 
+    /// Get the "current" task local storage from either the passed in
+    /// task, or fall back to the *thread* local stored storage.
+    static Storage* getCurrent(AsyncTask *task);
+
     void initializeLinkParent(AsyncTask *task, AsyncTask *parent);
 
     void pushValue(AsyncTask *task,
@@ -258,7 +262,9 @@ class TaskLocal {
     /// "pop" of the `B` value - it was spawned from a scope where only B was observable.
     void copyTo(AsyncTask *target);
 
-    void copyToOnlyOnlyFromCurrent(AsyncTask *target);
+    // FIXME(concurrency): We currently copy from "all" task groups we encounter
+    // however in practice we only
+    void copyToOnlyOnlyFromCurrentGroup(AsyncTask *target);
 
     /// Destroy and deallocate all items stored by this specific task.
     ///

include/swift/Runtime/Concurrency.h

@@ -668,18 +668,6 @@ void swift_task_localValuePop();
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_localsCopyTo(AsyncTask* target);
 
-/// Some task local bindings must be copied defensively when a child task is
-/// created in a task group. See task creation (swift_task_create_common) for
-/// a detailed discussion how and when this is used.
-///
-/// Its Swift signature is
-///
-/// \code
-/// func swift_task_localsCopyToTaskGroupChildTaskDefensively<Key>(AsyncTask* task)
-/// \endcode
-SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_task_localsCopyToTaskGroupChildTaskDefensively(AsyncTask* target);
-
 /// Switch the current task to a new executor if we aren't already
 /// running on a compatible executor.
 ///

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -378,12 +378,6 @@ OVERRIDE_TASK_LOCAL(task_localsCopyTo, void,
                     (AsyncTask *target),
                     (target))
 
-OVERRIDE_TASK_LOCAL(task_localsCopyToTaskGroupChildTaskDefensively, void,
-                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
-                    swift::,
-                    (AsyncTask *target),
-                    (target))
-
 OVERRIDE_TASK_STATUS(task_hasTaskGroupStatusRecord, bool,
                      SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                      swift::, , )

stdlib/public/Concurrency/Task.cpp

@@ -1017,7 +1017,7 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
       if (auto taskLocalHeadLinkType = ParentLocal.peekHeadLinkType()) {
         if (taskLocalHeadLinkType ==
             swift::TaskLocal::NextLinkType::IsNextCreatedInTaskGroupBody) {
-          swift_task_localsCopyToTaskGroupChildTaskDefensively(task);
+          ParentLocal.copyToOnlyOnlyFromCurrentGroup(task);
           taskLocalStorageInitialized = true;
         }
       }

stdlib/public/Concurrency/TaskLocal.cpp

@@ -149,39 +149,25 @@ static void swift_task_localValuePopImpl() {
 
 SWIFT_CC(swift)
 static void swift_task_localsCopyToImpl(AsyncTask *target) {
-  TaskLocal::Storage *Local = nullptr;
-
-  if (AsyncTask *task = swift_task_getCurrent()) {
-    Local = &task->_private().Local;
-  } else if (auto *storage = FallbackTaskLocalStorage::get()) {
-    Local = storage;
-  } else {
-    // bail out, there are no values to copy
-    return;
+  if (auto Current = TaskLocal::Storage::getCurrent(swift_task_getCurrent())) {
+    Current->copyTo(target);
   }
-
-  Local->copyTo(target);
 }
 
-SWIFT_CC(swift)
-static void swift_task_localsCopyToTaskGroupChildTaskDefensivelyImpl(AsyncTask *target) {
-  TaskLocal::Storage *Local = nullptr;
+// =============================================================================
+// ==== Initialization ---------------------------------------------------------
 
-  if (AsyncTask *task = swift_task_getCurrent()) {
-    Local = &task->_private().Local;
+TaskLocal::Storage*
+TaskLocal::Storage::getCurrent(AsyncTask *current) {
+  if (current) {
+    return &current->_private().Local;
   } else if (auto *storage = FallbackTaskLocalStorage::get()) {
-    Local = storage;
-  } else {
-    // bail out, there are no values to copy
-    return;
+    return storage;
   }
 
-  Local->copyToOnlyOnlyFromCurrent(target);
+  return nullptr;
 }
 
-// =============================================================================
-// ==== Initialization ---------------------------------------------------------
-
 void TaskLocal::Storage::initializeLinkParent(AsyncTask* task,
                                               AsyncTask* parent) {
   assert(!head && "initial task local storage was already initialized");
@@ -501,7 +487,9 @@ void TaskLocal::Storage::copyTo(AsyncTask *target) {
   }
 }
 
-void TaskLocal::Storage::copyToOnlyOnlyFromCurrent(AsyncTask *target) {
+// TODO(concurrency): This can be optimized to copy only from the CURRENT group,
+//  but we need to detect this, e.g. by more flags in the items made from a group?
+void TaskLocal::Storage::copyToOnlyOnlyFromCurrentGroup(AsyncTask *target) {
   assert(target && "task must not be null when copying values into it");
   assert(!(target->_private().Local.head) &&
       "Task must not have any task-local values bound before copying into it");
@@ -525,7 +513,7 @@ void TaskLocal::Storage::copyToOnlyOnlyFromCurrent(AsyncTask *target) {
         // 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.
-        copiedHead->relinkNext(item);
+        copiedHead->relinkTaskGroupLocalHeadToSafeNext(item);
         break;
       }
 

test/Concurrency/Runtime/async_task_locals_in_task_group_may_need_to_copy.swift

@@ -13,6 +13,8 @@ enum TL {
   static var two: Int = 2
   @TaskLocal
   static var three: Int = 3
+  @TaskLocal
+  static var four: Int = 4
 }
 
 // ==== ------------------------------------------------------------------------
@@ -96,6 +98,7 @@ func test() async {
     print("--")
   }
 
+  // Multiple nested task groups *in the same task*
   await TL.$one.withValue(11) {
     await withTaskGroup(of: Void.self) { group in
       await TL.$three.withValue(33) {
@@ -109,10 +112,52 @@ func test() async {
           }
         }
       }
+      print("--")
+    }
+  }
+
+  // Multiple "nested" task groups, however in different child tasks
+  await withTaskGroup(of: Void.self) { outer in
+    // Child task inside  'outer':
+    await TL.$one.withValue(11) { // DOES NOT have to be copied by the 'inner.addTask' (!)
+      outer.addTask {
+        await TL.$two.withValue(2222) {
+          await withTaskGroup(of: Void.self) { inner in
+            TL.$three.withValue(3333) { // MUST be copied by 'inner.addTask'
+              inner.addTask {
+                print("Survived, one: \(TL.one) @ \(#fileID):\(#line)") // CHECK: Survived, one: 11
+                print("Survived, two: \(TL.two) @ \(#fileID):\(#line)") // CHECK: Survived, two: 2222
+                print("Survived, three: \(TL.three) @ \(#fileID):\(#line)") // CHECK: Survived, three: 3333
+              }
+            }
+          }
+        }
+        print("--")
+      }
+    }
+  }
 
-      print("Survived, done") // CHECK: Survived, done
+  // Multiple "safe" task locals around a withTaskGroup, check if relinking works ok then
+  await TL.$one.withValue(11) {
+    await TL.$two.withValue(22) {
+      await TL.$three.withValue(33) {
+        await withTaskGroup(of: Void.self) { group in
+          await TL.$four.withValue(44) { // must copy
+            group.addTask {
+              print("Survived, one: \(TL.one) @ \(#fileID):\(#line)") // CHECK: Survived, one: 11
+              print("Survived, two: \(TL.two) @ \(#fileID):\(#line)") // CHECK: Survived, two: 22
+              print("Survived, three: \(TL.three) @ \(#fileID):\(#line)") // CHECK: Survived, three: 33
+              print("Survived, four: \(TL.four) @ \(#fileID):\(#line)") // CHECK: Survived, four: 44
+            }
+          }
+        }
+      }
     }
+    print("--")
   }
+
+
+  print("Survived, done") // CHECK: Survived, done
 }
 
 @main struct Main {