[6.0][Concurrency] Defensive copying of tasks locals in TaskGroups (rather than crashing)

include/swift/ABI/TaskLocal.h

@@ -44,6 +44,14 @@ class TaskLocal {
     /// lookups by skipping empty parent tasks during get(), and explained
     /// in depth in `createParentLink`.
     IsParent = 0b01,
+    /// 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,
   };
 
   class Item {
@@ -101,20 +109,55 @@ class TaskLocal {
     /// the Item linked list into the appropriate parent.
     static Item *createParentLink(AsyncTask *task, AsyncTask *parent);
 
+    static Item *createLink(AsyncTask *task,
+                            const HeapObject *key,
+                            const Metadata *valueType,
+                            bool inTaskGroupBody);
+
     static Item *createLink(AsyncTask *task,
                             const HeapObject *key,
                             const Metadata *valueType);
 
+    static Item *createLinkInTaskGroup(
+        AsyncTask *task,
+        const HeapObject *key,
+        const Metadata *valueType);
+
     void destroy(AsyncTask *task);
 
     Item *getNext() {
       return reinterpret_cast<Item *>(next & ~statusMask);
     }
 
+    void relinkTaskGroupLocalHeadToSafeNext(Item* nextOverride) {
+      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);
     }
 
+    bool isNextLinkPointer() const {
+      return static_cast<NextLinkType>(next & statusMask) ==
+             NextLinkType::IsNext;
+    }
+
+    bool isNextLinkPointerCreatedInTaskGroupBody() const {
+      return static_cast<NextLinkType>(next & statusMask) ==
+             NextLinkType::IsNextCreatedInTaskGroupBody;
+    }
+
     /// Item does not contain any actual value, and is only used to point at
     /// a specific parent item.
     bool isEmpty() const {
@@ -127,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) {
@@ -136,9 +179,9 @@ class TaskLocal {
       if (valueType) {
         size_t alignment = valueType->vw_alignment();
         return (offset + alignment - 1) & ~(alignment - 1);
-      } else {
-        return offset;
       }
+
+      return offset;
     }
 
     /// Determine the size of the item given a particular value type.
@@ -187,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,
@@ -200,6 +247,9 @@ class TaskLocal {
     /// can be safely disposed of.
     bool popValue(AsyncTask *task);
 
+    /// Peek at the head item and get its type.
+    std::optional<NextLinkType> peekHeadLinkType() const;
+
     /// Copy all task-local bindings to the target task.
     ///
     /// The new bindings allocate their own items and can out-live the current task.
@@ -212,6 +262,10 @@ class TaskLocal {
     /// "pop" of the `B` value - it was spawned from a scope where only B was observable.
     void copyTo(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.
     ///
     /// Items owned by a parent task are left untouched, since we do not own them.

include/swift/Runtime/Concurrency.h

@@ -663,7 +663,7 @@ 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);

stdlib/public/Concurrency/Actor.cpp

@@ -357,8 +357,6 @@ const char *__swift_runtime_env_useLegacyNonCrashingExecutorChecks() {
 #endif
 }
 
-#pragma clang diagnostic push
-#pragma ide diagnostic ignored "ConstantConditionsOC"
 // Done this way because of the interaction with the initial value of
 // 'unexpectedExecutorLogLevel'
 bool swift_bincompat_useLegacyNonCrashingExecutorChecks() {
@@ -378,7 +376,6 @@ bool swift_bincompat_useLegacyNonCrashingExecutorChecks() {
 
   return legacyMode;
 }
-#pragma clang diagnostic pop
 
 // Check override of executor checking mode.
 static void checkIsCurrentExecutorMode(void *context) {

stdlib/public/Concurrency/Task.cpp

@@ -729,17 +729,16 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
     assert(initialContextSize >= sizeof(FutureAsyncContext));
   }
 
-  // Add to the task group, if requested.
-  if (taskCreateFlags.addPendingGroupTaskUnconditionally()) {
-    assert(group && "Missing group");
-    swift_taskGroup_addPending(group, /*unconditionally=*/true);
-  }
-
-  AsyncTask *parent = nullptr;
   AsyncTask *currentTask = swift_task_getCurrent();
-  if (jobFlags.task_isChildTask()) {
-    parent = currentTask;
-    assert(parent != nullptr && "creating a child task with no active task");
+  AsyncTask *parent = jobFlags.task_isChildTask() ? currentTask : nullptr;
+
+  if (group) {
+    assert(parent && "a task created in a group must be a child task");
+    // Add to the task group, if requested.
+    if (taskCreateFlags.addPendingGroupTaskUnconditionally()) {
+      assert(group && "Missing group");
+      swift_taskGroup_addPending(group, /*unconditionally=*/true);
+    }
   }
 
   // Start with user specified priority at creation time (if any)
@@ -983,12 +982,51 @@ 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 we were going to copy ALL values anyway, we don't need to
+      // perform this defensive partial copying. In practice, we currently
+      // do not have child tasks which force copying, but we could.
+      assert(!taskCreateFlags.copyTaskLocals() &&
+             "Currently we don't have child tasks which force copying task "
+             "locals; unexpected attempt to combine the two!");
+
+      if (auto taskLocalHeadLinkType = ParentLocal.peekHeadLinkType()) {
+        if (taskLocalHeadLinkType ==
+            swift::TaskLocal::NextLinkType::IsNextCreatedInTaskGroupBody) {
+          ParentLocal.copyToOnlyOnlyFromCurrentGroup(task);
+          taskLocalStorageInitialized = true;
+        }
+      }
+    }
+
+    if (!taskLocalStorageInitialized) {
+      // just initialize the storage normally
+      task->_private().Local.initializeLinkParent(task, parent);
+    }
   }
 
   // Configure the initial context.
@@ -1050,7 +1088,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};