[Concurrency] Add asynchronous Task.sleep function

Author: drexin

include/swift/Runtime/Concurrency.h

@@ -486,6 +486,9 @@ void swift_task_enqueue(Job *job, ExecutorRef executor);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_enqueueGlobal(Job *job);
 
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_enqueueGlobalWithDelay(unsigned long long delay, Job *job);
+
 /// FIXME: only exists for the quick-and-dirty MainActor implementation.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_enqueueMainExecutor(Job *job);
@@ -499,6 +502,11 @@ void swift_MainActor_register(HeapObject *actor);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void (*swift_task_enqueueGlobal_hook)(Job *job);
 
+/// A hook to take over global enqueuing with delay.
+/// TODO: figure out a better abstraction plan than this.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void (*swift_task_enqueueGlobalWithDelay_hook)(unsigned long long delay, Job *job);
+
 /// Initialize the runtime storage for a default actor.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_defaultActor_initialize(DefaultActor *actor);

stdlib/public/Concurrency/GlobalExecutor.cpp

@@ -63,9 +63,27 @@ using namespace swift;
 SWIFT_CC(swift)
 void (*swift::swift_task_enqueueGlobal_hook)(Job *job) = nullptr;
 
+SWIFT_CC(swift)
+void (*swift::swift_task_enqueueGlobalWithDelay_hook)(unsigned long long delay, Job *job) = nullptr;
+
 #if SWIFT_CONCURRENCY_COOPERATIVE_GLOBAL_EXECUTOR
+
+#include <chrono>
+#include <thread>
+
 static Job *JobQueue = nullptr;
 
