In task-to-thread model, make sure that when we are waiting on an

async-let task, we run it in the context of the parent task.

Author: rokhinip

include/swift/Runtime/Concurrency.h

@@ -41,7 +41,7 @@
 
 // Does the runtime integrate with libdispatch?
 #ifndef SWIFT_CONCURRENCY_ENABLE_DISPATCH
-#if SWIFT_CONCURRENCY_COOPERATIVE_GLOBAL_EXECUTOR
+#if SWIFT_CONCURRENCY_COOPERATIVE_GLOBAL_EXECUTOR || SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
 #define SWIFT_CONCURRENCY_ENABLE_DISPATCH 0
 #else
 #define SWIFT_CONCURRENCY_ENABLE_DISPATCH 1

stdlib/public/Concurrency/Actor.cpp

@@ -249,6 +249,12 @@ AsyncTask *swift::_swift_task_clearCurrent() {
   return task;
 }
 
+AsyncTask *swift::_swift_task_setCurrent(AsyncTask *new_task) {
+  auto task = ActiveTask::get();
+  ActiveTask::set(new_task);
+  return task;
+}
+
 SWIFT_CC(swift)
 static ExecutorRef swift_task_getCurrentExecutorImpl() {
   auto currentTracking = ExecutorTrackingInfo::current();

stdlib/public/Concurrency/AsyncLet.cpp

@@ -167,7 +167,13 @@ void swift::swift_asyncLet_start(AsyncLet *alet,
                                  void *closureEntryPoint,
                                  HeapObject *closureContext) {
   auto flags = TaskCreateFlags();
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  // In the task to thread model, we don't want tasks to start running on
+  // separate threads - they will run in the context of the parent
+  flags.setEnqueueJob(false);
+#else
   flags.setEnqueueJob(true);
+#endif
 
   AsyncLetTaskOptionRecord asyncLetOptionRecord(alet);
   asyncLetOptionRecord.Parent = options;
@@ -191,7 +197,14 @@ void swift::swift_asyncLet_begin(AsyncLet *alet,
                        resultBuffer);
 
   auto flags = TaskCreateFlags();
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  // In the task to thread model, we don't want tasks to start running on
+  // separate threads - they will run in the context of the parent
+  flags.setEnqueueJob(false);
+#else
   flags.setEnqueueJob(true);
+#endif
+
 
   AsyncLetWithBufferTaskOptionRecord asyncLetOptionRecord(alet, resultBuffer);
   asyncLetOptionRecord.Parent = options;

stdlib/public/Concurrency/Executor.swift

@@ -93,7 +93,7 @@ func _checkExpectedExecutor(_filenameStart: Builtin.RawPointer,
     _filenameStart, _filenameLength, _filenameIsASCII, _line, _executor)
 }
 
-#if !SWIFT_STDLIB_SINGLE_THREADED_CONCURRENCY
+#if !SWIFT_STDLIB_SINGLE_THREADED_CONCURRENCY && !SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
 // This must take a DispatchQueueShim, not something like AnyObject,
 // or else SILGen will emit a retain/release in unoptimized builds,
 // which won't work because DispatchQueues aren't actually

stdlib/public/Concurrency/Task.cpp

@@ -187,6 +187,29 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
       escalatedPriority = waitingStatus.getStoredPriority();
     }
 
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+    // In the task to thread model, we will execute the task that we are waiting
+    // on, on the current thread itself. As a result, do not bother adding the
+    // waitingTask to any thread queue. Instead, we will clear the old task, run
+    // the new one and then reattempt to continue running the old task
+
+    auto oldTask = _swift_task_clearCurrent();
+    assert(oldTask == waitingTask);
+
+    SWIFT_TASK_DEBUG_LOG("[RunInline] Switching away from running %p to now running %p", oldTask, this);
+    // 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(this, ExecutorRef::generic());
+
+    SWIFT_TASK_DEBUG_LOG("[RunInline] Switching back from running %p to now running %p", this, oldTask);
+
+    // We now are back in the context of the waiting task and need to reevaluate
+    // our state
+    _swift_task_setCurrent(oldTask);
+    queueHead = fragment->waitQueue.load(std::memory_order_acquire);
+    continue;
+#else
     // Put the waiting task at the beginning of the wait queue.
     waitingTask->getNextWaitingTask() = queueHead.getTask();
     auto newQueueHead = WaitQueueItem::get(Status::Executing, waitingTask);
@@ -198,6 +221,7 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
       _swift_task_clearCurrent();
       return FutureFragment::Status::Executing;
     }
+#endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
   }
 }
 
