Change the logic for adding new task status records to a task

This change has two parts to it:

1. Add in a new interface (addStatusRecordWithChecks) for adding task
status records that also takes in a function ref. This function ref will
be used to evaluate if current state of the parent task has any changes
that need to be propagated to the child task that has been created.

This is necessary to prevent the following race between task creation
and concurrent cancellation and escalation:

a. Parent task create child task. It does lazy relaxed loads on its own
   state while doing so and propagates this state to the child.
b. Child task is created but has not been attached to the parent
   task/task group.
c. Parent task gets cancelled by another thread.
d. Child task gets linked into the parent’s task status records but no
   reevaluation has happened to account for changes that might have happened to
   the parent after (a).

2. Move status record management functions from the
Runtime/Concurrency.h to TaskPrivate.h. Remove any corresponding
overrides that are no longer needed. Remove unused tryAddStatusRecord
method whose functionality is provided by addStatusRecordWithChecks.

Radar-Id: rdar://problem/86347801

Author: rokhinip

include/swift/Runtime/Concurrency.h

@@ -17,9 +17,9 @@
 #ifndef SWIFT_RUNTIME_CONCURRENCY_H
 #define SWIFT_RUNTIME_CONCURRENCY_H
 
+#include "swift/ABI/AsyncLet.h"
 #include "swift/ABI/Task.h"
 #include "swift/ABI/TaskGroup.h"
-#include "swift/ABI/AsyncLet.h"
 #include "swift/ABI/TaskStatus.h"
 
 #pragma clang diagnostic push
@@ -466,40 +466,6 @@ void swift_asyncLet_consume_throwing(SWIFT_ASYNC_CONTEXT AsyncContext *,
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_taskGroup_hasTaskGroupRecord();
 
-/// Add a status record to a task.  The record should not be
-/// modified while it is registered with a task.
-///
-/// This must be called synchronously with the task.
-///
-/// If the task is already cancelled, returns `false` but still adds
-/// the status record.
-SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-bool swift_task_addStatusRecord(TaskStatusRecord *record);
-
-/// Add a status record to a task if the task has not already
-/// been cancelled.   The record should not be modified while it is
-/// registered with a task.
-///
-/// This must be called synchronously with the task.
-///
-/// If the task is already cancelled, returns `false` and does not
-/// add the status record.
-SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-bool swift_task_tryAddStatusRecord(TaskStatusRecord *record);
-
-/// Remove a status record from a task.  After this call returns,
-/// the record's memory can be freely modified or deallocated.
-///
-/// This must be called synchronously with the task.  The record must
-/// be registered with the task or else this may crash.
-///
-/// The given record need not be the last record added to
-/// the task, but the operation may be less efficient if not.
-///
-/// Returns false if the task has been cancelled.
-SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-bool swift_task_removeStatusRecord(TaskStatusRecord *record);
-
 /// Signifies whether the current task is in the middle of executing the
 /// operation block of a `with(Throwing)TaskGroup(...) { <operation> }`.
 ///

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -308,18 +308,6 @@ OVERRIDE_TASK_LOCAL(task_localsCopyTo, void,
                     (AsyncTask *target),
                     (target))
 
-OVERRIDE_TASK_STATUS(task_addStatusRecord, bool,
-                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
-                     swift::, (TaskStatusRecord *newRecord), (newRecord))
-
-OVERRIDE_TASK_STATUS(task_tryAddStatusRecord, bool,
-                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
-                     swift::, (TaskStatusRecord *newRecord), (newRecord))
-
-OVERRIDE_TASK_STATUS(task_removeStatusRecord, bool,
-                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
-                     swift::, (TaskStatusRecord *record), (record))
-
 OVERRIDE_TASK_STATUS(task_hasTaskGroupStatusRecord, bool,
                      SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                      swift::, , )

stdlib/public/Concurrency/AsyncLet.cpp

@@ -142,8 +142,15 @@ void swift::asyncLet_addImpl(AsyncTask *task, AsyncLet *asyncLet,
   auto record = impl->getTaskRecord();
   assert(impl == record && "the async-let IS the task record");
 
-  // ok, now that the group actually is initialized: attach it to the task
-  swift_task_addStatusRecord(record);
+  // ok, now that the async let task actually is initialized: attach it to the
+  // current task
+  bool addedRecord = swift_task_addStatusRecordWithChecks(
+      record, [&](ActiveTaskStatus parentStatus) {
+        swift_task_updateNewChildWithParentAndGroupState(task, parentStatus,
+                                                         NULL);
+        return true;
+      });
+  assert(addedRecord);
 }
 
 // =============================================================================

stdlib/public/Concurrency/Task.cpp

@@ -1056,12 +1056,22 @@ swift_task_addCancellationHandlerImpl(
   auto *record = new (allocation)
       CancellationNotificationStatusRecord(unsigned_handler, context);
 
-  if (swift_task_addStatusRecord(record))
-    return record;
-
-  // else, the task was already cancelled, so while the record was added,
-  // we must run it immediately here since no other task will trigger it.
-  record->run();
+  bool fireHandlerNow = false;
+
+  swift_task_addStatusRecordWithChecks(record,
+                                       [&](ActiveTaskStatus parentStatus) {
+                                         if (parentStatus.isCancelled()) {
+                                           fireHandlerNow = true;
+                                           /* We don't fire the cancellation
+                                            * handler here since this function
+                                            * needs to be idempotent */
+                                         }
+                                         return true;
+                                       });
+
+  if (fireHandlerNow) {
+    record->run();
+  }
   return record;
 }
 

stdlib/public/Concurrency/TaskGroup.cpp

@@ -470,11 +470,13 @@ static void swift_taskGroup_initializeImpl(TaskGroup *group, const Metadata *T)
   assert(impl == record && "the group IS the task record");
 
   // ok, now that the group actually is initialized: attach it to the task
-  bool notCancelled = swift_task_addStatusRecord(record);
-
-  // If the task has already been cancelled, reflect that immediately in
-  // the group status.
-  if (!notCancelled) impl->statusCancel();
+  swift_task_addStatusRecordWithChecks(record,
+                                       [&](ActiveTaskStatus parentStatus) {
+                                         if (parentStatus.isCancelled()) {
+                                           impl->statusCancel();
+                                         }
+                                         return true;
+                                       });
 }
 
 // =============================================================================

