Task group children should be executed on parent task at time of await

Author: rokhinip

stdlib/public/Concurrency/Task.cpp

@@ -704,6 +704,9 @@ static AsyncTaskAndContext swift_task_create_commonImpl(
     }
     case TaskOptionRecordKind::RunInline: {
       runInlineOption = cast<RunInlineTaskOptionRecord>(option);
+      // TODO (rokhinip): We seem to be creating runInline tasks like detached
+      // tasks but they need to maintain the voucher and priority of calling
+      // thread and therefore need to behave a bit more like SC child tasks.
       break;
     }
     }
@@ -974,6 +977,15 @@ static AsyncTaskAndContext swift_task_create_commonImpl(
   // Attach to the group, if needed.
   if (group) {
     swift_taskGroup_attachChild(group, task);
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+    // We need to take a retain here to keep the child task for the task group
+    // alive. In the non-task-to-thread model, we'd always take this retain
+    // below since we'd enqueue the child task. But since we're not going to be
+    // enqueueing the child task in this model, we need to take this +1 to
+    // balance out the release that exists after the task group child task
+    // creation
+    swift_retain(task);
+#endif
   }
 
   // If we're supposed to copy task locals, do so now.
@@ -988,6 +1000,9 @@ static AsyncTaskAndContext swift_task_create_commonImpl(
 
   // If we're supposed to enqueue the task, do so now.
   if (taskCreateFlags.enqueueJob()) {
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+    assert(false && "Should not be enqueuing tasks in task-to-thread model");
+#endif
     swift_retain(task);
     task->flagAsAndEnqueueOnExecutor(executor);
   }

stdlib/public/Concurrency/TaskGroup.cpp

@@ -284,15 +284,24 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
   };
 
 private:
-#if !SWIFT_STDLIB_SINGLE_THREADED_CONCURRENCY
+#if SWIFT_STDLIB_SINGLE_THREADED_CONCURRENCY || SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  // Synchronization is simple here. In a single threaded mode, all swift tasks
+  // run on a single thread so no coordination is needed. In a task-to-thread
+  // model, only the parent task which created the task group can
+  //
+  //   (a) add child tasks to a group
+  //   (b) run the child tasks
+  //
+  // So we shouldn't need to worry about coordinating between child tasks and
+  // parents in a task group
+  void lock() const {}
+  void unlock() const {}
+#else
   // TODO: move to lockless via the status atomic (make readyQueue an mpsc_queue_t<ReadyQueueItem>)
   mutable std::mutex mutex_;
 
   void lock() const { mutex_.lock(); }
   void unlock() const { mutex_.unlock(); }
-#else
-  void lock() const {}
-  void unlock() const {}
 #endif
 
   /// Used for queue management, counting number of waiting and ready tasks
@@ -437,6 +446,11 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
   /// or a `PollStatus::MustWait` result if there are tasks in flight
   /// and the waitingTask eventually be woken up by a completion.
   PollResult poll(AsyncTask *waitingTask);
+
+private:
+  // Enqueue the completed task onto ready queue if there are no waiting tasks
+  // yet
+  void enqueueCompletedTask(AsyncTask *completedTask, bool hadErrorResult);
 };
 
 } // end anonymous namespace
@@ -569,6 +583,22 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
   }
 }
 
+// TaskGroup is locked upon entry and exit
+void TaskGroupImpl::enqueueCompletedTask(AsyncTask *completedTask, bool hadErrorResult) {
+    // Retain the task while it is in the queue;
+    // it must remain alive until the task group is alive.
+    swift_retain(completedTask);
+
+    auto readyItem = ReadyQueueItem::get(
+        hadErrorResult ? ReadyStatus::Error : ReadyStatus::Success,
+        completedTask
+    );
+
+    assert(completedTask == readyItem.getTask());
+    assert(readyItem.getTask()->isFuture());
+    readyQueue.enqueue(readyItem);
+}
+
 void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
   assert(completedTask);
   assert(completedTask->isFuture());
