Concurrency: Task.startOnMainActor

This patch adds an SPI to run the first partial function of a MainActor
asynchronous function on the MainActor synchronously. This is
effectively like the asynchronous program entrypoint behavior. The first
partial function is run synchronously. Following continuations are
enqueued for execution like any other asynchronous function.

Author: etcwilde

include/swift/Runtime/Concurrency.h

@@ -899,6 +899,9 @@ void swift_task_reportUnexpectedExecutor(
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 JobPriority swift_task_getCurrentThreadPriority(void);
 
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_startOnMainActor(AsyncTask* job);
+
 #if SWIFT_CONCURRENCY_COOPERATIVE_GLOBAL_EXECUTOR
 
 /// Donate this thread to the global executor until either the

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -382,6 +382,10 @@ OVERRIDE_TASK_STATUS(task_escalate, JobPriority,
                      swift::, (AsyncTask *task, JobPriority newPriority),
                      (task, newPriority))
 
+OVERRIDE_TASK(task_startOnMainActor, void,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
+              swift::, (AsyncTask *task), (task))
+
 #undef OVERRIDE
 #undef OVERRIDE_ACTOR
 #undef OVERRIDE_TASK

stdlib/public/Concurrency/Task.cpp

@@ -1594,6 +1594,17 @@ void (*swift::swift_task_asyncMainDrainQueue_hook)(
     swift_task_asyncMainDrainQueue_original original,
     swift_task_asyncMainDrainQueue_override compatOverride) = nullptr;
 
+SWIFT_CC(swift)
+static void swift_task_startOnMainActorImpl(AsyncTask* task) {
+  AsyncTask * originalTask = _swift_task_clearCurrent();
+  ExecutorRef mainExecutor = swift_task_getMainExecutor();
+  if (swift_task_getCurrentExecutor() != swift_task_getMainExecutor())
+    swift_Concurrency_fatalError(0, "Not on the main executor");
+  swift_retain(task);
+  swift_job_run(task, mainExecutor);
+  _swift_task_setCurrent(originalTask);
+}
+
 #define OVERRIDE_TASK COMPATIBILITY_OVERRIDE
 
 #ifdef SWIFT_STDLIB_SUPPORT_BACK_DEPLOYMENT

stdlib/public/Concurrency/Task.swift

@@ -186,6 +186,42 @@ extension Task: Equatable {
   }
 }
 
+@available(SwiftStdlib 5.9, *)
+extension Task where Failure == Error {
+    @_spi(MainActorUtilities)
+    @MainActor
+    public static func startOnMainActor(
+        priority: TaskPriority? = nil,
+        @_inheritActorContext @_implicitSelfCapture _ work: __owned @Sendable @escaping @MainActor() async throws -> Success
+    ) -> Task<Success, Error> {
+        let flags = taskCreateFlags(priority: priority, isChildTask: false,
+                                    copyTaskLocals: true, inheritContext: true,
+                                    enqueueJob: false,
+                                    addPendingGroupTaskUnconditionally: false)
+        let (task, _) = Builtin.createAsyncTask(flags, work)
+        _startTaskOnMainActor(task)
+        return Task<Success, Error>(task)
+    }
+}
+
+@available(SwiftStdlib 5.9, *)
+extension Task where Failure == Never {
+    @_spi(MainActorUtilities)
+    @MainActor
+    public static func startOnMainActor(
+        priority: TaskPriority? = nil,
+        @_inheritActorContext @_implicitSelfCapture _ work: __owned @Sendable @escaping @MainActor() async -> Success
+    ) -> Task<Success, Never> {
+        let flags = taskCreateFlags(priority: priority, isChildTask: false,
+                                    copyTaskLocals: true, inheritContext: true,
+                                    enqueueJob: false,
+                                    addPendingGroupTaskUnconditionally: false)
+        let (task, _) = Builtin.createAsyncTask(flags, work)
+        _startTaskOnMainActor(task)
+        return Task(task)
+    }
+}
+
 // ==== Task Priority ----------------------------------------------------------
 
 /// The priority of a task.
@@ -880,6 +916,9 @@ extension UnsafeCurrentTask: Equatable {
 @_silgen_name("swift_task_getCurrent")
 func _getCurrentAsyncTask() -> Builtin.NativeObject?
 
+@_silgen_name("swift_task_startOnMainActor")
+fileprivate func _startTaskOnMainActor(_ task: Builtin.NativeObject) -> Builtin.NativeObject?
+
 @available(SwiftStdlib 5.1, *)
 @_silgen_name("swift_task_getJobFlags")
 func getJobFlags(_ task: Builtin.NativeObject) -> JobFlags

unittests/runtime/CompatibilityOverrideConcurrency.cpp

@@ -97,6 +97,11 @@ static void swift_task_enqueueMainExecutor_override(
   Ran = true;
 }
 
+SWIFT_CC(swift)
+static void swift_task_startOnMainActor_override(AsyncTask* task) {
+  Ran = true;
+}
+
 #ifdef RUN_ASYNC_MAIN_DRAIN_QUEUE_TEST
 [[noreturn]] SWIFT_CC(swift)
 static void swift_task_asyncMainDrainQueue_override_fn(
@@ -284,6 +289,10 @@ TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_escalate) {
   swift_task_escalate(nullptr, {});
 }
 
+TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_startOnMainActorImpl) {
+  swift_task_startOnMainActor(nullptr);
+}
+
 #if RUN_ASYNC_MAIN_DRAIN_QUEUE_TEST
 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_task_asyncMainDrainQueue) {