Merge pull request #40520 from apple/rokhinip/synchronize-group-child-creation-and-cancellation

Synchronize adding a child task to a TaskGroup with parent task's TaskStatusRecord lock

Author: rokhinip

include/swift/ABI/TaskGroup.h

@@ -42,6 +42,10 @@ class alignas(Alignment_TaskGroup) TaskGroup {
 
   /// Checks the cancellation status of the group.
   bool isCancelled();
+
+  // Add a child task to the group. Always called with the status record lock of
+  // the parent task held
+  void addChildTask(AsyncTask *task);
 };
 
 } // end namespace swift

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -241,7 +241,7 @@ OVERRIDE_TASK_GROUP(taskGroup_initialize, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                     swift::, (TaskGroup *group, const Metadata *T), (group, T))
 
-OVERRIDE_TASK_GROUP(taskGroup_attachChild, void,
+OVERRIDE_TASK_STATUS(taskGroup_attachChild, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                     swift::, (TaskGroup *group, AsyncTask *child),
                     (group, child))

stdlib/public/Concurrency/TaskGroup.cpp

@@ -472,15 +472,13 @@ static void swift_taskGroup_initializeImpl(TaskGroup *group, const Metadata *T)
 // =============================================================================
 // ==== add / attachChild ------------------------------------------------------
 
-SWIFT_CC(swift)
-static void swift_taskGroup_attachChildImpl(TaskGroup *group,
-                                            AsyncTask *child) {
+void TaskGroup::addChildTask(AsyncTask *child) {
   SWIFT_TASK_DEBUG_LOG("attach child task = %p to group = %p", child, group);
 
   // The counterpart of this (detachChild) is performed by the group itself,
   // when it offers the completed (child) task's value to a waiting task -
   // during the implementation of `await group.next()`.
-  auto groupRecord = asImpl(group)->getTaskRecord();
+  auto groupRecord = asImpl(this)->getTaskRecord();
   groupRecord->attachChild(child);
 }
 

stdlib/public/Concurrency/TaskStatus.cpp

@@ -468,6 +468,26 @@ swift_task_detachChildImpl(ChildTaskStatusRecord *record) {
   swift_task_removeStatusRecord(record);
 }
 
+SWIFT_CC(swift)
+static void swift_taskGroup_attachChildImpl(TaskGroup *group,
+                                            AsyncTask *child) {
+
+  // We are always called from the context of the parent
+  //
+  // Acquire the status record lock of parent - we want to synchronize with
+  // concurrent cancellation or escalation as we're adding new tasks to the
+  // group.
+
+  Optional<StatusRecordLockRecord> recordLockRecord;
+  auto parent = swift_task_getCurrent();
+  auto oldStatus =
+      acquireStatusRecordLock(parent, recordLockRecord, LockContext::OnTask);
+  group->addChildTask(child);
+
+  // Release the status record lock, restoring exactly the old status.
+  releaseStatusRecordLock(parent, oldStatus, recordLockRecord);
+}
+
 /****************************** CANCELLATION ******************************/
 /**************************************************************************/