[Concurrency] Implement only-when-necessary defensive TL copying in in task groups

Author: ktoso

include/swift/ABI/TaskLocal.h

@@ -125,6 +125,25 @@ class TaskLocal {
       return reinterpret_cast<Item *>(next & ~statusMask);
     }
 
+    void relinkNext(Item* nextOverride) {
+      fprintf(stderr, "[%s:%d](%s) this item = %p\n", __FILE_NAME__, __LINE__, __FUNCTION__, this);
+      fprintf(stderr, "[%s:%d](%s) NEXT = %p\n", __FILE_NAME__, __LINE__, __FUNCTION__, getNext());
+      fprintf(stderr, "[%s:%d](%s) storage = %p\n", __FILE_NAME__, __LINE__, __FUNCTION__, getStoragePtr());
+      fprintf(stderr, "[%s:%d](%s) next kind = %d\n", __FILE_NAME__, __LINE__, __FUNCTION__, getNextLinkType());
+      assert(!getNext() &&
+               "Can only relink task local item that was not pointing at anything yet");
+      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!");
+
+      next = reinterpret_cast<uintptr_t>(nextOverride) |
+             static_cast<uintptr_t>((nextOverride->isNextLinkPointer()
+                                         ? NextLinkType::IsNextCreatedInTaskGroupBody
+                                         : NextLinkType::IsParent));
+    }
+
     NextLinkType getNextLinkType() const {
       return static_cast<NextLinkType>(next & statusMask);
     }
@@ -134,7 +153,7 @@ class TaskLocal {
              NextLinkType::IsNext;
     }
 
-    bool IsNextCreatedInTaskGroupBody() const {
+    bool isNextCreatedInTaskGroupBody() const {
       return static_cast<NextLinkType>(next & statusMask) ==
              NextLinkType::IsNextCreatedInTaskGroupBody;
     }
@@ -151,7 +170,7 @@ class TaskLocal {
         reinterpret_cast<char *>(this) + storageOffset(valueType));
     }
 
-    void copyTo(AsyncTask *task);
+    TaskLocal::Item* copyTo(AsyncTask *task);
 
     /// Compute the offset of the storage from the base of the item.
     static size_t storageOffset(const Metadata *valueType) {
@@ -238,6 +257,8 @@ 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);
+
     /// 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

@@ -663,11 +663,15 @@ void swift_task_localValuePop();
 /// Its Swift signature is
 ///
 /// \code
-/// func _taskLocalValueGet<Key>(AsyncTask* task)
+/// func swift_task_localsCopyTo<Key>(AsyncTask* task)
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_localsCopyTo(AsyncTask* target);
 
+// FIXME no need for runtime API
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_localsOnlyCurrentCopyTo(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,6 +378,12 @@ OVERRIDE_TASK_LOCAL(task_localsCopyTo, void,
                     (AsyncTask *target),
                     (target))
 
+OVERRIDE_TASK_LOCAL(task_localsOnlyCurrentCopyTo, 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

@@ -732,24 +732,25 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
   AsyncTask *currentTask = swift_task_getCurrent();
   AsyncTask *parent = jobFlags.task_isChildTask() ? currentTask : nullptr;
 
+  /// FIXME: this is the fail path
   if (group) {
     assert(parent && "a task created in a group must be a child task");
-
-    // Prevent task-local misuse;
-    // We must not allow an addTask {} wrapped immediately with a withValue {}
-    auto ParentLocal = parent->_private().Local;
-    if (auto taskLocalHeadLinkType = ParentLocal.peekHeadLinkType()) {
-      if (taskLocalHeadLinkType ==
-          swift::TaskLocal::NextLinkType::IsNextCreatedInTaskGroupBody) {
-#if !SWIFT_CONCURRENCY_EMBEDDED
-        swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroup(
-            nullptr, 0, true, 0);
-#endif
-        // TODO(ktoso): If we were to keep this crash mode; offer a better failure for embedded swift
-        abort();
-      }
-    }
-
+//
+//    // Prevent task-local misuse;
+//    // We must not allow an addTask {} wrapped immediately with a withValue {}
+//    auto ParentLocal = parent->_private().Local;
+//    if (auto taskLocalHeadLinkType = ParentLocal.peekHeadLinkType()) {
+//      if (taskLocalHeadLinkType ==
+//          swift::TaskLocal::NextLinkType::IsNextCreatedInTaskGroupBody) {
+//#if !SWIFT_CONCURRENCY_EMBEDDED
+//        swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroup(
+//            nullptr, 0, true, 0);
+//#endif
+//        // TODO(ktoso): If we were to keep this crash mode; offer a better failure for embedded swift
+//        abort();
+//      }
+//    }
+//
     // Add to the task group, if requested.
     if (taskCreateFlags.addPendingGroupTaskUnconditionally()) {
       assert(group && "Missing group");
@@ -998,12 +999,44 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
     // In a task group we would not have allowed the `add` to create a child anymore,
     // however better safe than sorry and `async let` are not expressed as task groups,
     // so they may have been spawned in any case still.
-    if (swift_task_isCancelled(parent) ||
-        (group && group->isCancelled()))
+    if ((group && group->isCancelled()) || swift_task_isCancelled(parent))
       swift_task_cancel(task);
 
-    // Initialize task locals with a link to the parent task.
-    task->_private().Local.initializeLinkParent(task, parent);
+    // Initialize task locals storage
+    bool taskLocalStorageInitialized = false;
+
+    // Inside a task group, we may have to perform some defensive copying,
+    // check if doing so is necessary, and initialize storage using partial
+    // defensive copies if necessary.
+    if (group) {
+      assert(parent && "a task created in a group must be a child task");
+      // We are a child task in a task group; and it may happen that we are calling
+      // addTask specifically in such shape:
+      //
+      //     $local.withValue(theValue) { addTask {} }
+      //
+      // If this is the case, we MUST copy `theValue` (and any other such directly
+      // wrapping the addTask value bindings), because those values will be popped
+      // when withValue returns - breaking our structured concurrency guarantees
+      // that we rely on for the "link directly to parent's task local Item".
+      //
+      // Values set outside the task group are not subject to this problem, as
+      // their structural lifetime guarantee is upheld by the group scope
+      // out-living any addTask created tasks.
+      auto ParentLocal = parent->_private().Local;
+      if (auto taskLocalHeadLinkType = ParentLocal.peekHeadLinkType()) {
+        if (taskLocalHeadLinkType ==
+            swift::TaskLocal::NextLinkType::IsNextCreatedInTaskGroupBody) {
+          swift_task_localsOnlyCurrentCopyTo(task);
+          taskLocalStorageInitialized = true;
+        }
+      }
+    }
+
+    if (!taskLocalStorageInitialized) {
+      // just initialize the storage normally
+      task->_private().Local.initializeLinkParent(task, parent);
+    }
   }
 
   // Configure the initial context.
@@ -1065,7 +1098,7 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
 #endif
     swift_retain(task);
     task->flagAsAndEnqueueOnExecutor(
-        serialExecutor); // FIXME: pass the task executor explicitly?
+        serialExecutor);
   }
 
   return {task, initialContext};