@@ -610,52 +640,58 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
       if (waitQueue.compare_exchange_strong(
           waitingTask, nullptr,
           /*success*/ std::memory_order_release,
-          /*failure*/ std::memory_order_acquire) &&
-          statusCompletePendingReadyWaiting(assumed)) {
-        // Run the task.
-        auto result = PollResult::get(completedTask, hadErrorResult);
+          /*failure*/ std::memory_order_acquire)) {
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+         // We have completed a child task in a task group task and we know
+         // there is a waiting task who will reevaluate TaskGroupImpl::poll once
+         // we return, by virtue of being in the task-to-thread model.
+         // We want poll() to then satisfy the condition of having readyTasks()
+         // that it can dequeue from the readyQueue so we need to enqueue our
+         // completion.
+
+         // TODO (rokhinip): There's probably a more efficient way to deal with
+         // this since the child task can directly offer the result to the
+         // parent who will run next but that requires a fair bit of plumbing
+         enqueueCompletedTask(completedTask, hadErrorResult);
+         unlock(); // TODO: remove fragment lock, and use status for synchronization
+         return;
+#else /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
+        if (statusCompletePendingReadyWaiting(assumed)) {
+          // Run the task.
+          auto result = PollResult::get(completedTask, hadErrorResult);
 
-        unlock(); // TODO: remove fragment lock, and use status for synchronization
-
-        auto waitingContext =
-            static_cast<TaskFutureWaitAsyncContext *>(
-                waitingTask->ResumeContext);
+          unlock(); // TODO: remove fragment lock, and use status for synchronization
 
-        fillGroupNextResult(waitingContext, result);
-        detachChild(result.retainedTask);
+          auto waitingContext =
+              static_cast<TaskFutureWaitAsyncContext *>(
+                  waitingTask->ResumeContext);
 
-        _swift_tsan_acquire(static_cast<Job *>(waitingTask));
+          fillGroupNextResult(waitingContext, result);
+          detachChild(result.retainedTask);
 
-        // TODO: allow the caller to suggest an executor
-        waitingTask->flagAsAndEnqueueOnExecutor(ExecutorRef::generic());
-        return;
-      } // else, try again
+          _swift_tsan_acquire(static_cast<Job *>(waitingTask));
+          // TODO: allow the caller to suggest an executor
+          waitingTask->flagAsAndEnqueueOnExecutor(ExecutorRef::generic());
+          return;
+        } // else, try again
+#endif
+      }
     }
-
     llvm_unreachable("should have enqueued and returned.");
+  } else {
+    // ==== b) enqueue completion ------------------------------------------------
+    //
+    // else, no-one was waiting (yet), so we have to instead enqueue to the message
+    // queue when a task polls during next() it will notice that we have a value
+    // ready for it, and will process it immediately without suspending.
+    assert(!waitQueue.load(std::memory_order_relaxed));
+
+    SWIFT_TASK_DEBUG_LOG("group has no waiting tasks, RETAIN and store ready task = %p",
+                         completedTask);
+    enqueueCompletedTask(completedTask, hadErrorResult);
+    unlock(); // TODO: remove fragment lock, and use status for synchronization
   }
 
-  // ==== b) enqueue completion ------------------------------------------------
-  //
-  // else, no-one was waiting (yet), so we have to instead enqueue to the message
-  // queue when a task polls during next() it will notice that we have a value
-  // ready for it, and will process it immediately without suspending.
-  assert(!waitQueue.load(std::memory_order_relaxed));
-  SWIFT_TASK_DEBUG_LOG("group has no waiting tasks, RETAIN and store ready task = %p",
-                       completedTask);
-  // Retain the task while it is in the queue;
-  // it must remain alive until the task group is alive.
-  swift_retain(completedTask);
-
-  auto readyItem = ReadyQueueItem::get(
-      hadErrorResult ? ReadyStatus::Error : ReadyStatus::Success,
-      completedTask
-  );
-
-  assert(completedTask == readyItem.getTask());
-  assert(readyItem.getTask()->isFuture());
-  readyQueue.enqueue(readyItem);
-  unlock(); // TODO: remove fragment lock, and use status for synchronization
   return;
 }
 