stdlib/public/Concurrency/TaskPrivate.h

@@ -438,6 +438,40 @@ inline bool AsyncTask::localValuePop() {
   return _private().Local.popValue(this);
 }
 
+/*************** Methods for Status records manipulation ******************/
+
+/// Remove a status record from a task.  After this call returns,
+/// the record's memory can be freely modified or deallocated.
+///
+/// This must be called synchronously with the task.  The record must
+/// be registered with the task or else this may crash.
+///
+/// The given record need not be the last record added to
+/// the task, but the operation may be less efficient if not.
+///
+/// Returns false if the task has been cancelled.
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift) bool swift_task_removeStatusRecord(TaskStatusRecord *record);
+
+/// Add a status record to a task. This must be called synchronously with the
+/// task.
+///
+/// This function also takes in a function_ref which is given the task status of
+/// the task we're adding the record to, to determine if the current status of
+/// the task permits adding the status record. This function_ref may be called
+/// multiple times and must be idempotent.
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift)
+bool swift_task_addStatusRecordWithChecks(
+    TaskStatusRecord *record,
+    llvm::function_ref<bool(ActiveTaskStatus)> testAddRecord);
+
+/// A helper method for updating a new child task that is created with
+/// information from the parent or the group that it was going to be added to.
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift)
+void swift_task_updateNewChildWithParentAndGroupState(
+    AsyncTask *child, ActiveTaskStatus parentStatus, TaskGroup *group);
 } // end namespace swift
 
 #endif

stdlib/public/Concurrency/TaskStatus.cpp