@@ -255,8 +279,13 @@ void AsyncTask::completeFuture(AsyncContext *context) {
   // Schedule every waiting task on the executor.
   auto waitingTask = queueHead.getTask();
 
-  if (!waitingTask)
+  if (!waitingTask) {
     SWIFT_TASK_DEBUG_LOG("task %p had no waiting tasks", this);
+  } else {
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+    assert(false && "Task should have no waiters in task-to-thread model");
+#endif
+  }
 
   while (waitingTask) {
     // Find the next waiting task before we invalidate it by resuming

stdlib/public/Concurrency/TaskPrivate.h

@@ -79,6 +79,8 @@ void asyncLet_addImpl(AsyncTask *task, AsyncLet *asyncLet,
 
 /// Clear the active task reference for the current thread.
 AsyncTask *_swift_task_clearCurrent();
+/// Set the active task reference for the current thread.
+AsyncTask *_swift_task_setCurrent(AsyncTask *newTask);
 
 /// release() establishes a happens-before relation with a preceding acquire()
 /// on the same address.
@@ -89,7 +91,7 @@ void _swift_tsan_release(void *addr);
 /// executors.
 #define DISPATCH_QUEUE_GLOBAL_EXECUTOR (void *)1
 
-#if !SWIFT_STDLIB_SINGLE_THREADED_CONCURRENCY
+#if SWIFT_CONCURRENCY_ENABLE_DISPATCH
 inline SerialExecutorWitnessTable *
 _swift_task_getDispatchQueueSerialExecutorWitnessTable() {
   extern SerialExecutorWitnessTable wtable
@@ -712,7 +714,9 @@ retry:;
 ///
 /// rdar://88366470 (Direct handoff behaviour when tasks switch executors)
 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
   SWIFT_TASK_DEBUG_LOG("%p->flagAsAndEnqueueOnExecutor()", this);
   auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   auto newStatus = oldStatus;
@@ -746,7 +750,7 @@ inline void AsyncTask::flagAsAndEnqueueOnExecutor(ExecutorRef newExecutor) {
         oldStatus.getStoredPriority(), this);
       swift_dispatch_lock_override_end((qos_class_t) oldStatus.getStoredPriority());
     }
-#endif
+#endif /* SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION */
     swift_task_exitThreadLocalContext((char *)&_private().ExclusivityAccessSet[0]);
     restoreTaskVoucher(this);
   }
@@ -759,6 +763,7 @@ inline void AsyncTask::flagAsAndEnqueueOnExecutor(ExecutorRef newExecutor) {
       Flags.task_isAsyncLetTask());
 
   swift_task_enqueue(this, newExecutor);
+#endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
 }
 
 inline void AsyncTask::flagAsSuspended() {

test/Concurrency/execute_child_on_await.swift

@@ -0,0 +1,74 @@
+// RUN: %target-run-simple-swift(-parse-as-library  -Xfrontend -disable-availability-checking -Xfrontend -concurrency-model=task-to-thread)
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// REQUIRES: freestanding
+// REQUIRES: concurrency_runtime
+
+@_spi(_TaskToThreadModel) import _Concurrency
+import StdlibUnittest
+import Darwin
+
+var a = 0;
+
+func foo() async -> Int {
+  a += 1
+  return a
+}
+
+func bar() async -> Int {
+  a += 2
+  return a
+}
+
+func asyncFib(_ n: Int, _ expected_thread: pthread_t) async -> Int {
+  expectEqual(pthread_self(), expected_thread);
+
+  if n == 0 || n == 1 {
+    return n
+  }
+
+  async let first = asyncFib(n-2, expected_thread)
+  async let second = asyncFib(n-1, expected_thread)
+
+  return (await second) + (await first)
+}
+
+func fib(_ n: Int) -> Int {
+  var first = 0
+  var second = 1
+  for _ in 0..<n {
+    let temp = first
+    first = second
+    second = temp + first
+  }
+  return first
+}
+
+@main struct Main {
+  static func main() {
+
+    let tests = TestSuite("Execute child task on await")
+    tests.test("Basic async let only execute on await") {
+      Task.runInline {
+         async let a = foo()
+         async let b = bar()
+
+         expectEqual(await b, 2)
+         expectEqual(await a, 3)
+         // Re-querying a completed task should return original result and not
+         // reexecute it
+         expectEqual(await b, 2)
+      }
+    }
+
+    tests.test("Nested async lets that only execute on await") {
+      Task.runInline {
+        let result = await asyncFib(10, pthread_self())
+        expectEqual(result, fib(10))
+      }
+    }
+
+    runAllTests()
+  }
+}