[Concurrency] Implement Task.Handle.get() in terms of an async runtime call.

Switch the contract between the runtime operation `swift_future_task_wait`
and Task.Handle.get() pver to an asynchronous call, so that the
compiler will set up the resumption frame for us. This allows us to
correctly wait on futures.

Update our "basic" future test to perform both normal returns and
throwing returns from a future, either having to wait on the queue or
coming by afterward.

Author: DougGregor

include/swift/Runtime/Concurrency.h

@@ -110,31 +110,31 @@ swift_task_escalate(AsyncTask *task, JobPriority newPriority);
 
 /// The result of waiting for a task future.
 struct TaskFutureWaitResult {
-  enum Kind : uintptr_t {
-    /// The waiting task has been added to the future's wait queue, and will
-    /// be scheduled once the future has completed.
-    Waiting,
-
-    /// The future succeeded and produced a result value. \c storage points
-    /// at that value.
-    Success,
-
-    /// The future finished by throwing an error. \c storage is that error
-    /// existential.
-    Error,
-  };
-
-  Kind kind;
+  /// Whether the storage represents the error result vs. the successful
+  /// result.
+  bool hadErrorResult;
+
+  /// Storage for the result of the future.
+  ///
+  /// When the future completed normally, this is a pointer to the storage
+  /// of the result value, which lives inside the future task itself.
+  ///
+  /// When the future completed by throwing an error, this is the error
+  /// object itself.
   OpaqueValue *storage;
 };
 
 /// Wait for a future task to complete.
 ///
-/// This can be called from any thread.
+/// This can be called from any thread. Its Swift signature is
 ///
+/// \code
+/// func swift_task_future_wait(on task: Builtin.NativeObject) async
+///     -> TaskFutureWaitResult
+/// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-TaskFutureWaitResult
-swift_task_future_wait(AsyncTask *task, AsyncTask *waitingTask);
+AsyncFunctionType<TaskFutureWaitResult(AsyncTask *task)>
+swift_task_future_wait;
 
 /// Add a status record to a task.  The record should not be
 /// modified while it is registered with a task.

stdlib/public/Concurrency/Task.cpp

@@ -73,6 +73,56 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask) {
   }
 }
 
