[Concurrency] Fix a race when using cancelAll on a task group concurrently with child tasks being removed.

_swift_taskGroup_cancelAllChildren relies on there being no concurrent modification when called from the owning task, but this is not guaranteed.

Rearrange things to always take the owning task's status record lock when walking the group's children. Split _swift_taskGroup_cancelAllChildren into two functions, one which assumes/requires the lock is already held, and one which acquires the lock. We don't have the owning task in this case, but we can either get it from the current task, or by looking at the parent of the child task we're working on.

rdar://147172991

Author: mikeash

include/swift/ABI/TaskStatus.h

@@ -133,10 +133,13 @@ class ChildTaskStatusRecord : public TaskStatusRecord {
 /// Group child tasks DO NOT have their own `ChildTaskStatusRecord` entries,
 /// and are only tracked by their respective `TaskGroupTaskStatusRecord`.
 class TaskGroupTaskStatusRecord : public TaskStatusRecord {
-public:
-  AsyncTask *FirstChild;
+  // FirstChild may be read concurrently to check for the presence of children,
+  // so it needs to be atomic. The pointer is never dereferenced in that case,
+  // so we can universally use memory_order_relaxed on it.
+  std::atomic<AsyncTask *> FirstChild;
   AsyncTask *LastChild;
 
+public:
   TaskGroupTaskStatusRecord()
       : TaskStatusRecord(TaskStatusRecordKind::TaskGroup),
         FirstChild(nullptr),
@@ -155,7 +158,9 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
   /// Return the first child linked by this record.  This may be null;
   /// if not, it (and all of its successors) are guaranteed to satisfy
   /// `isChildTask()`.
-  AsyncTask *getFirstChild() const { return FirstChild; }
+  AsyncTask *getFirstChild() const {
+    return FirstChild.load(std::memory_order_relaxed);
+  }
 
   /// Attach the passed in `child` task to this group.
   void attachChild(AsyncTask *child) {
@@ -165,9 +170,9 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
     auto oldLastChild = LastChild;
     LastChild = child;
 
-    if (!FirstChild) {
+    if (!getFirstChild()) {
       // This is the first child we ever attach, so store it as FirstChild.
-      FirstChild = child;
+      FirstChild.store(child, std::memory_order_relaxed);
       return;
     }
 
@@ -176,15 +181,18 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
 
   void detachChild(AsyncTask *child) {
     assert(child && "cannot remove a null child from group");
-    if (FirstChild == child) {
-      FirstChild = getNextChildTask(child);
-      if (FirstChild == nullptr) {
+
+    AsyncTask *prev = getFirstChild();
+
+    if (prev == child) {
+      AsyncTask *next = getNextChildTask(child);
+      FirstChild.store(next, std::memory_order_relaxed);
+      if (next == nullptr) {
         LastChild = nullptr;
       }
       return;
     }
 
-    AsyncTask *prev = FirstChild;
     // Remove the child from the linked list, i.e.:
     //     prev -> afterPrev -> afterChild
     //                 ==

stdlib/public/Concurrency/TaskGroup.cpp

@@ -474,9 +474,8 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
   bool statusCancel();
 
   /// Cancel the group and all of its child tasks recursively.
-  /// This also sets
-  bool cancelAll();
-
+  /// This also sets the cancelled bit in the group status.
+  bool cancelAll(AsyncTask *task);
 };
 
 #if !SWIFT_CONCURRENCY_EMBEDDED
@@ -1378,7 +1377,8 @@ void DiscardingTaskGroup::offer(AsyncTask *completedTask, AsyncContext *context)
     // Discarding results mode immediately treats a child failure as group cancellation.
     // "All for one, one for all!" - any task failing must cause the group and all sibling tasks to be cancelled,
     // such that the discarding group can exit as soon as possible.
-    cancelAll();
+    auto parent = completedTask->childFragment()->getParent();
+    cancelAll(parent);
 
     if (afterComplete.hasWaitingTask() && afterComplete.pendingTasks(this) == 0) {
       // We grab the waiting task while holding the group lock, because this
@@ -2097,10 +2097,12 @@ static bool swift_taskGroup_isCancelledImpl(TaskGroup *group) {
 
 SWIFT_CC(swift)
 static void swift_taskGroup_cancelAllImpl(TaskGroup *group) {
-  asBaseImpl(group)->cancelAll();
+  // TaskGroup is not a Sendable type, so this can only be called from the
+  // owning task.
+  asBaseImpl(group)->cancelAll(swift_task_getCurrent());
 }
 
-bool TaskGroupBase::cancelAll() {
+bool TaskGroupBase::cancelAll(AsyncTask *owningTask) {
   SWIFT_TASK_DEBUG_LOG("cancel all tasks in group = %p", this);
 
   // Flag the task group itself as cancelled.  If this was already
@@ -2114,8 +2116,8 @@ bool TaskGroupBase::cancelAll() {
 
   // Cancel all the child tasks.  TaskGroup is not a Sendable type,
   // so cancelAll() can only be called from the owning task.  This
-  // satisfies the precondition on cancelAllChildren().
-  _swift_taskGroup_cancelAllChildren(asAbstract(this));
+  // satisfies the precondition on cancelAllChildren_unlocked().
+  _swift_taskGroup_cancelAllChildren_unlocked(asAbstract(this), owningTask);
 
   return true;
 }
@@ -2124,22 +2126,8 @@ SWIFT_CC(swift)
 static void swift_task_cancel_group_child_tasksImpl(TaskGroup *group) {
   // TaskGroup is not a Sendable type, and so this operation (which is not
   // currently exposed in the API) can only be called from the owning
-  // task.  This satisfies the precondition on cancelAllChildren().
-  _swift_taskGroup_cancelAllChildren(group);
-}
-
-/// Cancel all the children of the given task group.
-///
-/// The caller must guarantee that this is either called from the
-/// owning task of the task group or while holding the owning task's
-/// status record lock.
-void swift::_swift_taskGroup_cancelAllChildren(TaskGroup *group) {
-  // Because only the owning task of the task group can modify the
-  // child list of a task group status record, and it can only do so
-  // while holding the owning task's status record lock, we do not need
-  // any additional synchronization within this function.
-  for (auto childTask: group->getTaskRecord()->children())
-    swift_task_cancel(childTask);
+  // task.  This satisfies the precondition on cancelAllChildren_unlocked().
+  _swift_taskGroup_cancelAllChildren_unlocked(group, swift_task_getCurrent());
 }
 
 // =============================================================================

stdlib/public/Concurrency/TaskPrivate.h

@@ -90,11 +90,16 @@ AsyncTask *_swift_task_setCurrent(AsyncTask *newTask);
 
 /// Cancel all the child tasks that belong to `group`.
 ///
-/// The caller must guarantee that this is either called from the
-/// owning task of the task group or while holding the owning task's
-/// status record lock.
+/// The caller must guarantee that this is called while holding the owning
+/// task's status record lock.
 void _swift_taskGroup_cancelAllChildren(TaskGroup *group);
 
+/// Cancel all the child tasks that belong to `group`.
+///
+/// The caller must guarantee that this is called from the owning task.
+void _swift_taskGroup_cancelAllChildren_unlocked(TaskGroup *group,
+                                                 AsyncTask *owningTask);
+
 /// Remove the given task from the given task group.
 ///
 /// This is an internal API; clients outside of the TaskGroup implementation

stdlib/public/Concurrency/TaskStatus.cpp

@@ -943,6 +943,35 @@ void swift::_swift_taskGroup_detachChild(TaskGroup *group,
   });
 }
 
+/// Cancel all the child tasks that belong to `group`.
+///
+/// The caller must guarantee that this is called while holding the owning
+/// task's status record lock.
+void swift::_swift_taskGroup_cancelAllChildren(TaskGroup *group) {
+  // Because only the owning task of the task group can modify the
+  // child list of a task group status record, and it can only do so
+  // while holding the owning task's status record lock, we do not need
+  // any additional synchronization within this function.
+  for (auto childTask : group->getTaskRecord()->children())
+    swift_task_cancel(childTask);
+}
+
+/// Cancel all the child tasks that belong to `group`.
+///
+/// The caller must guarantee that this is called from the owning task.
+void swift::_swift_taskGroup_cancelAllChildren_unlocked(TaskGroup *group,
+                                                        AsyncTask *owningTask) {
+  // Early out. If there are no children, there's nothing to do. We can safely
+  // check this without locking, since this can only be concurrently mutated
+  // from a child task. If there are no children then no more can be added.
+  if (!group->getTaskRecord()->getFirstChild())
+    return;
+
+  withStatusRecordLock(owningTask, [&group](ActiveTaskStatus status) {
+    _swift_taskGroup_cancelAllChildren(group);
+  });
+}
+
 /**************************************************************************/
 /****************************** CANCELLATION ******************************/
 /**************************************************************************/