Add support to addStatusRecord and removeStatusRecord so that the flags

field of an ActiveTaskStatus can also be modified while the
TaskStatusRecord list is being modified. Make the StatusRecordLock
reentrant.

Radar-Id: rdar://problem/88093007

Author: rokhinip

include/swift/ABI/TaskStatus.h

@@ -22,21 +22,26 @@
 
 #include "swift/ABI/MetadataValues.h"
 #include "swift/ABI/Task.h"
+#include "swift/Runtime/HeapObject.h"
 
 namespace swift {
 
 /// The abstract base class for all status records.
 ///
-/// TaskStatusRecords are typically allocated on the stack (possibly
-/// in the task context), partially initialized, and then atomically
-/// added to the task with `swift_task_addTaskStatusRecord`.  While
-/// registered with the task, a status record should only be
-/// modified in ways that respect the possibility of asynchronous
-/// 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 `removeStatusRecord`,
-/// at which point the memory can be returned to the system.
+/// TaskStatusRecords are typically allocated on the stack (possibly in the task
+/// context), partially initialized, and then atomically added to the task with
+/// `addStatusRecord`.
+///
+/// Status records currently are only added to a task by itself but they may be
+/// removed from the task by other threads, or accessed by other threads
+/// asynchronous to the task for other purposes.
+///
+/// As a result, while registered with the task, a status record should only be
+/// modified in ways that respect the possibility of asynchronous access by a
+/// different 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 `removeStatusRecord`, at
+/// which point the memory can be returned to the system.
 class TaskStatusRecord {
 public:
   TaskStatusRecordFlags Flags;

stdlib/public/Concurrency/AsyncLet.cpp

@@ -150,7 +150,7 @@ void swift::asyncLet_addImpl(AsyncTask *task, AsyncLet *asyncLet,
   // ok, now that the async let task actually is initialized: attach it to the
   // current task
   bool addedRecord =
-      addStatusRecord(record, [&](ActiveTaskStatus parentStatus) {
+      addStatusRecord(record, [&](ActiveTaskStatus parentStatus, ActiveTaskStatus& newStatus) {
         updateNewChildWithParentAndGroupState(task, parentStatus, NULL);
         return true;
       });

stdlib/public/Concurrency/Task.cpp

@@ -1533,8 +1533,8 @@ swift_task_addCancellationHandlerImpl(
 
   bool fireHandlerNow = false;
 
-  addStatusRecord(record, [&](ActiveTaskStatus parentStatus) {
-    if (parentStatus.isCancelled()) {
+  addStatusRecord(record, [&](ActiveTaskStatus oldStatus, ActiveTaskStatus& newStatus) {
+    if (oldStatus.isCancelled()) {
       fireHandlerNow = true;
       // We don't fire the cancellation handler here since this function needs
       // to be idempotent

stdlib/public/Concurrency/TaskGroup.cpp

@@ -888,10 +888,10 @@ static void swift_taskGroup_initializeWithFlagsImpl(size_t rawGroupFlags,
   assert(record->getKind() == swift::TaskStatusRecordKind::TaskGroup);
 
   // ok, now that the group actually is initialized: attach it to the task
-  addStatusRecord(record, [&](ActiveTaskStatus parentStatus) {
+  addStatusRecord(record, [&](ActiveTaskStatus oldStatus, ActiveTaskStatus& newStatus) {
     // If the task has already been cancelled, reflect that immediately in
     // the group's status.
-    if (parentStatus.isCancelled()) {
+    if (oldStatus.isCancelled()) {
       impl->statusCancel();
     }
     return true;

stdlib/public/Concurrency/TaskPrivate.h

@@ -51,6 +51,7 @@ namespace swift {
 
 class AsyncTask;
 class TaskGroup;
+class ActiveTaskStatus;
 
 /// Allocate task-local memory on behalf of a specific task,
 /// not necessarily the current one.  Generally this should only be
@@ -115,6 +116,63 @@ _swift_task_getDispatchQueueSerialExecutorWitnessTable() {
 }
 #endif
 
+/*************** Methods for Status records manipulation ******************/
+
+/// Remove the status record from input task which may not be the current task.
+/// This may be called asynchronously from the current task.  After this call
+/// returns, the record's memory can be freely modified or deallocated.  The
+/// record must be registered with the task. If it isn't, this function will
+/// crash.
+///
+/// This function also takes in a function_ref which is given the old
+/// ActiveTaskStatus on the task and a reference to the new ActiveTaskStatus
+/// that is to be set on the task that we are removing the record from. It may
+/// modify the new ActiveTaskStatus that is to be set on the task. This function
+/// may be called multiple times inside a RMW loop and must be therefore be
+/// idempotent. The new status passed to `fn` is freshly derived from the
+/// current status and does not include modifications made by previous runs
+/// through the loop
+SWIFT_CC(swift)
+void removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn);
+
+/// Remove a status record from the current task.  After this call returns,
+/// the record's memory can be freely modified or deallocated.
+///
+/// This must be called synchronously with the task.
+///
+/// The given record need not be the last record added to
+/// the task, but the operation may be less efficient if not.
+SWIFT_CC(swift)
+void removeStatusRecord(TaskStatusRecord *record);
+
+/// Add a status record to the current task. This must be called synchronously
+/// with the task.
+///
+/// This function also takes in a function_ref which is given the old
+/// ActiveTaskStatus on the task and a reference to the new ActiveTaskStatus
+/// that is to be set on the task that we are adding the record to.
+///
+/// This can be used by the function to
+/// (1) to determine if the current status of the task permits adding the status
+/// record
+/// (2) modify the active task status that is to be set if needed
+///
+/// If the function_ref returns false, the status record is not added to the
+/// task. This function_ref may be called multiple times and must be idempotent.
+/// The new status passed to `fn` is freshly derived from the current status and
+/// does not include modifications made by previous runs through the loop
+SWIFT_CC(swift)
+bool addStatusRecord(TaskStatusRecord *record,
+     llvm::function_ref<bool(ActiveTaskStatus, 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);
+
 // ==== ------------------------------------------------------------------------
 
 namespace {
@@ -851,39 +909,6 @@ 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