Merge pull request #38386 from rjmccall/track-task-is-running

Track whether a task is actively running

Author: rjmccall

include/swift/ABI/MetadataValues.h

@@ -48,7 +48,7 @@ enum {
   NumWords_DefaultActor = 12,
 
   /// The number of words in a task.
-  NumWords_AsyncTask = 16,
+  NumWords_AsyncTask = 24,
 
   /// The number of words in a task group.
   NumWords_TaskGroup = 32,
@@ -2211,7 +2211,10 @@ class AsyncContextFlags : public FlagSet<uint32_t> {
     Kind_width          = 8,
 
     CanThrow            = 8,
-    ShouldNotDeallocate = 9
+
+    // Kind-specific flags should grow down from 31.
+
+    Continuation_IsExecutorSwitchForced = 31,
   };
 
   explicit AsyncContextFlags(uint32_t bits) : FlagSet(bits) {}
@@ -2227,17 +2230,10 @@ class AsyncContextFlags : public FlagSet<uint32_t> {
   /// Whether this context is permitted to throw.
   FLAGSET_DEFINE_FLAG_ACCESSORS(CanThrow, canThrow, setCanThrow)
 
-  /// Whether a function should avoid deallocating its context before
-  /// returning.  It should still pass its caller's context to its
-  /// return continuation.
-  ///
-  /// This flag can be set in the caller to optimize context allocation,
-  /// e.g. if the callee's context size is known statically and simply
-  /// allocated as part of the caller's context, or if the callee will
-  /// be called multiple times.
-  FLAGSET_DEFINE_FLAG_ACCESSORS(ShouldNotDeallocate,
-                                shouldNotDeallocateInCallee,
-                                setShouldNotDeallocateInCallee)
+  /// See AsyncContinuationFlags::isExecutorSwitchForced.
+  FLAGSET_DEFINE_FLAG_ACCESSORS(Continuation_IsExecutorSwitchForced,
+                                continuation_isExecutorSwitchForced,
+                                continuation_setIsExecutorSwitchForced)
 };
 
 /// Flags passed to swift_continuation_init.
@@ -2247,6 +2243,7 @@ class AsyncContinuationFlags : public FlagSet<size_t> {
     CanThrow            = 0,
     HasExecutorOverride = 1,
     IsPreawaited        = 2,
+    IsExecutorSwitchForced = 3,
   };
 
   explicit AsyncContinuationFlags(size_t bits) : FlagSet(bits) {}
@@ -2262,10 +2259,27 @@ class AsyncContinuationFlags : public FlagSet<size_t> {
                                 hasExecutorOverride,
                                 setHasExecutorOverride)
 
+  /// Whether the switch to the target executor should be forced
+  /// by swift_continuation_await.  If this is not set, and
+  /// swift_continuation_await finds that the continuation has
+  /// already been resumed, then execution will continue on the
+  /// current executor.  This has no effect in combination with
+  /// pre-awaiting.
+  ///
+  /// Setting this flag when you know statically that you're
+  /// already on the right executor is suboptimal.  In particular,
+  /// there's no good reason to set this if you're not also using
+  /// an executor override.
+  FLAGSET_DEFINE_FLAG_ACCESSORS(IsExecutorSwitchForced,
+                                isExecutorSwitchForced,
+                                setIsExecutorSwitchForced)
+
   /// Whether the continuation is "pre-awaited".  If so, it should
   /// be set up in the already-awaited state, and so resumptions
   /// will immediately schedule the continuation to begin
-  /// asynchronously.
+  /// asynchronously.  The continuation must not be subsequently
+  /// awaited if this is set.  The task is immediately treated as
+  /// suspended.
   FLAGSET_DEFINE_FLAG_ACCESSORS(IsPreawaited,
                                 isPreawaited,
                                 setIsPreawaited)

include/swift/ABI/Task.h

@@ -203,7 +203,7 @@ class AsyncTask : public Job {
 
   /// Private storage for the use of the runtime.
   struct alignas(2 * alignof(void*)) OpaquePrivateStorage {
-    void *Storage[6];
+    void *Storage[14];
 
     /// Initialize this storage during the creation of a task.
     void initialize(AsyncTask *task);
@@ -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;
@@ -632,6 +650,10 @@ class ContinuationAsyncContext : public AsyncContext {
     ErrorResult = error;
   }
 
+  bool isExecutorSwitchForced() const {
+    return Flags.continuation_isExecutorSwitchForced();
+  }
+
   static bool classof(const AsyncContext *context) {
     return context->Flags.getKind() == AsyncContextKind::Continuation;
   }

include/swift/Runtime/Atomic.h

@@ -53,10 +53,15 @@ class alignas(Size) atomic_impl {
   constexpr atomic_impl(Value value) : value(value) {}
 
   /// Force clients to always pass an order.
-  Value load(std::memory_order order) {
+  Value load(std::memory_order order) const {
     return value.load(order);
   }
 
+  /// Force clients to always pass an order.
+  void store(Value newValue, std::memory_order order) {
+    return value.store(newValue, order);
+  }
+
   /// Force clients to always pass an order.
   bool compare_exchange_weak(Value &oldValue, Value newValue,
                              std::memory_order successOrder,
@@ -75,14 +80,14 @@ class alignas(Size) atomic_impl {
 /// AMD processors that lack cmpxchg16b, so we just use the intrinsic.
 template <class Value>
 class alignas(2 * sizeof(void*)) atomic_impl<Value, 2 * sizeof(void*)> {
-  volatile Value atomicValue;
+  mutable volatile Value atomicValue;
 public:
   constexpr atomic_impl(Value initialValue) : atomicValue(initialValue) {}
 
   atomic_impl(const atomic_impl &) = delete;
   atomic_impl &operator=(const atomic_impl &) = delete;
 
-  Value load(std::memory_order order) {
+  Value load(std::memory_order order) const {
     assert(order == std::memory_order_relaxed ||
            order == std::memory_order_acquire ||
            order == std::memory_order_consume);
@@ -107,6 +112,17 @@ class alignas(2 * sizeof(void*)) atomic_impl<Value, 2 * sizeof(void*)> {
     return reinterpret_cast<Value &>(resultArray);
   }
 
+  void store(Value newValue, std::memory_order order) {
+    assert(order == std::memory_order_relaxed ||
+           order == std::memory_order_release);
+    Value oldValue = load(std::memory_order_relaxed);
+    while (!compare_exchange_weak(oldValue, newValue,
+                                  /*success*/ order,
+                                  /*failure*/ std::memory_order_relaxed)) {
+      // try again
+    }
+  }
+
   bool compare_exchange_weak(Value &oldValue, Value newValue,
                              std::memory_order successOrder,
                              std::memory_order failureOrder) {

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

@@ -78,10 +78,6 @@ OVERRIDE_ACTOR(job_run, void,
                swift::, (class Job *job, ExecutorRef executor),
                (job, executor))
 
-OVERRIDE_ACTOR(task_getCurrent, AsyncTask *,
-               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
-               swift::, ,)
-
 OVERRIDE_ACTOR(task_getCurrentExecutor, ExecutorRef,
                SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                swift::, ,)
@@ -159,6 +155,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,