[Concurrency] Reduce overhead of Task.yield and Task.sleep (#37090)

* [Concurrency] Reduce overhead of Task.yield and Task.sleep

Instead of creating a new task, we create a simple job that wraps a Builtin.RawUnsafeContinuation and resumes the continuation when it is executed. The job instance is allocated on the task local allocator, meaning we don't malloc anything.

* Update stdlib/public/Concurrency/Task.swift

Co-authored-by: Konrad `ktoso` Malawski <[email protected]>

Co-authored-by: Konrad `ktoso` Malawski <[email protected]>

Author: drexin

include/swift/ABI/MetadataValues.h

@@ -1983,7 +1983,8 @@ enum class JobKind : size_t {
 
   DefaultActorInline = First_Reserved,
   DefaultActorSeparate,
-  DefaultActorOverride
+  DefaultActorOverride,
+  NullaryContinuation
 };
 
 /// The priority of a job.  Higher priorities are larger values.

include/swift/ABI/Task.h

@@ -166,6 +166,25 @@ class ActiveTaskStatus {
   }
 };
 
+class NullaryContinuationJob : public Job {
+
+private:
+  AsyncTask* Task;
+  AsyncTask* Continuation;
+
+public:
+  NullaryContinuationJob(AsyncTask *task, JobPriority priority, AsyncTask *continuation)
+    : Job({JobKind::NullaryContinuation, priority}, &process),
+      Task(task), Continuation(continuation) {}
+
+  SWIFT_CC(swiftasync)
+  static void process(Job *job);
+
+  static bool classof(const Job *job) {
+    return job->Flags.getKind() == JobKind::NullaryContinuation;
+  }
+};
+
 /// An asynchronous task.  Tasks are the analogue of threads for
 /// asynchronous functions: that is, they are a persistent identity
 /// for the overall async computation.
@@ -588,6 +607,10 @@ class ContinuationAsyncContext : public AsyncContext {
   /// The executor that should be resumed to.
   ExecutorRef ResumeToExecutor;
 
+  void setErrorResult(SwiftError *error) {
+    ErrorResult = error;
+  }
+
   static bool classof(const AsyncContext *context) {
     return context->Flags.getKind() == AsyncContextKind::Continuation;
   }

include/swift/Runtime/Concurrency.h

@@ -17,6 +17,7 @@
 #ifndef SWIFT_RUNTIME_CONCURRENCY_H
 #define SWIFT_RUNTIME_CONCURRENCY_H
 
+#include "swift/ABI/Task.h"
 #include "swift/ABI/TaskGroup.h"
 #include "swift/ABI/AsyncLet.h"
 #include "swift/ABI/TaskStatus.h"
@@ -419,6 +420,13 @@ SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_removeCancellationHandler(
     CancellationNotificationStatusRecord *record);
 
+/// Create a NullaryContinuationJob from a continuation.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+NullaryContinuationJob*
+swift_task_createNullaryContinuationJob(
+    size_t priority,
+    AsyncTask *continuation);
+
 /// Report error about attempting to bind a task-local value from an illegal context.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroup(

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -146,6 +146,11 @@ OVERRIDE_TASK(task_removeCancellationHandler, void,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
               (CancellationNotificationStatusRecord *record), (record))
 
+OVERRIDE_TASK(task_createNullaryContinuationJob, NullaryContinuationJob *,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
+              (size_t priority,
+               AsyncTask *continuation), (priority, continuation))
+
 OVERRIDE_TASK(task_asyncMainDrainQueue, void,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
               , )

stdlib/public/Concurrency/Task.cpp

@@ -96,6 +96,20 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask) {
   }
 }
 
