Merge pull request #61911 from apple/rokhinip/99977665-continuations-support

Implement continuations in task-to-thread model.

Author: swift-ci

include/swift/ABI/Task.h

@@ -25,6 +25,8 @@
 #include "swift/Runtime/Config.h"
 #include "swift/Runtime/VoucherShims.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/Threading/ConditionVariable.h"
+#include "swift/Threading/Mutex.h"
 #include "bitset"
 #include "queue" // TODO: remove and replace with our own mpsc
 
@@ -711,6 +713,16 @@ class ContinuationAsyncContext : public AsyncContext {
   /// Public ABI.
   ExecutorRef ResumeToExecutor;
 
+#if defined(SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY)
+  /// In a task-to-thread model, instead of voluntarily descheduling the task
+  /// from the thread, we will block the thread (and therefore task).
+  /// This condition variable is lazily allocated on the stack only if the
+  /// continuation has not been resumed by the point of await. The mutex in the
+  /// condition variable is therefore not really protecting any state as all
+  /// coordination is done via the AwaitSynchronization atomic
+  ConditionVariable *Cond;
+#endif
+
   void setErrorResult(SwiftError *error) {
     ErrorResult = error;
   }

lib/IRGen/IRGenModule.cpp

@@ -683,14 +683,28 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
       *this, "swift.async_task_and_context",
       { SwiftTaskPtrTy, SwiftContextPtrTy });
 
-  ContinuationAsyncContextTy = createStructType(
-      *this, "swift.continuation_context",
-      {SwiftContextTy,       // AsyncContext header
-       SizeTy,               // flags
-       SizeTy,               // await synchronization
-       ErrorPtrTy,           // error result pointer
-       OpaquePtrTy,          // normal result address
-       SwiftExecutorTy});    // resume to executor
+  if (Context.LangOpts.isConcurrencyModelTaskToThread()) {
+    ContinuationAsyncContextTy = createStructType(
+        *this, "swift.continuation_context",
+        {SwiftContextTy,       // AsyncContext header
+         SizeTy,               // flags
+         SizeTy,               // await synchronization
+         ErrorPtrTy,           // error result pointer
+         OpaquePtrTy,          // normal result address
+         SwiftExecutorTy,      // resume to executor
+         SizeTy                // pointer to condition variable
+         });
+  } else {
+    ContinuationAsyncContextTy = createStructType(
+        *this, "swift.continuation_context",
+        {SwiftContextTy,       // AsyncContext header
+         SizeTy,               // flags
+         SizeTy,               // await synchronization
+         ErrorPtrTy,           // error result pointer
+         OpaquePtrTy,          // normal result address
+         SwiftExecutorTy       // resume to executor
+         });
+  }
   ContinuationAsyncContextPtrTy =
     ContinuationAsyncContextTy->getPointerTo(DefaultAS);
 

stdlib/public/Concurrency/Actor.cpp

@@ -845,7 +845,7 @@ class DefaultActorImpl : public HeapObject {
   // the future
   alignas(sizeof(ActiveActorStatus)) char StatusStorage[sizeof(ActiveActorStatus)];
 #endif
-  // TODO(rokhinip): Make this a flagset
+  // TODO (rokhinip): Make this a flagset
   bool isDistributedRemoteActor;
 
 public:

stdlib/public/Concurrency/Task.cpp

@@ -1311,7 +1311,9 @@ static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
                                         ? ContinuationStatus::Awaited
                                         : ContinuationStatus::Pending,
                                       std::memory_order_relaxed);
-
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  context->Cond = nullptr;
+#endif
   AsyncTask *task;
 
   // A preawait immediately suspends the task.
@@ -1351,8 +1353,7 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
          "awaiting a corrupt or already-awaited continuation");
 
   // If the status is already Resumed, we can resume immediately.
