Merge pull request #80096 from mikeash/task-group-cancellation-lock-fix

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

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 ******************************/
 /**************************************************************************/