+void NullaryContinuationJob::process(Job *_job) {
+  auto *job = cast<NullaryContinuationJob>(_job);
+
+  auto *task = job->Task;
+  auto *continuation = job->Continuation;
+
+  _swift_task_dealloc_specific(task, job);
+
+  auto *context = cast<ContinuationAsyncContext>(continuation->ResumeContext);
+
+  context->setErrorResult(nullptr);
+  swift_continuation_resume(continuation);
+}
+
 void AsyncTask::completeFuture(AsyncContext *context) {
   using Status = FutureFragment::Status;
   using WaitQueueItem = FutureFragment::WaitQueueItem;
@@ -962,6 +976,21 @@ static void swift_task_removeCancellationHandlerImpl(
   swift_task_dealloc(record);
 }
 
+SWIFT_CC(swift)
+static NullaryContinuationJob*
+swift_task_createNullaryContinuationJobImpl(
+    size_t priority,
+    AsyncTask *continuation) {
+  void *allocation =
+      swift_task_alloc(sizeof(NullaryContinuationJob));
+  auto *job =
+      new (allocation) NullaryContinuationJob(
+        swift_task_getCurrent(), static_cast<JobPriority>(priority),
+        continuation);
+
+  return job;
+}
+
 SWIFT_CC(swift)
 void swift::swift_continuation_logFailedCheck(const char *message) {
   swift_reportError(0, message);

stdlib/public/Concurrency/Task.swift

@@ -597,19 +597,13 @@ extension Task {
   ///
   /// This function does _not_ block the underlying thread.
   public static func sleep(_ duration: UInt64) async {
-    // Set up the job flags for a new task.
-    var flags = Task.JobFlags()
-    flags.kind = .task
-    flags.priority = .default
-    flags.isFuture = true
+    let currentTask = Builtin.getCurrentAsyncTask()
+    let priority = getJobFlags(currentTask).priority ?? Task.currentPriority._downgradeUserInteractive
 
-    // Create the asynchronous task future.
-    let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, {})
-
-    // Enqueue the resulting job.
-    _enqueueJobGlobalWithDelay(duration, Builtin.convertTaskToJob(task))
-
-    await Handle<Void, Never>(task).get()
+    return await Builtin.withUnsafeContinuation { (continuation: Builtin.RawUnsafeContinuation) -> Void in
+      let job = _taskCreateNullaryContinuationJob(priority: priority.rawValue, continuation: continuation)
+      _enqueueJobGlobalWithDelay(duration, job)
+    }
   }
 }
 
@@ -625,22 +619,13 @@ extension Task {
   /// if the task is the highest-priority task in the system, it might go
   /// immediately back to executing.
   public static func yield() async {
-    // Prepare the job flags
-    var flags = JobFlags()
-    flags.kind = .task
-    flags.priority = .default
-    flags.isFuture = true
-
-    // Create the asynchronous task future, it will do nothing, but simply serves
-    // as a way for us to yield our execution until the executor gets to it and
-    // resumes us.
-    // TODO: consider if it would be useful for this task to be a child task
-    let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, {})
-
-    // Enqueue the resulting job.
-    _enqueueJobGlobal(Builtin.convertTaskToJob(task))
-
-    let _ = await Handle<Void, Never>(task).get()
+    let currentTask = Builtin.getCurrentAsyncTask()
+    let priority = getJobFlags(currentTask).priority ?? Task.currentPriority._downgradeUserInteractive
+
+    return await Builtin.withUnsafeContinuation { (continuation: Builtin.RawUnsafeContinuation) -> Void in
+      let job = _taskCreateNullaryContinuationJob(priority: priority.rawValue, continuation: continuation)
+      _enqueueJobGlobal(job)
+    }
   }
 }
 
@@ -818,6 +803,10 @@ func _taskCancel(_ task: Builtin.NativeObject)
 @_silgen_name("swift_task_isCancelled")
 func _taskIsCancelled(_ task: Builtin.NativeObject) -> Bool
 
+@available(SwiftStdlib 5.5, *)
+@_silgen_name("swift_task_createNullaryContinuationJob")
+func _taskCreateNullaryContinuationJob(priority: Int, continuation: Builtin.RawUnsafeContinuation) -> Builtin.Job
+
 @available(SwiftStdlib 5.5, *)
 @usableFromInline
 @_silgen_name("swift_task_isCurrentExecutor")