+class DelayedJob {
+public:
+  Job *job;
+  unsigned long long when;
+  DelayedJob *next;
+
+  DelayedJob(Job *job, unsigned long long when) : job(job), when(when), next(nullptr) {}
+};
+
+static DelayedJob *DelayedJobQueue = nullptr;
+
 /// Get the next-in-queue storage slot.
 static Job *&nextInQueue(Job *cur) {
   return reinterpret_cast<Job*&>(cur->SchedulerPrivate);
@@ -89,13 +107,58 @@ static void insertIntoJobQueue(Job *newJob) {
   *position = newJob;
 }
 
+static unsigned long long currentNanos() {
+  auto now = std::chrono::steady_clock::now();
+  auto nowNanos = std::chrono::time_point_cast<std::chrono::nanoseconds>(now);
+  auto value = std::chrono::duration_cast<std::chrono::nanoseconds>(nowNanos.time_since_epoch());
+  return value.count();
+}
+
+/// Insert a job into the cooperative global queue.
+static void insertIntoDelayedJobQueue(unsigned long long delay, Job *job) {
+  DelayedJob **position = &DelayedJobQueue;
+  DelayedJob *newJob = new DelayedJob(job, currentNanos() + delay);
+
+  while (auto cur = *position) {
+    // If we find a job with lower priority, insert here.
+    if (cur->when > newJob->when) {
+      newJob->next = cur;
+      *position = newJob;
+      return;
+    }
+
+    // Otherwise, keep advancing through the queue.
+    position = &cur->next;
+  }
+  *position = newJob;
+}
+
 /// Claim the next job from the cooperative global queue.
 static Job *claimNextFromJobQueue() {
-  if (auto job = JobQueue) {
-    JobQueue = nextInQueue(job);
-    return job;
+  // Check delayed jobs first
+  while (true) {
+    if (auto delayedJob = DelayedJobQueue) {
+      if (delayedJob->when < currentNanos()) {
+        DelayedJobQueue = delayedJob->next;
+        auto job = delayedJob->job;
+        
+        delete delayedJob;
+
+        return job;
+      }
+    }
+    if (auto job = JobQueue) {
+      JobQueue = nextInQueue(job);
+      return job;
+    }
+    // there are only delayed jobs left, but they are not ready,
+    // so we sleep until the first one is
+    if (auto delayedJob = DelayedJobQueue) {
+      std::this_thread::sleep_for(std::chrono::nanoseconds(delayedJob->when - currentNanos()));
+      continue;
+    }
+    return nullptr;
   }
-  return nullptr;
 }
 
 void swift::donateThreadToGlobalExecutorUntil(bool (*condition)(void *),
@@ -177,6 +240,30 @@ void swift::swift_task_enqueueGlobal(Job *job) {
 #endif
 }
 
+void swift::swift_task_enqueueGlobalWithDelay(unsigned long long delay, Job *job) {
+  assert(job && "no job provided");
+
+  // If the hook is defined, use it.
+  if (swift_task_enqueueGlobalWithDelay_hook)
+    return swift_task_enqueueGlobalWithDelay_hook(delay, job);
+
+#if SWIFT_CONCURRENCY_COOPERATIVE_GLOBAL_EXECUTOR
+  insertIntoDelayedJobQueue(delay, job);
+#else
+
+  dispatch_function_t dispatchFunction = &__swift_run_job;
+  void *dispatchContext = job;
+
+  JobPriority priority = job->getPriority();
+
+  // TODO: cache this to avoid the extra call
+  auto queue = dispatch_get_global_queue((dispatch_qos_class_t) priority,
+                                         /*flags*/ 0);
+  dispatch_time_t when = dispatch_time(DISPATCH_TIME_NOW, delay);
+  dispatch_after_f(when, queue, dispatchContext, dispatchFunction);
+#endif
+}
+
 
 /// Enqueues a task on the main executor.
 /// FIXME: only exists for the quick-and-dirty MainActor implementation.

stdlib/public/Concurrency/Task.swift

@@ -472,6 +472,27 @@ public func _runAsyncHandler(operation: @escaping () async -> ()) {
   )
 }
 
+
+  /// Suspends the current task for _at least_ the given duration
+  /// in nanoseconds.
+  ///
+  /// 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 = JobFlags()
+    flags.kind = .task
+    flags.priority = .default
+    flags.isFuture = true
+
+    // Create the asynchronous task future.
+    let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, nil, {})
+
+    // Enqueue the resulting job.
+    _enqueueJobGlobalWithDelay(duration, Builtin.convertTaskToJob(task))
+
+    let _ = await Handle<Void, Never>(task).get()
+  }
 // ==== UnsafeCurrentTask ------------------------------------------------------
 
 extension Task {
@@ -571,6 +592,10 @@ func getJobFlags(_ task: Builtin.NativeObject) -> Task.JobFlags
 @usableFromInline
 func _enqueueJobGlobal(_ task: Builtin.Job)
 
+@_silgen_name("swift_task_enqueueGlobalWithDelay")
+@usableFromInline
+func _enqueueJobGlobalWithDelay(_ delay: UInt64, _ task: Builtin.Job)
+
 @available(*, deprecated)
 @_silgen_name("swift_task_runAndBlockThread")
 public func runAsyncAndBlock(_ asyncFun: @escaping () async -> ())

test/Concurrency/Runtime/async_task_sleep.swift

@@ -0,0 +1,40 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library)
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+
+import _Concurrency
+// FIXME: should not depend on Dispatch
+import Dispatch
+
+@main struct Main {
+  static let pause = 500_000_000 // 500ms
+  
+  static func main() async {
+    await testSleepDuration()
+    await testSleepDoesNotBlock()
+  }
+
+  static func testSleepDuration() async {
+    let start = DispatchTime.now()
+
+    await Task.sleep(UInt64(pause))
+
+    let stop = DispatchTime.now()
+
+    // assert that at least the specified time passed since calling `sleep`
+    assert(stop >= (start + .nanoseconds(pause)))
+  }
+
+  static func testSleepDoesNotBlock() async {
+    // FIXME: Should run on main executor
+    let task = Task.runDetached {
+      print("Run first")
+    }
+
+    await Task.sleep(UInt64(pause))
+
+    print("Run second")
+
+    await task.get()
+  }
+}