@@ -694,7 +730,7 @@ static void swift_taskGroup_wait_next_throwingImpl(
     OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
     TaskGroup *_group,
     ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
-    AsyncContext *rawContext) {
+    AsyncContext *rawContext) SWIFT_OPTNONE {
   auto waitingTask = swift_task_getCurrent();
   waitingTask->ResumeTask = task_group_wait_resume_adapter;
   waitingTask->ResumeContext = rawContext;
@@ -719,9 +755,9 @@ static void swift_taskGroup_wait_next_throwingImpl(
 #ifdef __ARM_ARCH_7K__
     return workaround_function_swift_taskGroup_wait_next_throwingImpl(
         resultPointer, callerContext, _group, resumeFunction, rawContext);
-#else
+#else /* __ARM_ARCH_7K__ */
     return;
-#endif
+#endif /* __ARM_ARCH_7K__ */
 
   case PollStatus::Empty:
   case PollStatus::Error:
@@ -739,16 +775,25 @@ static void swift_taskGroup_wait_next_throwingImpl(
   }
 }
 
-PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
+PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) SWIFT_OPTNONE {
   lock(); // TODO: remove group lock, and use status for synchronization
   SWIFT_TASK_DEBUG_LOG("poll group = %p", this);
-  auto assumed = statusMarkWaitingAssumeAcquire();
 
   PollResult result;
   result.storage = nullptr;
   result.successType = nullptr;
   result.retainedTask = nullptr;
 
+  // Have we suspended the task?
+  bool hasSuspended = false;
+  bool haveRunOneChildTaskInline = false;
+
+reevaluate_if_taskgroup_has_results:;
+  auto assumed = statusMarkWaitingAssumeAcquire();
+  if (haveRunOneChildTaskInline) {
+    assert(assumed.readyTasks());
+  }
+
   // ==== 1) bail out early if no tasks are pending ----------------------------
   if (assumed.isEmpty()) {
     SWIFT_TASK_DEBUG_LOG("poll group = %p, group is empty, no pending tasks", this);
@@ -762,9 +807,6 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
     return result;
   }
 
-  // Have we suspended the task?
-  bool hasSuspended = false;
-
   auto waitHead = waitQueue.load(std::memory_order_acquire);
 
   // ==== 2) Ready task was polled, return with it immediately -----------------
@@ -779,17 +821,17 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
         /*success*/ std::memory_order_relaxed,
         /*failure*/ std::memory_order_acquire)) {
 
-      // Success! We are allowed to poll.
-      ReadyQueueItem item;
-      bool taskDequeued = readyQueue.dequeue(item);
-      assert(taskDequeued); (void) taskDequeued;
-
       // We're going back to running the task, so if we suspended before,
       // we need to flag it as running again.
       if (hasSuspended) {
         waitingTask->flagAsRunning();
       }
 
+      // Success! We are allowed to poll.
+      ReadyQueueItem item;
+      bool taskDequeued = readyQueue.dequeue(item);
+      assert(taskDequeued); (void) taskDequeued;
+
       assert(item.getTask()->isFuture());
       auto futureFragment = item.getTask()->futureFragment();
 
@@ -845,6 +887,28 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
         /*success*/ std::memory_order_release,
         /*failure*/ std::memory_order_acquire)) {
       unlock(); // TODO: remove fragment lock, and use status for synchronization
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+      // The logic here is paired with the logic in TaskGroupImpl::offer. Once
+      // we run the
+      auto oldTask = _swift_task_clearCurrent();
+      assert(oldTask == waitingTask);
+
+      auto childTask = getTaskRecord()->getFirstChild();
+      assert(childTask != NULL);
+
+      SWIFT_TASK_DEBUG_LOG("[RunInline] Switching away from running %p to now running %p", oldTask, childTask);
+      // Run the new task on the same thread now - this should run the new task to
+      // completion. All swift tasks in task-to-thread model run on generic
+      // executor
+      swift_job_run(childTask, ExecutorRef::generic());
+      haveRunOneChildTaskInline = true;
+
+      SWIFT_TASK_DEBUG_LOG("[RunInline] Switching back from running %p to now running %p", childTask, oldTask);
+      // We are back to being the parent task and now that we've run the child
+      // task, we should reevaluate parent task
+      _swift_task_setCurrent(oldTask);
+      goto reevaluate_if_taskgroup_has_results;
+#endif
       // no ready tasks, so we must wait.
       result.status = PollStatus::MustWait;
       _swift_task_clearCurrent();

stdlib/public/Concurrency/TaskGroup.swift

@@ -238,11 +238,19 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
     operation: __owned @Sendable @escaping () async -> ChildTaskResult
   ) {
 #if compiler(>=5.5) && $BuiltinCreateAsyncTaskInGroup
+#if SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
+    let flags = taskCreateFlags(
+      priority: priority, isChildTask: true, copyTaskLocals: false,
+      inheritContext: false, enqueueJob: false,
+      addPendingGroupTaskUnconditionally: true
+    )
+#else
     let flags = taskCreateFlags(
       priority: priority, isChildTask: true, copyTaskLocals: false,
       inheritContext: false, enqueueJob: true,
       addPendingGroupTaskUnconditionally: true
     )
+#endif
 
     // Create the task in this group.
     _ = Builtin.createAsyncTaskInGroup(flags, _group, operation)
@@ -272,12 +280,19 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
       // the group is cancelled and is not accepting any new work
       return false
     }
