Merge pull request #63916 from apple/rokhinip/105930505-updateStatusRecord-and-friends

Improve functionality available for manipulating status records.

Author: rokhinip

include/swift/ABI/TaskStatus.h

@@ -61,25 +61,10 @@ class TaskStatusRecord {
 
   TaskStatusRecord *getParent() const { return Parent; }
 
-  /// Change the parent of this unregistered status record to the
-  /// given record.
-  ///
-  /// This should be used when the record has been previously initialized
-  /// without knowing what the true parent is.  If we decide to cache
-  /// important information (e.g. the earliest timeout) in the innermost
-  /// status record, this is the method that should fill that in
-  /// from the parent.
+  /// Change the parent of this status record to the given record.
   void resetParent(TaskStatusRecord *newParent) {
     Parent = newParent;
-    // TODO: cache
   }
-
-  /// Splice a record out of the status-record chain.
-  ///
-  /// Unlike resetParent, this assumes that it's just removing one or
-  /// more records from the chain and that there's no need to do any
-  /// extra cache manipulation.
-  void spliceParent(TaskStatusRecord *newParent) { Parent = newParent; }
 };
 
 /// A status record which states that a task has one or

stdlib/public/Concurrency/AsyncLet.cpp

@@ -334,7 +334,7 @@ static void swift_asyncLet_endImpl(AsyncLet *alet) {
 
   // Remove the child record from the parent task
   auto record = asImpl(alet)->getTaskRecord();
-  removeStatusRecord(record);
+  removeStatusRecordFromSelf(record);
 
   // TODO: we need to implicitly await either before the end or here somehow.
 
@@ -362,7 +362,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();
-  removeStatusRecord(record);
+  removeStatusRecordFromSelf(record);
 
   // and finally, release the task and destroy the async-let
   assert(swift_task_getCurrent() && "async-let must have a parent task");

stdlib/public/Concurrency/Task.cpp

@@ -1520,7 +1520,7 @@ swift_task_addCancellationHandlerImpl(
 SWIFT_CC(swift)
 static void swift_task_removeCancellationHandlerImpl(
     CancellationNotificationStatusRecord *record) {
-  removeStatusRecord(record);
+  removeStatusRecordFromSelf(record);
   swift_task_dealloc(record);
 }
 

stdlib/public/Concurrency/TaskGroup.cpp

@@ -985,7 +985,7 @@ void AccumulatingTaskGroup::destroy() {
   assert(this->isEmpty() && "Attempted to destroy non-empty task group!");
 
   // First, remove the group from the task and deallocate the record
-  removeStatusRecord(getTaskRecord());
+  removeStatusRecordFromSelf(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.
@@ -1008,7 +1008,7 @@ void DiscardingTaskGroup::destroy() {
   assert(this->isEmpty() && "Attempted to destroy non-empty task group!");
 
   // First, remove the group from the task and deallocate the record
-  removeStatusRecord(getTaskRecord());
+  removeStatusRecordFromSelf(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.

stdlib/public/Concurrency/TaskPrivate.h

@@ -124,34 +124,6 @@ _swift_task_getDispatchQueueSerialExecutorWitnessTable() {
 
 /*************** 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 input task.
 ///
 /// Clients can optionally pass in the status of the task if they have already
@@ -160,6 +132,12 @@ void removeStatusRecord(TaskStatusRecord *record);
 /// perform. This status will be updated with the last status on the task prior
 /// to updating it with the new status if the input function_ref allows so.
 ///
+/// Clients can optionally pass in the status of the task if they have already
+/// done the load on it or if they require the oldStatus on the task prior to
+/// the atomic ActiveTaskStatus update that addStatusRecord will do. This status
+/// will be updated with the last status on the task prior to updating it with
+/// the new status if the input function_ref allows so.
+///
 /// 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.
@@ -195,6 +173,59 @@ SWIFT_CC(swift)
 bool addStatusRecordToSelf(TaskStatusRecord *record,  ActiveTaskStatus& taskStatus,
      llvm::function_ref<bool(ActiveTaskStatus, ActiveTaskStatus&)> testAddRecord);
 
+/// 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.
+///
+/// The given record need not be the last record added to
+/// the task, but the operation may be less efficient if not.
+///
+/// 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, ActiveTaskStatus& status,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn = nullptr);
+
+SWIFT_CC(swift)
+void removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn = nullptr);
+
+/// Remove a status record from the current task. This must be called
+/// synchronously with the task.
+SWIFT_CC(swift)
+void removeStatusRecordFromSelf(TaskStatusRecord *record, ActiveTaskStatus &status,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn = nullptr);
+
+SWIFT_CC(swift)
+void removeStatusRecordFromSelf(TaskStatusRecord *record,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn = nullptr);
+
+/// Update the specified input status record while holding the status record
+/// lock of the task. The status record must already be registered with the
+/// task - if it isn't, this API provides no additional protections.
+///
+/// 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 when we are unlocking the task status record
+/// lock. 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 updateStatusRecord(AsyncTask *task, TaskStatusRecord *record,
+     llvm::function_ref<void()>updateRecord,
+     ActiveTaskStatus& status,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn = nullptr);
+
 /// 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)
@@ -842,7 +873,7 @@ inline void AsyncTask::flagAsRunning() {
     SWIFT_TASK_DEBUG_LOG("[Dependency] %p->flagAsRunning() and remove dependencyRecord %p",
                     this, dependencyRecord);
 
-    removeStatusRecord(this, dependencyRecord, [&](ActiveTaskStatus oldStatus,
+    removeStatusRecord(this, dependencyRecord, oldStatus, [&](ActiveTaskStatus unused,
                        ActiveTaskStatus& newStatus) {
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
@@ -863,7 +894,6 @@ inline void AsyncTask::flagAsRunning() {
       newStatus = newStatus.withoutEnqueued();
       newStatus = newStatus.withoutTaskDependency();
     });
-
     this->destroyTaskDependency(dependencyRecord);
 
     adoptTaskVoucher(this);

stdlib/public/Concurrency/TaskStatus.cpp

@@ -299,6 +299,7 @@ bool swift::addStatusRecord(AsyncTask *task, TaskStatusRecord *newRecord,
     ActiveTaskStatus& oldStatus,
     llvm::function_ref<bool(ActiveTaskStatus, ActiveTaskStatus&)> shouldAddRecord) {
 
+  SWIFT_TASK_DEBUG_LOG("Adding %p record to task %p", newRecord, task);
   while (true) {
     // Wait for any active lock to be released.
     if (oldStatus.isStatusRecordLocked()) {
@@ -343,16 +344,22 @@ bool swift::addStatusRecordToSelf(TaskStatusRecord *record,
   return addStatusRecord(swift_task_getCurrent(), record, testAddRecord);
 }
 
+SWIFT_CC(swift)
+bool swift::addStatusRecordToSelf(TaskStatusRecord *record, ActiveTaskStatus &status,
+     llvm::function_ref<bool(ActiveTaskStatus, ActiveTaskStatus&)> testAddRecord) {
+  return addStatusRecord(swift_task_getCurrent(), record, status, testAddRecord);
+}
+
 static void removeStatusRecordLocked(ActiveTaskStatus status, TaskStatusRecord *record) {
   bool removedRecord = false;
   auto cur = status.getInnermostRecord();
   assert(cur->getKind() == TaskStatusRecordKind::Private_RecordLock);
 
-  // Splice the record out.
+  // Cut the record out.
   while (cur != nullptr) {
     auto next = cur->getParent();
     if (next == record) {
-      cur->spliceParent(record->getParent());
+      cur->resetParent(record->getParent());
       removedRecord = true;
       break;
     }
@@ -365,13 +372,12 @@ static void removeStatusRecordLocked(ActiveTaskStatus status, TaskStatusRecord *
 // modify some flags in the ActiveTaskStatus at the same time.
 SWIFT_CC(swift)
 void swift::removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
-     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn = nullptr) {
+     ActiveTaskStatus& oldStatus,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn) {
 
   SWIFT_TASK_DEBUG_LOG("remove status record = %p, from task = %p",
                        record, task);
 
-  auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
-
   if (oldStatus.isStatusRecordLocked() &&
         waitForStatusRecordUnlockIfNotSelfLocked(task, oldStatus)) {
     SWIFT_TASK_DEBUG_LOG("[StatusRecordLock] Lock for task %p is already owned by thread", task);
@@ -446,16 +452,46 @@ void swift::removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
 
 }
 
-// Convenience wrapper for when we don't care to make any modifications to the
-// flags fields of the active task status when we are removing a task status
-// record
+// Convenience wrapper for when client hasnt already done the load of the status
 SWIFT_CC(swift)
-void swift::removeStatusRecord(TaskStatusRecord *record) {
+void swift::removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn) {
+  auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
+  return removeStatusRecord(task, record, oldStatus, fn);
+}
 
-  auto task = swift_task_getCurrent();
-  SWIFT_TASK_DEBUG_LOG("remove status record = %p, from current task = %p",
-                       record, task);
-  return removeStatusRecord(task, record);
+// Convenience wrapper for modifications on current task
+SWIFT_CC(swift)
+void swift::removeStatusRecordFromSelf(TaskStatusRecord *record, ActiveTaskStatus &status,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn) {
+  return removeStatusRecord(swift_task_getCurrent(), record, status, fn);
+}
+
+SWIFT_CC(swift)
+void swift::removeStatusRecordFromSelf(TaskStatusRecord *record,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn) {
+  return removeStatusRecord(swift_task_getCurrent(), record, fn);
+}
+
+SWIFT_CC(swift)
+void swift::updateStatusRecord(AsyncTask *task, TaskStatusRecord *record,
+     llvm::function_ref<void()>updateRecord,
+     ActiveTaskStatus& status,
+     llvm::function_ref<void(ActiveTaskStatus, ActiveTaskStatus&)>fn) {
+
+  withStatusRecordLock(task, status, [&](ActiveTaskStatus lockedStatus) {
+#if NDEBUG
+    bool foundRecord = false;
+    for (auto cur: lockedStatus.records()) {
+      if (cur == record) {
+        foundRecord = true;
+        break;
+      }
+    }
+    assert(foundRecord);
+#endif
+    updateRecord();
+  }, fn);
 }
 
 SWIFT_CC(swift)