Merge pull request #63919 from apple/rokhinip/105932276-track-task-enqueued-executor

Track  the executor that a task is enqueued in, in dependency records

Author: rokhinip

include/swift/ABI/Executor.h

@@ -67,6 +67,7 @@ class ExecutorRef {
     : Identity(identity), Implementation(implementation) {}
 
 public:
+
   /// A generic execution environment.  When running in a generic
   /// environment, it's presumed to be okay to switch synchronously
   /// to an actor.  As an executor request, this represents a request

include/swift/ABI/TaskStatus.h

@@ -22,6 +22,7 @@
 
 #include "swift/ABI/MetadataValues.h"
 #include "swift/ABI/Task.h"
+#include "swift/ABI/Executor.h"
 #include "swift/Runtime/HeapObject.h"
 
 namespace swift {
@@ -278,22 +279,23 @@ class EscalationNotificationStatusRecord : public TaskStatusRecord {
 // This record is allocated for a task to record what it is dependent on before
 // the task can make progress again.
 class TaskDependencyStatusRecord : public TaskStatusRecord {
-  // A word sized storage which references what this task is suspended waiting
-  // for. Note that this is different from the waitQueue in the future fragment
-  // of a task since that denotes all the tasks which this specific task, will
-  // unblock.
+  // A word sized storage which references what this task is waiting for. Note
+  // that this is different from the waitQueue in the future fragment of a task
+  // since that denotes all the tasks which this specific task, will unblock.
   //
   // This field is only really pointing to something valid when the
-  // ActiveTaskStatus specifies that the task is suspended. It can be accessed
-  // asynchronous to the task during escalation which will therefore require the
-  // task status record lock for synchronization.
+  // ActiveTaskStatus specifies that the task is suspended or enqueued. It can
+  // be accessed asynchronous to the task during escalation which will therefore
+  // require the task status record lock for synchronization.
   //
   // When a task has TaskDependencyStatusRecord in the status record list, it
   // must be the innermost status record, barring the status record lock which
   // could be taken while this record is present.
   //
   // The type of thing we are waiting on, is specified in the enum below
-  union {
+  union Dependent {
+    constexpr Dependent() {}
+
     // This task is suspended waiting on another task. This could be an async
     // let child task or it could be another unstructured task.
     AsyncTask *Task;
@@ -310,40 +312,68 @@ class TaskDependencyStatusRecord : public TaskStatusRecord {
     // the duration of the wait. We do not need to take an additional +1 on this
     // task group in this dependency record.
     TaskGroup *TaskGroup;
-  } WaitingOn;
+
+    // The task is enqueued waiting on an executor. It could be any kind of
+    // executor - the generic executor, the default actor's executor, or an
+    // actor with a custom executor.
+    //
+    // This information is helpful to know *where* a task is enqueued into
+    // (potentially intrusively), so that the appropriate escalation effect
+    // (which may be different for each type of executor) can happen if a task
+    // is escalated while enqueued.
+    ExecutorRef Executor;
+  } DependentOn;
 
   // Enum specifying the type of dependency this task has
   enum DependencyKind {
     WaitingOnTask = 1,
     WaitingOnContinuation,
     WaitingOnTaskGroup,
+
+    EnqueuedOnExecutor,
   } DependencyKind;
 
+  // The task that has this task status record - ie a backpointer from the
+  // record to the task with the record. This is not its own +1, we rely on the
+  // fact that since this status record is linked into a task, the task is
+  // already alive and maintained by someone and we can safely borrow the
+  // reference.
+  AsyncTask *WaitingTask;
+
 public:
-  TaskDependencyStatusRecord(AsyncTask *task) :
+  TaskDependencyStatusRecord(AsyncTask *waitingTask, AsyncTask *task) :
     TaskStatusRecord(TaskStatusRecordKind::TaskDependency),
-        DependencyKind(WaitingOnTask) {
-      WaitingOn.Task = task;
+        DependencyKind(WaitingOnTask), WaitingTask(waitingTask) {
+      DependentOn.Task = task;
   }
 
-  TaskDependencyStatusRecord(ContinuationAsyncContext *context) :
+  TaskDependencyStatusRecord(AsyncTask *waitingTask, ContinuationAsyncContext *context) :
     TaskStatusRecord(TaskStatusRecordKind::TaskDependency),
-        DependencyKind(WaitingOnContinuation) {
-      WaitingOn.Continuation = context;
+        DependencyKind(WaitingOnContinuation), WaitingTask(waitingTask) {
+      DependentOn.Continuation = context;
   }
 
-  TaskDependencyStatusRecord(TaskGroup *taskGroup) :
+  TaskDependencyStatusRecord(AsyncTask *waitingTask, TaskGroup *taskGroup) :
     TaskStatusRecord(TaskStatusRecordKind::TaskDependency),
-        DependencyKind(WaitingOnTaskGroup) {
-      WaitingOn.TaskGroup = taskGroup;
+        DependencyKind(WaitingOnTaskGroup), WaitingTask(waitingTask){
+      DependentOn.TaskGroup = taskGroup;
   }
 
-  void destroy() { }
+  TaskDependencyStatusRecord(AsyncTask *waitingTask, ExecutorRef executor) :
+    TaskStatusRecord(TaskStatusRecordKind::TaskDependency),
+        DependencyKind(EnqueuedOnExecutor), WaitingTask(waitingTask) {
+      DependentOn.Executor = executor;
+  }
 
   static bool classof(const TaskStatusRecord *record) {
     return record->getKind() == TaskStatusRecordKind::TaskDependency;
   }
 
+  void updateDependencyToEnqueuedOn(ExecutorRef executor) {
+    DependencyKind = EnqueuedOnExecutor;
+    DependentOn.Executor = executor;
+  }
+
   void performEscalationAction(JobPriority newPriority);
 };
 

stdlib/public/Concurrency/Actor.cpp

@@ -1808,19 +1808,21 @@ static void swift_task_enqueueImpl(Job *job, ExecutorRef executor) {
     return swift_task_enqueueGlobal(job);
 
   if (executor.isDefaultActor()) {
-#if SWIFT_CONCURRENCY_ACTORS_AS_LOCKS
-    assert(false && "Should not enqueue tasks on actors in actors as locks model");
-#else
-    return asImpl(executor.getDefaultActor())->enqueue(job, job->getPriority());
-#endif
+    return swift_defaultActor_enqueue(job, executor.getDefaultActor());
   }
 
+  // For main actor or actors with custom executors
   auto wtable = executor.getSerialExecutorWitnessTable();
   auto executorObject = executor.getIdentity();
   auto executorType = swift_getObjectType(executorObject);
   _swift_task_enqueueOnExecutor(job, executorObject, executorType, wtable);
 }
 
+SWIFT_CC(swift)
+void swift::swift_executor_escalate(ExecutorRef executor, AsyncTask *task,
+  JobPriority newPriority) {
+}
+
 #define OVERRIDE_ACTOR COMPATIBILITY_OVERRIDE
 #include COMPATIBILITY_OVERRIDE_INCLUDE_PATH
 

stdlib/public/Concurrency/TaskPrivate.h

@@ -122,6 +122,12 @@ _swift_task_getDispatchQueueSerialExecutorWitnessTable() {
 }
 #endif
 
+// The task listed as argument is escalated to a new priority. Pass that
+// inforamtion along to the executor that it is enqueued into.
+SWIFT_CC(swift)
+void
+swift_executor_escalate(ExecutorRef executor, AsyncTask *task, JobPriority newPriority);
+
 /*************** Methods for Status records manipulation ******************/
 
 /// Add a status record to the input task.
@@ -839,17 +845,19 @@ inline void AsyncTask::flagAsRunning() {
 
   if (!oldStatus.hasTaskDependency()) {
     SWIFT_TASK_DEBUG_LOG("%p->flagAsRunning() with no task dependency", this);
+    assert(_private().dependencyRecord == nullptr);
+
     while (true) {
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
-    // Task's priority is greater than the thread's - do a self escalation
-    qos_class_t maxTaskPriority = (qos_class_t) oldStatus.getStoredPriority();
-    if (threadOverrideInfo.can_override && (maxTaskPriority > overrideFloor)) {
-      SWIFT_TASK_DEBUG_LOG("[Override] Self-override thread with oq_floor %#x to match task %p's max priority %#x",
-          overrideFloor, this, maxTaskPriority);
-
-      (void) swift_dispatch_thread_override_self(maxTaskPriority);
-      overrideFloor = maxTaskPriority;
-    }
+      // Task's priority is greater than the thread's - do a self escalation
+      qos_class_t maxTaskPriority = (qos_class_t) oldStatus.getStoredPriority();
+      if (threadOverrideInfo.can_override && (maxTaskPriority > overrideFloor)) {
+        SWIFT_TASK_DEBUG_LOG("[Override] Self-override thread with oq_floor %#x to match task %p's max priority %#x",
+            overrideFloor, this, maxTaskPriority);
+
+        (void) swift_dispatch_thread_override_self(maxTaskPriority);
+        overrideFloor = maxTaskPriority;
+      }
 #endif
       // Set self as executor and remove escalation bit if any - the task's
       // priority escalation has already been reflected on the thread.
@@ -923,83 +931,75 @@ inline void AsyncTask::flagAsRunning() {
 inline void AsyncTask::flagAsAndEnqueueOnExecutor(ExecutorRef newExecutor) {
 #if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
   assert(false && "Should not enqueue any tasks to execute in task-to-thread model");
-#else
+#else /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
   auto oldStatus = _private()._status().load(std::memory_order_relaxed);
-  JobPriority priority = JobPriority::Unspecified;
-
-  if (oldStatus.isRunning()) {
-    SWIFT_TASK_DEBUG_LOG("%p->flagAsAndEnqueueOnExecutor() running to enqueued", this);
-    // Case 1:
-    //  running -> enqueued
-    //  Most likely due to task running into actor contention
-    //  TODO: Need to record a new task dependency
-    while (true) {
-      // Drop execution lock and any override the thread might have received as
-      // a result of executing it previously. Mark the task as being enqueued
-      auto newStatus = oldStatus.withRunning(false);
-      newStatus = newStatus.withoutStoredPriorityEscalation();
-      newStatus = newStatus.withEnqueued();
+  assert(!oldStatus.isEnqueued());
 
-      if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
-              /* success */std::memory_order_relaxed,
-              /* failure */std::memory_order_relaxed)) {
-        newStatus.traceStatusChanged(this);
-        priority = newStatus.getStoredPriority();
-        break;
-      }
-    }
-#if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
-    // The thread was previously running the task, now that we aren't, we need
-    // to remove any task escalation on the thread as a result of the task.
-    if (oldStatus.isStoredPriorityEscalated()) {
-      SWIFT_TASK_DEBUG_LOG("[Override] Reset override %#x on thread from task %p",
-        oldStatus.getStoredPriority(), this);
-      swift_dispatch_lock_override_end((qos_class_t) oldStatus.getStoredPriority());
-    }
-#endif /* SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION */
-    swift_task_exitThreadLocalContext((char *)&_private().ExclusivityAccessSet[0]);
-    restoreTaskVoucher(this);
-
-  } else if (oldStatus.hasTaskDependency()) {
-    //  Case 2: suspended -> enqueued
-    //  Subcase 2a: Task had a dependency which is now cleared - remove task
-    //  dependency record and destroy it
-    auto dependencyRecord = _private().dependencyRecord;
+  if (!oldStatus.isRunning() && oldStatus.hasTaskDependency()) {
+    // Task went from suspended --> enqueued and has a previous
+    // dependency record.
+    //
+    // Atomically update the existing dependency record with new dependency
+    // information.
+    TaskDependencyStatusRecord *dependencyRecord = _private().dependencyRecord;
     assert(dependencyRecord != nullptr);
-    SWIFT_TASK_DEBUG_LOG("%p->flagAsAndEnqueueOnExecutor() suspended to enqueued and remove dependency %p", this, dependencyRecord);
 
-    removeStatusRecord(this, dependencyRecord, [&](ActiveTaskStatus unused,
-                          ActiveTaskStatus& newStatus) {
+    SWIFT_TASK_DEBUG_LOG("[Dependency] %p->flagAsAndEnqueueOnExecutor() and update dependencyRecord %p",
+      this, dependencyRecord);
+
+    updateStatusRecord(this, dependencyRecord, [&] {
+
+      // Update dependency record to the new dependency
+      dependencyRecord->updateDependencyToEnqueuedOn(newExecutor);
+
+    }, oldStatus, [&](ActiveTaskStatus unused, ActiveTaskStatus &newStatus) {
+
+      // Remove escalation bits + set enqueued bit
       newStatus = newStatus.withoutStoredPriorityEscalation();
       newStatus = newStatus.withEnqueued();
-      newStatus = newStatus.withoutTaskDependency();
-
-      priority = newStatus.getStoredPriority();
+      assert(newStatus.hasTaskDependency());
     });
-    this->destroyTaskDependency(dependencyRecord);
-
   } else {
-    // Case 2:  suspended -> enqueued
-    // Subcase 2b: Task is newly created and enqueued to run - no task
-    // dependency record present
-    SWIFT_TASK_DEBUG_LOG("%p->flagAsAndEnqueueOnExecutor() suspended to enqueued with no dependency", this);
-    auto newStatus = oldStatus;
-    while (true) {
+    // 2 subcases:
+    // * Task went from running on this thread --> enqueued on executor
+    // * Task went from suspended to enqueued on this executor and has no
+    // dependency record (Eg. newly created)
+    assert(_private().dependencyRecord == nullptr);
+
+    void *allocation = _swift_task_alloc_specific(this, sizeof(class TaskDependencyStatusRecord));
+    TaskDependencyStatusRecord *dependencyRecord = _private().dependencyRecord = ::new (allocation) TaskDependencyStatusRecord(this, newExecutor);
+    SWIFT_TASK_DEBUG_LOG("[Dependency] %p->flagAsAndEnqueueOnExecutor() with dependencyRecord %p", this,
+      dependencyRecord);
+
+    addStatusRecord(this, dependencyRecord, oldStatus, [&](ActiveTaskStatus unused,
+                    ActiveTaskStatus &newStatus) {
+
+      newStatus = newStatus.withRunning(false);
       newStatus = newStatus.withoutStoredPriorityEscalation();
       newStatus = newStatus.withEnqueued();
+      newStatus = newStatus.withTaskDependency();
 
-      if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
-              /* success */std::memory_order_relaxed,
-              /* failure */std::memory_order_relaxed)) {
-        newStatus.traceStatusChanged(this);
-        priority = newStatus.getStoredPriority();
-        break;
+      return true;
+    });
+
+    if (oldStatus.isRunning()) {
+#if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
+      // The thread was previously running the task, now that we aren't and
+      // we've successfully escalated the thing the task is waiting on. We need
+      // to remove any task escalation on the thread as a result of the task.
+      if (oldStatus.isStoredPriorityEscalated()) {
+        SWIFT_TASK_DEBUG_LOG("[Override] Reset override %#x on thread from task %p",
+          oldStatus.getStoredPriority(), this);
+        swift_dispatch_lock_override_end((qos_class_t) oldStatus.getStoredPriority());
       }
+#endif /* SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION */
+      swift_task_exitThreadLocalContext((char *)&_private().ExclusivityAccessSet[0]);
+      restoreTaskVoucher(this);
     }
   }
 
   // Set up task for enqueue to next location by setting the Job priority field
-  Flags.setPriority(priority);
+  Flags.setPriority(oldStatus.getStoredPriority());
   concurrency::trace::task_flags_changed(
       this, static_cast<uint8_t>(Flags.getPriority()), Flags.task_isChildTask(),
       Flags.task_isFuture(), Flags.task_isGroupChildTask(),
@@ -1036,9 +1036,8 @@ void AsyncTask::flagAsSuspended(TaskDependencyStatusRecord *dependencyStatusReco
     // Note that we have to do this escalation while adding the status record
     // and not after - we are not guaranteed to be able to have a valid
     // reference to the dependencyStatusRecord or its contents, once we have
-    // published it in the ActiveTaskStatus.
-    SWIFT_TASK_DEBUG_LOG("[Dependency] Escalate the dependency %p of task %p",
-                  dependencyStatusRecord, this);
+    // published it in the ActiveTaskStatus since someone else could
+    // concurrently made us runnable.
     dependencyStatusRecord->performEscalationAction(newStatus.getStoredPriority());
 
     // Always add the dependency status record
@@ -1062,8 +1061,6 @@ void AsyncTask::flagAsSuspended(TaskDependencyStatusRecord *dependencyStatusReco
 
 inline void AsyncTask::destroyTaskDependency(TaskDependencyStatusRecord *dependencyRecord) {
   assert(_private().dependencyRecord == dependencyRecord);
-
-  dependencyRecord->destroy();
   _swift_task_dealloc_specific(this, dependencyRecord);
 
   _private().dependencyRecord = nullptr;
@@ -1075,7 +1072,7 @@ inline void AsyncTask::flagAsSuspendedOnTask(AsyncTask *task) {
   assert(_private().dependencyRecord == nullptr);
 
   void *allocation = _swift_task_alloc_specific(this, sizeof(class TaskDependencyStatusRecord));
-  auto record = ::new (allocation) TaskDependencyStatusRecord(task);
+  auto record = ::new (allocation) TaskDependencyStatusRecord(this, task);
   SWIFT_TASK_DEBUG_LOG("[Dependency] Create a dependencyRecord %p for dependency on task %p", allocation, task);
   _private().dependencyRecord = record;
 
@@ -1086,7 +1083,7 @@ inline void AsyncTask::flagAsSuspendedOnContinuation(ContinuationAsyncContext *c
   assert(_private().dependencyRecord == nullptr);
 
   void *allocation = _swift_task_alloc_specific(this, sizeof(class TaskDependencyStatusRecord));
-  auto record = ::new (allocation) TaskDependencyStatusRecord(context);
+  auto record = ::new (allocation) TaskDependencyStatusRecord(this, context);
   SWIFT_TASK_DEBUG_LOG("[Dependency] Create a dependencyRecord %p for dependency on continuation %p", allocation, context);
   _private().dependencyRecord = record;
 
@@ -1097,7 +1094,7 @@ inline void AsyncTask::flagAsSuspendedOnTaskGroup(TaskGroup *taskGroup) {
   assert(_private().dependencyRecord == nullptr);
 
   void *allocation = _swift_task_alloc_specific(this, sizeof(class TaskDependencyStatusRecord));
-  auto record = ::new (allocation) TaskDependencyStatusRecord(taskGroup);
+  auto record = ::new (allocation) TaskDependencyStatusRecord(this, taskGroup);
   SWIFT_TASK_DEBUG_LOG("[Dependency] Create a dependencyRecord %p for dependency on taskGroup %p", allocation, taskGroup);
   _private().dependencyRecord = record;