-
+#if SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
+    let flags = taskCreateFlags(
+      priority: priority, isChildTask: true, copyTaskLocals: false,
+      inheritContext: false, enqueueJob: false,
+      addPendingGroupTaskUnconditionally: false
+    )
+#else
     let flags = taskCreateFlags(
       priority: priority, isChildTask: true, copyTaskLocals: false,
       inheritContext: false, enqueueJob: true,
       addPendingGroupTaskUnconditionally: false
     )
+#endif
 
     // Create the task in this group.
     _ = Builtin.createAsyncTaskInGroup(flags, _group, operation)
@@ -634,8 +649,8 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   ///     guard let result = await group.nextResult() else {
   ///         return  // No task to wait on, which won't happen in this example.
   ///     }
-  ///     
-  ///     switch result { 
+  ///
+  ///     switch result {
   ///     case .success(let value): print(value)
   ///     case .failure(let error): print("Failure: \(error)")
   ///     }
@@ -810,15 +825,15 @@ extension ThrowingTaskGroup: AsyncSequence {
   ///     group.addTask { 1 }
   ///     group.addTask { throw SomeError }
   ///     group.addTask { 2 }
-  ///     
-  ///     do { 
+  ///
+  ///     do {
   ///         // Assuming the child tasks complete in order, this prints "1"
   ///         // and then throws an error.
   ///         for try await r in group { print(r) }
   ///     } catch {
   ///         // Resolve the error.
   ///     }
-  ///     
+  ///
   ///     // Assuming the child tasks complete in order, this prints "2".
   ///     for try await r in group { print(r) }
   ///
@@ -847,7 +862,7 @@ extension ThrowingTaskGroup: AsyncSequence {
     /// this iterator is guaranteed to never produce more values.
     ///
     /// For more information about the iteration order and semantics,
-    /// see `ThrowingTaskGroup.next()` 
+    /// see `ThrowingTaskGroup.next()`
     ///
     /// - Throws: The error thrown by the next child task that completes.
     ///