+
+namespace {
+
+/// An asynchronous context within a task that describes a general "Future".
+/// task.
+///
+/// This type matches the ABI of a function `<T> () async throws -> T`, which
+/// is the type used by `Task.runDetached` and `Task.group.add` to create
+/// futures.
+class TaskFutureWaitAsyncContext : public AsyncContext {
+public:
+  // Error result is always present.
+  SwiftError *errorResult = nullptr;
+
+  // No indirect results.
+
+  TaskFutureWaitResult result;
+
+  // FIXME: Currently, this is always here, but it isn't technically
+  // necessary.
+  void* Self;
+
+  // Arguments.
+  AsyncTask *task;
+
+  using AsyncContext::AsyncContext;
+};
+
+}
+
+/// Run the given task, privoding it with the result of the future.
+static void runTaskWithFutureResult(
+    AsyncTask *waitingTask, ExecutorRef executor,
+    FutureFragment *futureFragment, bool hadErrorResult) {
+  auto waitingTaskContext =
+      static_cast<TaskFutureWaitAsyncContext *>(waitingTask->ResumeContext);
+
+  waitingTaskContext->result.hadErrorResult = hadErrorResult;
+  if (hadErrorResult) {
+    waitingTaskContext->result.storage =
+        reinterpret_cast<OpaqueValue *>(futureFragment->getError());
+  } else {
+    waitingTaskContext->result.storage = futureFragment->getStoragePtr();
+  }
+
+  // TODO: schedule this task on the executor rather than running it
+  // directly.
+  waitingTask->run(executor);
+}
+
 void AsyncTask::completeFuture(AsyncContext *context, ExecutorRef executor) {
   using Status = FutureFragment::Status;
   using WaitQueueItem = FutureFragment::WaitQueueItem;
@@ -103,9 +153,8 @@ void AsyncTask::completeFuture(AsyncContext *context, ExecutorRef executor) {
     // Find the next waiting task.
     auto nextWaitingTask = waitingTask->getNextWaitingTask();
 
-    // TODO: schedule this task on the executor rather than running it
-    // directly.
-    waitingTask->run(executor);
+    // Run the task.
+    runTaskWithFutureResult(waitingTask, executor, fragment, hadErrorResult);
 
     // Move to the next task.
     waitingTask = nextWaitingTask;
@@ -264,21 +313,37 @@ AsyncTaskAndContext swift::swift_task_create_future_f(
   return {task, initialContext};
 }
 
-TaskFutureWaitResult
-swift::swift_task_future_wait(AsyncTask *task, AsyncTask *waitingTask) {
+void swift::swift_task_future_wait(
+    AsyncTask *waitingTask, ExecutorRef executor,
+    AsyncContext *rawContext) {
+  // Suspend the waiting task.
+  waitingTask->ResumeTask = rawContext->ResumeParent;
+  waitingTask->ResumeContext = rawContext;
+
+  // Wait on the future.
+  auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
+  auto task = context->task;
   assert(task->isFuture());
   switch (task->waitFuture(waitingTask)) {
   case FutureFragment::Status::Executing:
-    return TaskFutureWaitResult{TaskFutureWaitResult::Waiting, nullptr};
+    // The waiting task has been queued on the future.
+    return;
 
   case FutureFragment::Status::Success:
-    return TaskFutureWaitResult{
-        TaskFutureWaitResult::Success, task->futureFragment()->getStoragePtr()};
-
-   case FutureFragment::Status::Error:
-      return TaskFutureWaitResult{
-          TaskFutureWaitResult::Error,
-          reinterpret_cast<OpaqueValue *>(task->futureFragment()->getError())};
+    // Run the task with a successful result.
+    // FIXME: Want to guarantee a tail call here
+    runTaskWithFutureResult(
+        waitingTask, executor, task->futureFragment(),
+        /*hadErrorResult=*/false);
+    return;
+
+ case FutureFragment::Status::Error:
+    // Run the task with an error result.
+    // FIXME: Want to guarantee a tail call here
+    runTaskWithFutureResult(
+        waitingTask, executor, task->futureFragment(),
+        /*hadErrorResult=*/true);
+    return;
   }
 }
 

stdlib/public/Concurrency/Task.swift

@@ -121,18 +121,16 @@ extension Task {
     /// and throwing a specific error or using `checkCancellation` the error
     /// thrown out of the task will be re-thrown here.
     public func get() async throws -> Success {
-      let rawResult = taskFutureWait(
-        on: task, waiting: Builtin.getCurrentAsyncTask())
-      switch TaskFutureWaitResult<Success>(raw: rawResult) {
-      case .executing:
-        fatalError("don't know how to synchronously return")
-
-      case .success(let result):
-        return result
-
-      case .failure(let error):
-        throw error
+      let rawResult = await taskFutureWait(on: task)
+      if rawResult.hadErrorResult {
+        // Throw the result on error.
+        throw unsafeBitCast(rawResult.storage, to: Error.self)
       }
+
+      // Take the value on success
+      let storagePtr =
+        rawResult.storage.bindMemory(to: Success.self, capacity: 1)
+      return UnsafeMutablePointer<Success>(mutating: storagePtr).pointee
     }
 
     /// Attempt to cancel the task.
@@ -269,8 +267,7 @@ extension Task {
     flags.isFuture = true
 
     // Create the asynchronous task future.
-    let (task, context) =
-      Builtin.createAsyncTaskFuture(flags.bits, nil, operation)
+    let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, nil, operation)
 
     return Handle<T>(task: task)
   }
@@ -316,8 +313,7 @@ extension Task {
     flags.isFuture = true
 
     // Create the asynchronous task future.
-    let (task, context) =
-      Builtin.createAsyncTaskFuture(flags.bits, nil, operation)
+    let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, nil, operation)
 
     return Handle<T>(task: task)
   }
@@ -423,47 +419,12 @@ public func runAsync(_ asyncFun: @escaping () async -> ()) {
   runTask(childTask.0)
 }
 
-/// Describes the result of waiting for a future.
-enum TaskFutureWaitResult<T> {
-  /// The future is still executing, and our waiting task has been placed
-  /// on its queue for when the future completes.
-  case executing
-
-  /// The future has succeeded with the given value.
-  case success(T)
-
-  /// The future has thrown the given error.
-  case failure(Error)
-
-  /// Initialize this instance from a raw result, taking any instance within
-  /// that result.
-  init(raw: RawTaskFutureWaitResult) {
-    switch raw.kind {
-    case 0:
-      self = .executing
-
-    case 1:
-      // Take the value on success
-      let storagePtr = raw.storage.bindMemory(to: T.self, capacity: 1)
-      self = .success(UnsafeMutablePointer<T>(mutating: storagePtr).move())
-
-    case 2:
-      // Take the error on error.
-      self = .failure(unsafeBitCast(raw.storage, to: Error.self))
-
-    default:
-      assert(false)
-      self = .executing
-    }
-  }
-}
-
 struct RawTaskFutureWaitResult {
-  let kind: Int
+  let hadErrorResult: Bool
   let storage: UnsafeRawPointer
 }
 
 @_silgen_name("swift_task_future_wait")
 func taskFutureWait(
-  on task: Builtin.NativeObject, waiting waitingTask: Builtin.NativeObject
-) -> RawTaskFutureWaitResult
+  on task: Builtin.NativeObject
+) async -> RawTaskFutureWaitResult

test/Concurrency/Runtime/basic_future.swift

@@ -1,12 +1,13 @@
-// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency) | %FileCheck %s
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency)
+
 // REQUIRES: executable_test
 // REQUIRES: concurrency
 // REQUIRES: OS=macosx
 
 import Dispatch
 
 extension DispatchQueue {
-  func async<R>(execute: @escaping () async -> R) -> Task.Handle<R> {
+  func async<R>(execute: @escaping () async throws -> R) -> Task.Handle<R> {
     let handle = Task.runDetached(operation: execute)
 
     // Run the task
@@ -16,30 +17,83 @@ extension DispatchQueue {
   }
 }
 
+enum HomeworkError: Error, Equatable {
+  case dogAteIt(String)
+}
+
 func formGreeting(name: String) async -> String {
   return "Hello \(name) from async world"
 }
 
-func test(name: String) {
-  let taskHandle = DispatchQueue.main.async { () async -> String in
+func testSimple(
+  name: String, dogName: String, shouldThrow: Bool, doSuspend: Bool
+) {
+  print("Testing name: \(name), dog: \(dogName), shouldThrow: \(shouldThrow) doSuspend: \(doSuspend)")
+
+  let queue = DispatchQueue(label: "concurrent", attributes: .concurrent)
+  let group = DispatchGroup()
+  var completed = false
+
+  group.enter()
+  let taskHandle = queue.async { () async throws -> String in
+    defer {
+      group.leave()
+    }
+
     let greeting = await formGreeting(name: name)
+
+    // If the intent is to test suspending, wait a bit so the second task
+    // can complete.
+    if doSuspend {
+      print("- Future sleeping")
+      sleep(1)
+    }
+
+    if (shouldThrow) {
+      print("- Future throwing")
+      throw HomeworkError.dogAteIt(dogName + " the dog")
+    }
+
+    print("- Future returning normally")
     return greeting + "!"
   }
 
-  _ = DispatchQueue.main.async { () async in
-    // CHECK: Sleeping
-    print("Sleeping...")
-    sleep(2)
-    let result = await try! taskHandle.get()
-    // CHECK: Hello Ted from async world
-    print(result)
-    assert(result == "Hello Ted from async world!")
-    exit(0)
+  group.enter()
+  _ = queue.async { () async in
+    defer {
+      group.leave()
+    }
+
+    // If the intent is not to test suspending, wait a bit so the first task
+    // can complete.
+    if !doSuspend {
+      print("+ Reader sleeping")
+      sleep(1)
+    }
+
+    do {
+      print("+ Reader waiting for the result")
+      let result = await try taskHandle.get()
+      completed = true
+      print("+ Normal return: \(result)")
+      assert(result == "Hello \(name) from async world!")
+    } catch HomeworkError.dogAteIt(let badDog) {
+      completed = true
+      print("+ Error return: HomeworkError.dogAteIt(\(badDog))")
+      assert(badDog == dogName + " the dog")
+    } catch {
+      fatalError("Caught a different exception?")
+    }
   }
 
-  print("Main task")
+  group.wait()
+  assert(completed)
+  print("Finished test")
 }
 
-test(name: "Ted")
+testSimple(name: "Ted", dogName: "Hazel", shouldThrow: false, doSuspend: false)
+testSimple(name: "Ted", dogName: "Hazel", shouldThrow: true, doSuspend: false)
+testSimple(name: "Ted", dogName: "Hazel", shouldThrow: false, doSuspend: true)
+testSimple(name: "Ted", dogName: "Hazel", shouldThrow: true, doSuspend: true)
 
-dispatchMain()
+print("Done")