Track whether a task is actively running.

Tracking this as a single bit is actually largely uninteresting
to the runtime.  To handle priority escalation properly, we really
need to track this at a finer grain of detail: recording that the
task is running on a specific thread, enqueued on a specific actor,
or so on.  But starting by tracking a single bit is important for
two reasons:

- First, it's more realistic about the performance overheads of
  tasks: we're going to be doing this tracking eventually, and
  the cost of that tracking will be dominated by the atomic
  access, so doing that access now sets the baseline about right.

- Second, it ensures that we've actually got runtime involvement
  in all the right places to do this tracking.

A propos of the latter: there was no runtime involvement with
awaiting a continuation, which is a point at which the task
potentially transitions from running to suspended.  We must do
the tracking as part of this transition, rather than recognizing
in the run-loops that a task is still active and treating it as
having suspended, because the latter point potentially races with
the resumption of the task.  To do this, I've had to introduce
a runtime function, swift_continuation_await, to do this awaiting
rather than inlining the atomic operation on the continuation.

As part of doing this work, I've also fixed a bug where we failed
to load-acquire in swift_task_escalate before walking the task
status records to invoke escalation actions.

I've also fixed several places where the handling of task statuses
may have accidentally allowed the task to revert to uncancelled.

Author: rjmccall

include/swift/ABI/Task.h

@@ -260,7 +260,25 @@ class AsyncTask : public Job {
   void runInFullyEstablishedContext() {
     return ResumeTask(ResumeContext); // 'return' forces tail call
   }
-  
+
+  /// Flag that this task is now running.  This can update
+  /// the priority stored in the job flags if the priority has been
+  /// escalated.
+  ///
+  /// Generally this should be done immediately after updating
+  /// ActiveTask.
+  void flagAsRunning();
+  void flagAsRunning_slow();
+
+  /// Flag that this task is now suspended.  This can update the
+  /// priority stored in the job flags if the priority hsa been
+  /// escalated.  Generally this should be done immediately after
+  /// clearing ActiveTask and immediately before enqueuing the task
+  /// somewhere.  TODO: record where the task is enqueued if
+  /// possible.
+  void flagAsSuspended();
+  void flagAsSuspended_slow();
+
   /// Check whether this task has been cancelled.
   /// Checking this is, of course, inherently race-prone on its own.
   bool isCancelled() const;

include/swift/Runtime/Concurrency.h

@@ -605,20 +605,28 @@ void swift_defaultActor_enqueue(Job *job, DefaultActor *actor);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_distributed_actor_is_remote(DefaultActor *actor);
 
+/// Do a primitive suspension of the current task, as if part of
+/// a continuation, although this does not provide any of the
+/// higher-level continuation semantics.  The current task is returned;
+/// its ResumeFunction and ResumeContext will need to be initialized,
+/// and then it will need to be enqueued or run as a job later.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+AsyncTask *swift_task_suspend();
+
 /// Prepare a continuation in the current task.
 ///
 /// The caller should initialize the Parent, ResumeParent,
 /// and NormalResult fields.  This function will initialize the other
-/// fields with appropriate defaaults; the caller may then overwrite
+/// fields with appropriate defaults; the caller may then overwrite
 /// them if desired.
-///
-/// This function is provided as a code-size and runtime-usage
-/// optimization; calling it is not required if code is willing to
-/// do all its work inline.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 AsyncTask *swift_continuation_init(ContinuationAsyncContext *context,
                                    AsyncContinuationFlags flags);
 
+/// Await an initialized continuation.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
+void swift_continuation_await(ContinuationAsyncContext *continuationContext);
+
 /// Resume a task from a non-throwing continuation, given a normal
 /// result which has already been stored into the continuation.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)

include/swift/Runtime/RuntimeFunctions.def

@@ -1582,6 +1582,14 @@ FUNCTION(ContinuationInit,
          ARGS(ContinuationAsyncContextPtrTy, SizeTy),
          ATTRS(NoUnwind))
 
