Allow swift_continuation_await to be forced to honor an executor.

Author: rjmccall

include/swift/ABI/MetadataValues.h

@@ -2211,7 +2211,11 @@ class AsyncContextFlags : public FlagSet<uint32_t> {
     Kind_width          = 8,
 
     CanThrow            = 8,
-    ShouldNotDeallocate = 9
+    ShouldNotDeallocate = 9,
+
+    // Kind-specific flags should grow down from 31.
+
+    Continuation_IsExecutorSwitchForced = 31,
   };
 
   explicit AsyncContextFlags(uint32_t bits) : FlagSet(bits) {}
@@ -2238,6 +2242,11 @@ class AsyncContextFlags : public FlagSet<uint32_t> {
   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 +2256,7 @@ class AsyncContinuationFlags : public FlagSet<size_t> {
     CanThrow            = 0,
     HasExecutorOverride = 1,
     IsPreawaited        = 2,
+    IsExecutorSwitchForced = 3,
   };
 
   explicit AsyncContinuationFlags(size_t bits) : FlagSet(bits) {}
@@ -2262,10 +2272,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

@@ -650,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;
   }

stdlib/public/Concurrency/Task.cpp

@@ -948,6 +948,8 @@ static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
                                               AsyncContinuationFlags flags) {
   context->Flags = AsyncContextKind::Continuation;
   if (flags.canThrow()) context->Flags.setCanThrow(true);
+  if (flags.isExecutorSwitchForced())
+    context->Flags.continuation_setIsExecutorSwitchForced(true);
   context->ErrorResult = nullptr;
 
   // Set the current executor as the target executor unless there's
@@ -993,9 +995,12 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
 
   // 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.
+  if (oldStatus != ContinuationStatus::Pending) {
+    if (context->isExecutorSwitchForced())
+      return swift_task_switch(context, context->ResumeParent,
+                               context->ResumeToExecutor);
     return context->ResumeParent(context);
+  }
 
   // Load the current task (we alreaady did this in assertions builds).
 #ifdef NDEBUG
@@ -1021,7 +1026,10 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
 
   // Restore the running state of the task and resume it.
   task->flagAsRunning();
-  return task->runInFullyEstablishedContext();
+  if (context->isExecutorSwitchForced())
+    return swift_task_switch(context, context->ResumeParent,
+                             context->ResumeToExecutor);
+  return context->ResumeParent(context);
 }
 
 static void resumeTaskAfterContinuation(AsyncTask *task,