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

Diff for: include/swift/ABI/TaskStatus.h

@@ -35,7 +35,7 @@ namespace swift {
 /// access by a cancelling thread.  In particular, the chain of
 /// status records must not be disturbed.  When the task leaves
 /// the scope that requires the status record, the record can
-/// be unregistered from the task with `swift_task_removeStatusRecord`,
+/// be unregistered from the task with `removeStatusRecord`,
 /// at which point the memory can be returned to the system.
 class TaskStatusRecord {
 public:
@@ -256,7 +256,7 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
 ///
 /// The end of any call to the function will be ordered before the
 /// end of a call to unregister this record from the task.  That is,
-/// code may call `swift_task_removeStatusRecord` and freely
+/// code may call `removeStatusRecord` and freely
 /// assume after it returns that this function will not be
 /// subsequently used.
 class CancellationNotificationStatusRecord : public TaskStatusRecord {
@@ -284,7 +284,7 @@ class CancellationNotificationStatusRecord : public TaskStatusRecord {
 ///
 /// The end of any call to the function will be ordered before the
 /// end of a call to unregister this record from the task.  That is,
-/// code may call `swift_task_removeStatusRecord` and freely
+/// code may call `removeStatusRecord` and freely
 /// assume after it returns that this function will not be
 /// subsequently used.
 class EscalationNotificationStatusRecord : public TaskStatusRecord {

Diff for: 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> }`.
 ///

Diff for: 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::, , )

Diff for: stdlib/public/Concurrency/AsyncLet.cpp

@@ -142,8 +142,14 @@ 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 =
+      addStatusRecord(record, [&](ActiveTaskStatus parentStatus) {
+        updateNewChildWithParentAndGroupState(task, parentStatus, NULL);
+        return true;
+      });
+  assert(addedRecord);
 }
 
 // =============================================================================
@@ -309,7 +315,7 @@ static void swift_asyncLet_endImpl(AsyncLet *alet) {
 
   // Remove the child record from the parent task
   auto record = asImpl(alet)->getTaskRecord();
-  swift_task_removeStatusRecord(record);
+  removeStatusRecord(record);
 
   // TODO: we need to implicitly await either before the end or here somehow.
 
@@ -337,7 +343,7 @@ static void asyncLet_finish_after_task_completion(SWIFT_ASYNC_CONTEXT AsyncConte
 
   // Remove the child record from the parent task
   auto record = asImpl(alet)->getTaskRecord();
-  swift_task_removeStatusRecord(record);
+  removeStatusRecord(record);
 
   // and finally, release the task and destroy the async-let
   assert(swift_task_getCurrent() && "async-let must have a parent task");

Diff for: stdlib/public/Concurrency/Task.cpp

@@ -1056,19 +1056,27 @@ swift_task_addCancellationHandlerImpl(
   auto *record = new (allocation)
       CancellationNotificationStatusRecord(unsigned_handler, context);
 
-  if (swift_task_addStatusRecord(record))
-    return record;
+  bool fireHandlerNow = false;
 
-  // 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();
+  addStatusRecord(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;
 }
 
 SWIFT_CC(swift)
 static void swift_task_removeCancellationHandlerImpl(
     CancellationNotificationStatusRecord *record) {
-  swift_task_removeStatusRecord(record);
+  removeStatusRecord(record);
   swift_task_dealloc(record);
 }
 

Diff for: stdlib/public/Concurrency/TaskGroup.cpp

@@ -474,11 +474,14 @@ 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();
+  addStatusRecord(record, [&](ActiveTaskStatus parentStatus) {
+    // If the task has already been cancelled, reflect that immediately in
+    // the group's status.
+    if (parentStatus.isCancelled()) {
+      impl->statusCancel();
+    }
+    return true;
+  });
 }
 
 // =============================================================================
@@ -505,7 +508,7 @@ void TaskGroupImpl::destroy() {
   SWIFT_TASK_DEBUG_LOG("destroying task group = %p", this);
 
   // First, remove the group from the task and deallocate the record
-  swift_task_removeStatusRecord(getTaskRecord());
+  removeStatusRecord(getTaskRecord());
 
   // No need to drain our queue here, as by the time we call destroy,
   // all tasks inside the group must have been awaited on already.

Diff for: stdlib/public/Concurrency/TaskPrivate.h

@@ -238,12 +238,23 @@ class alignas(sizeof(void*) * 2) ActiveTaskStatus {
   bool isStoredPriorityEscalated() const {
     return Flags & IsEscalated;
   }
+
+  /// Creates a new active task status for a task with the specified priority
+  /// and masks away any existing priority related flags on the task status. All
+  /// other flags about the task are unmodified. This is only safe to use to
+  /// generate an initial task status for a new task that is not yet running.
+  ActiveTaskStatus withNewPriority(JobPriority priority) const {
+    return ActiveTaskStatus(Record,
+                            (Flags & ~PriorityMask) | uintptr_t(priority));
+  }
+
   ActiveTaskStatus withEscalatedPriority(JobPriority priority) const {
     assert(priority > getStoredPriority());
     return ActiveTaskStatus(Record,
                             (Flags & ~PriorityMask)
                                | IsEscalated | uintptr_t(priority));
   }
+
   ActiveTaskStatus withoutStoredPriorityEscalation() const {
     assert(isStoredPriorityEscalated());
     return ActiveTaskStatus(Record, Flags & ~IsEscalated);
@@ -438,6 +449,39 @@ 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_CC(swift)
+bool 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_CC(swift)
+bool addStatusRecord(TaskStatusRecord *record,
+                     llvm::function_ref<bool(ActiveTaskStatus)> testAddRecord);
+
+/// A helper function 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_CC(swift)
+void updateNewChildWithParentAndGroupState(AsyncTask *child,
+                                           ActiveTaskStatus parentStatus,
+                                           TaskGroup *group);
+
 } // end namespace swift
 
 #endif