-  // Comparing against Pending may be very slightly more compact.
-  if (oldStatus != ContinuationStatus::Pending) {
+  if (oldStatus == ContinuationStatus::Resumed) {
     if (context->isExecutorSwitchForced())
       return swift_task_switch(context, context->ResumeParent,
                                context->ResumeToExecutor);
@@ -1364,28 +1365,72 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
   auto task = swift_task_getCurrent();
 #endif
 
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  // In the task to thread model, we do not suspend the task that is waiting on
+  // the continuation resumption. Instead we simply block the thread on a
+  // condition variable keep the task alive on the thread.
+  //
+  // This condition variable can be allocated on the stack of the blocking
+  // thread - with the address of it published to the resuming thread via the
+  // context.
+  ConditionVariable Cond;
+
+  context->Cond = &Cond;
+#else /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
   // Flag the task as suspended.
   task->flagAsSuspended();
+#endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
 
-  // Try to transition to Awaited.
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  // If the cmpxchg is successful, the store release also publishes the write to
+  // the Cond in the ContinuationAsyncContext to any concurrent accessing
+  // thread.
+  //
+  // If it failed, then someone concurrently resumed the continuation in which
+  // case, we don't care about publishing the Cond in the
+  // ContinuationAsyncContext anyway.
+#endif
+  // Try to transition to Awaited
   bool success =
     sync.compare_exchange_strong(oldStatus, ContinuationStatus::Awaited,
                                  /*success*/ std::memory_order_release,
                                  /*failure*/ std::memory_order_acquire);
 
-  // If that succeeded, we have nothing to do.
   if (success) {
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+    // This lock really protects nothing but we need to hold it
+    // while calling the condition wait
+    Cond.lock();
+
+    // Condition variables can have spurious wakeups so we need to check this in
+    // a do-while loop.
+    do {
+      Cond.wait();
+      oldStatus = sync.load(std::memory_order_relaxed);
+    } while (oldStatus != ContinuationStatus::Resumed);
+
+    Cond.unlock();
+#else
+    // If that succeeded, we have nothing to do since we've successfully
+    // suspended the task
     _swift_task_clearCurrent();
     return;
+#endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
   }
 
   // If it failed, it should be because someone concurrently resumed
   // (note that the compare-exchange above is strong).
   assert(oldStatus == ContinuationStatus::Resumed &&
          "continuation was concurrently corrupted or awaited");
 
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  // Since the condition variable is stack allocated, we don't need to do
+  // anything here to clean up
+#else
   // Restore the running state of the task and resume it.
   task->flagAsRunning();
+#endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
+
   if (context->isExecutorSwitchForced())
     return swift_task_switch(context, context->ResumeParent,
                              context->ResumeToExecutor);
@@ -1397,6 +1442,7 @@ static void resumeTaskAfterContinuation(AsyncTask *task,
   continuationChecking::willResume(context);
 
   auto &sync = context->AwaitSynchronization;
+
   auto status = sync.load(std::memory_order_acquire);
   assert(status != ContinuationStatus::Resumed &&
          "continuation was already resumed");
@@ -1405,27 +1451,41 @@ static void resumeTaskAfterContinuation(AsyncTask *task,
   // restarting.
   _swift_tsan_release(static_cast<Job *>(task));
 
-  // The status should be either Pending or Awaited.  If it's Awaited,
-  // which is probably the most likely option, then we should immediately
-  // enqueue; we don't need to update the state because there shouldn't
-  // be a racing attempt to resume the continuation.  If it's Pending,
-  // we need to set it to Resumed; if that fails (with a strong cmpxchg),
-  // it should be because the original thread concurrently set it to
-  // Awaited, and so we need to enqueue.
+  // The status should be either Pending or Awaited.
+  //
+  // Case 1: Status is Pending
+  // No one has awaited us, we just need to set it to Resumed; if that fails
+  // (with a strong cmpxchg), it should be because the original thread
+  // concurrently set it to Awaited, in which case, we fall into Case 2.
+  //
+  // Case 2: Status is Awaited
+  // This is probably the more frequently hit case.
+  // In task-to-thread model, we update status to be Resumed and signal the
+  // waiting thread. In regular model, we immediately enqueue the task and can
+  // skip updates to the continuation state since there shouldn't be a racing
+  // attempt to resume the continuation.
   if (status == ContinuationStatus::Pending &&
       sync.compare_exchange_strong(status, ContinuationStatus::Resumed,
                                    /*success*/ std::memory_order_release,
-                                   /*failure*/ std::memory_order_relaxed)) {
+                                   /*failure*/ std::memory_order_acquire)) {
     return;
   }
   assert(status == ContinuationStatus::Awaited &&
          "detected concurrent attempt to resume continuation");
+#if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
+  // If we see status == ContinuationStatus::Awaited, then we should also be
+  // seeing a pointer to the cond var since we're doing a load acquire on sync
+  // which pairs with the store release in swift_continuation_awaitImpl
+  assert(context->Cond != nullptr);
 
+  sync.store(ContinuationStatus::Resumed, std::memory_order_relaxed);
+  context->Cond->signal();
+#else
   // TODO: maybe in some mode we should set the status to Resumed here
   // to make a stronger best-effort attempt to catch racing attempts to
   // resume the continuation?
-
   task->flagAsAndEnqueueOnExecutor(context->ResumeToExecutor);
+#endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
 }
 
 SWIFT_CC(swift)

test/Concurrency/blocking_continuations.swift

@@ -0,0 +1,66 @@
+// RUN: %target-run-simple-swift(-parse-as-library -Xfrontend -disable-availability-checking -Xfrontend -concurrency-model=task-to-thread -g -Xlinker -object_path_lto -Xlinker /tmp/abc.o)
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// REQUIRES: concurrency_runtime
+// REQUIRES: freestanding
+// UNSUPPORTED: threading_none
+
+@_spi(_TaskToThreadModel) import _Concurrency
+import StdlibUnittest
+import Darwin
+
+var globalContinuation : CheckedContinuation<Int, Never>? = nil
+
+func waitOnContinuation(_unused : UnsafeMutableRawPointer) -> UnsafeMutableRawPointer? {
+    Task.runInline {
+      let result = await withCheckedContinuation { continuation in
+          globalContinuation = continuation
+      }
+      print("Continuation successfully resumed")
+      expectEqual(result, 10)
+    }
+    return nil
+}
+
+func resumeContinuation(_unused : UnsafeMutableRawPointer) -> UnsafeMutableRawPointer? {
+    Task.runInline {
+      while (globalContinuation == nil) {}
+      globalContinuation!.resume(returning: 10)
+    }
+    return nil
+}
+
+@main struct Main {
+  static func main() {
+
+    let tests = TestSuite("Continuations in task-to-thread")
+    tests.test("Basic continuations - no blocking") {
+      Task.runInline {
+        await withCheckedContinuation { continuation in
+          continuation.resume()
+        }
+      }
+    }
+
+    tests.test("Continuations - with blocking") {
+      var thread1 : pthread_t? = nil
+      guard pthread_create(&thread1, nil, waitOnContinuation, nil) == 0 else {
+        fatalError("pthread_create failed")
+      }
+
+      var thread2 : pthread_t? = nil
+      guard pthread_create(&thread2, nil, resumeContinuation, nil) == 0 else {
+        fatalError("pthread_create failed")
+      }
+
+      guard pthread_join(thread1!, nil) == 0 else {
+        fatalError("pthread_join failed")
+      }
+      guard pthread_join(thread2!, nil) == 0 else {
+        fatalError("pthread_join failed")
+      }
+    }
+
+    runAllTests()
+  }
+}