+// void swift_continuation_await(AsyncContext *continuationContext);
+FUNCTION(ContinuationAwait,
+         swift_continuation_await, SwiftAsyncCC,
+         ConcurrencyAvailability,
+         RETURNS(VoidTy),
+         ARGS(ContinuationAsyncContextPtrTy),
+         ATTRS(NoUnwind))
+
 // void swift_continuation_resume(AsyncTask *continuation);
 FUNCTION(ContinuationResume,
          swift_continuation_resume, SwiftCC,

lib/IRGen/GenConcurrency.cpp

@@ -21,6 +21,7 @@
 #include "ExtraInhabitants.h"
 #include "GenProto.h"
 #include "GenType.h"
+#include "IRGenDebugInfo.h"
 #include "IRGenFunction.h"
 #include "IRGenModule.h"
 #include "LoadableTypeInfo.h"
@@ -241,3 +242,34 @@ void irgen::emitDestroyTaskGroup(IRGenFunction &IGF, llvm::Value *group) {
 
   IGF.Builder.CreateLifetimeEnd(group);
 }
+
+llvm::Function *IRGenModule::getAwaitAsyncContinuationFn() {
+  StringRef name = "__swift_continuation_await_point";
+  if (llvm::GlobalValue *F = Module.getNamedValue(name))
+    return cast<llvm::Function>(F);
+
+  // The parameters here match the extra arguments passed to
+  // @llvm.coro.suspend.async by emitAwaitAsyncContinuation.
+  llvm::Type *argTys[] = { ContinuationAsyncContextPtrTy };
+  auto *suspendFnTy =
+    llvm::FunctionType::get(VoidTy, argTys, false /*vaargs*/);
+
+  llvm::Function *suspendFn =
+      llvm::Function::Create(suspendFnTy, llvm::Function::InternalLinkage,
+                             name, &Module);
+  suspendFn->setCallingConv(SwiftAsyncCC);
+  suspendFn->setDoesNotThrow();
+  IRGenFunction suspendIGF(*this, suspendFn);
+  if (DebugInfo)
+    DebugInfo->emitArtificialFunction(suspendIGF, suspendFn);
+  auto &Builder = suspendIGF.Builder;
+
+  llvm::Value *context = suspendFn->getArg(0);
+  auto *call = Builder.CreateCall(getContinuationAwaitFn(), { context });
+  call->setDoesNotThrow();
+  call->setCallingConv(SwiftAsyncCC);
+  call->setTailCallKind(AsyncTailCallKind);
+
+  Builder.CreateRetVoid();
+  return suspendFn;
+}

lib/IRGen/IRGenFunction.cpp

@@ -665,50 +665,8 @@ void IRGenFunction::emitAwaitAsyncContinuation(
   assert(AsyncCoroutineCurrentContinuationContext && "no active continuation");
   auto pointerAlignment = IGM.getPointerAlignment();
 
-  // Check whether the continuation has already been resumed.
-  // If so, we can just immediately continue with the control flow.
-  // Otherwise, we need to suspend, and resuming the continuation will
-  // trigger the function to resume.
-  //
-  // We do this by atomically trying to change the synchronization field
-  // in the continuation context from 0 (the state it was initialized
-  // with) to 1.  If this fails, the continuation must already have been
-  // resumed, so we can bypass the suspension point and immediately
-  // start interpreting the result stored in the continuation.
-  // Note that we use a strong compare-exchange (the default for the LLVM
-  // cmpxchg instruction), so spurious failures are disallowed; we can
-  // therefore trust that a failure means that the continuation has
-  // already been resumed.
-
-  auto contAwaitSyncAddr =
-      Builder.CreateStructGEP(AsyncCoroutineCurrentContinuationContext, 1);
-
-  auto pendingV = llvm::ConstantInt::get(
-      contAwaitSyncAddr->getType()->getPointerElementType(),
-      unsigned(ContinuationStatus::Pending));
-  auto awaitedV = llvm::ConstantInt::get(
-      contAwaitSyncAddr->getType()->getPointerElementType(),
-      unsigned(ContinuationStatus::Awaited));
-  auto results = Builder.CreateAtomicCmpXchg(
-      contAwaitSyncAddr, pendingV, awaitedV,
-      llvm::AtomicOrdering::Release /*success ordering*/,
-      llvm::AtomicOrdering::Acquire /* failure ordering */,
-      llvm::SyncScope::System);
-  auto firstAtAwait = Builder.CreateExtractValue(results, 1);
-  auto contBB = createBasicBlock("await.async.resume");
-  auto abortBB = createBasicBlock("await.async.abort");
-  Builder.CreateCondBr(firstAtAwait, abortBB, contBB);
-  Builder.emitBlock(abortBB);
-  {
-    // We were the first to the sync point. "Abort" (return from the
-    // coroutine partial function, without making a tail call to anything)
-    // because the continuation result is not available yet. When the
-    // continuation is later resumed, the task will get scheduled
-    // starting from the suspension point.
-    emitCoroutineOrAsyncExit();
-  }
-
-  Builder.emitBlock(contBB);
+  // Call swift_continuation_await to check whether the continuation
+  // has already been resumed.
   {
     // Set up the suspend point.
     SmallVector<llvm::Value *, 8> arguments;
@@ -718,15 +676,10 @@ void IRGenFunction::emitAwaitAsyncContinuation(
     auto resumeProjFn = getOrCreateResumePrjFn();
     arguments.push_back(
         Builder.CreateBitOrPointerCast(resumeProjFn, IGM.Int8PtrTy));
-    // The dispatch function just calls the resume point.
-    auto resumeFnPtr =
-        getFunctionPointerForResumeIntrinsic(AsyncCoroutineCurrentResume);
     arguments.push_back(Builder.CreateBitOrPointerCast(
-        createAsyncDispatchFn(resumeFnPtr, {IGM.Int8PtrTy}),
+        IGM.getAwaitAsyncContinuationFn(),
         IGM.Int8PtrTy));
-    arguments.push_back(AsyncCoroutineCurrentResume);
-    arguments.push_back(Builder.CreateBitOrPointerCast(
-        AsyncCoroutineCurrentContinuationContext, IGM.Int8PtrTy));
+    arguments.push_back(AsyncCoroutineCurrentContinuationContext);
     auto resultTy =
         llvm::StructType::get(IGM.getLLVMContext(), {IGM.Int8PtrTy}, false /*packed*/);
     emitSuspendAsyncCall(swiftAsyncContextIndex, resultTy, arguments);

lib/IRGen/IRGenModule.h

@@ -1330,6 +1330,7 @@ private:                            \
   llvm::Constant *getFixLifetimeFn();
 
   llvm::Constant *getFixedClassInitializationFn();
+  llvm::Function *getAwaitAsyncContinuationFn();
 
   /// The constructor used when generating code.
   ///

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -159,6 +159,21 @@ OVERRIDE_TASK(task_asyncMainDrainQueue, void,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
               , )
 
+OVERRIDE_TASK(task_suspend, AsyncTask *,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
+              swift::, ,)
+
+OVERRIDE_TASK(continuation_init, AsyncTask *,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
+              swift::, (ContinuationAsyncContext *context,
+                        AsyncContinuationFlags flags),
+              (context, flags))
+
+OVERRIDE_TASK(continuation_await, void,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swiftasync),
+              swift::, (ContinuationAsyncContext *context),
+              (context))
+
 OVERRIDE_ASYNC_LET(asyncLet_wait, void, SWIFT_EXPORT_FROM(swift_Concurrency),
                    SWIFT_CC(swiftasync), swift::,
                    (OpaqueValue *result,

stdlib/public/Concurrency/Actor.cpp

@@ -231,10 +231,11 @@ void swift::runJobInEstablishedExecutorContext(Job *job) {
     // Update the active task in the current thread.
     ActiveTask::set(task);
 
-    // FIXME: update the task status to say that it's running
-    // on the current thread.  If the task suspends itself to run
-    // on an actor, it should update the task status appropriately;
-    // we don't need to update it afterwards.
+    // Update the task status to say that it's running on the
+    // current thread.  If the task suspends somewhere, it should
+    // update the task status appropriately; we don't need to update
+    // it afterwards.
+    task->flagAsRunning();
 
     task->runInFullyEstablishedContext();
 
@@ -1845,6 +1846,9 @@ SWIFT_CC(swiftasync)
 static void runOnAssumedThread(AsyncTask *task, ExecutorRef executor,
                                ExecutorTrackingInfo *oldTracking,
                                RunningJobInfo runner) {
+  // Note that this doesn't change the active task and so doesn't
+  // need to either update ActiveTask or flagAsRunning/flagAsSuspended.
+
   // If there's alreaady tracking info set up, just change the executor
   // there and tail-call the task.  We don't want these frames to
   // potentially accumulate linearly.
@@ -1932,6 +1936,7 @@ static void swift_task_switchImpl(SWIFT_ASYNC_CONTEXT AsyncContext *resumeContex
   fprintf(stderr, "[%lu] switch failed, task %p enqueued on executor %p\n",
           _swift_get_thread_id(), task, newExecutor.getIdentity());
 #endif
+  task->flagAsSuspended();
   swift_task_enqueue(task, newExecutor);
 }
 

stdlib/public/Concurrency/Task.cpp

@@ -89,6 +89,7 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
               _swift_get_thread_id(), waitingTask, this);
 #endif
       _swift_tsan_acquire(static_cast<Job *>(this));
+      if (contextIntialized) waitingTask->flagAsRunning();
       // The task is done; we don't need to wait.
       return queueHead.getStatus();
 
@@ -98,7 +99,7 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
               _swift_get_thread_id(), waitingTask, this);
 #endif
       _swift_tsan_release(static_cast<Job *>(waitingTask));
-      // Task is now complete. We'll need to add ourselves to the queue.
+      // Task is not complete. We'll need to add ourselves to the queue.
       break;
     }
 
@@ -110,6 +111,7 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
       context->successResultPointer = result;
       context->ResumeParent = resumeFn;
       context->Parent = callerContext;
+      waitingTask->flagAsSuspended();
     }
 
     // Put the waiting task at the beginning of the wait queue.
@@ -934,8 +936,16 @@ size_t swift::swift_task_getJobFlags(AsyncTask *task) {
   return task->Flags.getOpaqueValue();
 }
 
-AsyncTask *swift::swift_continuation_init(ContinuationAsyncContext *context,
-                                          AsyncContinuationFlags flags) {
+SWIFT_CC(swift)
+static AsyncTask *swift_task_suspendImpl() {
+  auto task = swift_task_getCurrent();
+  task->flagAsSuspended();
+  return task;
+}
+
+SWIFT_CC(swift)
+static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
+                                              AsyncContinuationFlags flags) {
   context->Flags = AsyncContextKind::Continuation;
   if (flags.canThrow()) context->Flags.setCanThrow(true);
   context->ErrorResult = nullptr;
@@ -957,9 +967,63 @@ AsyncTask *swift::swift_continuation_init(ContinuationAsyncContext *context,
   task->ResumeContext = context;
   task->ResumeTask = context->ResumeParent;
 
+  // A preawait immediately suspends the task.
+  if (flags.isPreawaited()) {
+    task->flagAsSuspended();
+  }
+
   return task;
 }
 
+SWIFT_CC(swiftasync)
+static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
+#ifndef NDEBUG
+  auto task = swift_task_getCurrent();
+  assert(task && "awaiting continuation without a task");
+  assert(task->ResumeContext == context);
+  assert(task->ResumeTask == context->ResumeParent);
+#endif
+
+  auto &sync = context->AwaitSynchronization;
+
+  auto oldStatus = sync.load(std::memory_order_acquire);
+  assert((oldStatus == ContinuationStatus::Pending ||
+          oldStatus == ContinuationStatus::Resumed) &&
+         "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)
+    // This is fine given how swift_continuation_init sets it up.
+    return context->ResumeParent(context);
+
+  // Load the current task (we alreaady did this in assertions builds).
+#ifdef NDEBUG
+  auto task = swift_task_getCurrent();
+#endif
+
+  // Flag the task as suspended.
+  task->flagAsSuspended();
+
+  // 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) return;
+
+  // 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");
+
+  // Restore the running state of the task and resume it.
+  task->flagAsRunning();
+  return task->runInFullyEstablishedContext();
+}
+
 static void resumeTaskAfterContinuation(AsyncTask *task,
                                         ContinuationAsyncContext *context) {
   auto &sync = context->AwaitSynchronization;