@@ -256,11 +256,14 @@ static bool withStatusRecordLock(AsyncTask *task,
 /*************************** RECORD MANAGEMENT ****************************/
 /**************************************************************************/
 
+SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swift)
-static bool swift_task_addStatusRecordImpl(TaskStatusRecord *newRecord) {
-  auto task = swift_task_getCurrent();
+bool swift_task_addStatusRecordWithChecks(
+    TaskStatusRecord *newRecord,
+    llvm::function_ref<bool(ActiveTaskStatus status)> shouldAddRecord) {
 
-  // Load the current state.  We can use a relaxed load because we're
+  auto task = swift_task_getCurrent();
+  // Load the current state. We can use a relaxed load because we're
   // synchronous with the task.
   auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
 
@@ -273,53 +276,28 @@ static bool swift_task_addStatusRecordImpl(TaskStatusRecord *newRecord) {
     newRecord->resetParent(oldStatus.getInnermostRecord());
 
     // Set the record as the new innermost record.
-    // We have to use a release on success to make the initialization of
-    // the new record visible to the cancelling thread.
     ActiveTaskStatus newStatus = oldStatus.withInnermostRecord(newRecord);
-    if (task->_private().Status.compare_exchange_weak(oldStatus, newStatus,
-           /*success*/ std::memory_order_release,
-           /*failure*/ std::memory_order_relaxed))
-      return !oldStatus.isCancelled();
-  }
-}
-
-SWIFT_CC(swift)
-static bool swift_task_tryAddStatusRecordImpl(TaskStatusRecord *newRecord) {
-  auto task = swift_task_getCurrent();
-
-  // Load the current state.  We can use a relaxed load because we're
-  // synchronous with the task.
-  auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
 
-  while (true) {
-    // If the old info is already cancelled, do nothing.
-    if (oldStatus.isCancelled())
+    if (shouldAddRecord(newStatus)) {
+      // We have to use a release on success to make the initialization of
+      // the new record visible to the cancelling thread.
+      if (task->_private().Status.compare_exchange_weak(
+              oldStatus, newStatus,
+              /*success*/ std::memory_order_release,
+              /*failure*/ std::memory_order_relaxed)) {
+        return true;
+      } else {
+        /* Retry */
+      }
+    } else {
       return false;
-
-    // Wait for any active lock to be released.
-    if (oldStatus.isLocked()) {
-      waitForStatusRecordUnlock(task, oldStatus);
-
-      if (oldStatus.isCancelled())
-        return false;
     }
-
-    // Reset the parent of the new record.
-    newRecord->resetParent(oldStatus.getInnermostRecord());
-
-    // Set the record as the new innermost record.
-    // We have to use a release on success to make the initialization of
-    // the new record visible to the cancelling thread.
-    ActiveTaskStatus newStatus = oldStatus.withInnermostRecord(newRecord);
-    if (task->_private().Status.compare_exchange_weak(oldStatus, newStatus,
-           /*success*/ std::memory_order_release,
-           /*failure*/ std::memory_order_relaxed))
-      return true;
   }
 }
 
+SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swift)
-static bool swift_task_removeStatusRecordImpl(TaskStatusRecord *record) {
+bool swift_task_removeStatusRecord(TaskStatusRecord *record) {
   auto task = swift_task_getCurrent();
   SWIFT_TASK_DEBUG_LOG("remove status record = %p, from current task = %p",
                        record, task);
@@ -411,7 +389,14 @@ swift_task_attachChildImpl(AsyncTask *child) {
   auto record = new (allocation) swift::ChildTaskStatusRecord(child);
   SWIFT_TASK_DEBUG_LOG("attach child task = %p, record = %p, to current task = %p",
                        child, record, swift_task_getCurrent());
-  swift_task_addStatusRecord(record);
+
+  bool added_record = swift_task_addStatusRecordWithChecks(
+      record, [&](ActiveTaskStatus parentStatus) {
+        swift_task_updateNewChildWithParentAndGroupState(child, parentStatus,
+                                                         NULL);
+        return true;
+      });
+  assert(added_record);
   return record;
 }
 
@@ -421,6 +406,43 @@ swift_task_detachChildImpl(ChildTaskStatusRecord *record) {
   swift_task_removeStatusRecord(record);
 }
 
+/* Called in the path of linking a child into a parent/group synchronously with
+ * the parent task.
+ *
+ * When called to link a child into a parent directly, this does not hold the
+ * parent's task status record lock. When called to link a child into a task
+ * group, this holds the parent's task status record lock.
+ */
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift)
+void swift_task_updateNewChildWithParentAndGroupState(
+    AsyncTask *child, ActiveTaskStatus parentStatus, TaskGroup *group) {
+  /*
+   * We can take the fast path of just modifying the ActiveTaskStatus in the
+   * child task since we know that it won't have any task status records and
+   * cannot be accessed by anyone else since it hasn't been linked in yet.
+   * Avoids the extra logic in `swift_task_cancel` and `swift_task_escalate`
+   */
+  auto oldChildTaskStatus =
+      child->_private().Status.load(std::memory_order_relaxed);
+  assert(oldChildTaskStatus.getInnermostRecord() == NULL);
+
+  auto newChildTaskStatus = oldChildTaskStatus;
+
+  /* Parent task is cancelled or group the child task is part of (if any) is
+   * cancelled */
+  if (parentStatus.isCancelled() || (group && group->isCancelled())) {
+    newChildTaskStatus = newChildTaskStatus.withCancelled();
+  }
+
+  /* Parent task got escalated, make sure to propagate it to child. */
+  if (parentStatus.isStoredPriorityEscalated()) {
+    newChildTaskStatus = newChildTaskStatus.withEscalatedPriority(
+        parentStatus.getStoredPriority());
+  }
+  child->_private().Status.store(newChildTaskStatus, std::memory_order_relaxed);
+}
+
 SWIFT_CC(swift)
 static void swift_taskGroup_attachChildImpl(TaskGroup *group,
                                             AsyncTask *child) {
@@ -430,12 +452,25 @@ static void swift_taskGroup_attachChildImpl(TaskGroup *group,
   // 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.
-
   auto parent = swift_task_getCurrent();
   withStatusRecordLock(parent, LockContext::OnTask,
-                       [&](ActiveTaskStatus &status) {
-    group->addChildTask(child);
-  });
+                       [&](ActiveTaskStatus &parentStatus) {
+                         group->addChildTask(child);
+                         /*
+                          * After getting parent's status record lock, do some
+                          * sanity checks to see if parent task or group has
+                          * state changes that need to be propagated to the
+                          * child.
+                          *
+                          * This is the same logic that we would do if we were
+                          * adding a child task status record - see also
+                          * asyncLet_addImpl. Since we attach a child task to a
+                          * TaskGroupRecord instead, we synchronize on the
+                          * parent's task status and then update the child.
+                          */
+                         swift_task_updateNewChildWithParentAndGroupState(
+                             child, parentStatus, group);
+                       });
 }
 
 /****************************** CANCELLATION ******************************/

unittests/runtime/CMakeLists.txt

@@ -140,6 +140,8 @@ if(("${SWIFT_HOST_VARIANT_SDK}" STREQUAL "${SWIFT_PRIMARY_VARIANT_SDK}") AND
 
   target_include_directories(SwiftRuntimeTests BEFORE PRIVATE
     ${SWIFT_SOURCE_DIR}/stdlib/include)
+  target_include_directories(SwiftRuntimeTests BEFORE PRIVATE
+    ${SWIFT_SOURCE_DIR}/stdlib/public)
 
   # FIXME: cross-compile for all variants.
   target_link_libraries(SwiftRuntimeTests

unittests/runtime/CompatibilityOverrideConcurrency.cpp

@@ -227,18 +227,6 @@ TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_localsCopyTo) {
   swift_task_localsCopyTo(nullptr);
 }
 
-TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_addStatusRecord) {
-  swift_task_addStatusRecord(nullptr);
-}
-
-TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_tryAddStatusRecord) {
-  swift_task_tryAddStatusRecord(nullptr);
-}
-
-TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_removeStatusRecord) {
-  swift_task_removeStatusRecord(nullptr);
-}
-
 TEST_F(CompatibilityOverrideConcurrencyTest, task_hasTaskGroupStatusRecord) {
   swift_task_hasTaskGroupStatusRecord();
 }