Merge pull request #64054 from etcwilde/ewilde/startOnMainActor

Concurrency: swift_task_startOnMainActor

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,44 @@ extension Task: Equatable {
   }
 }
 
+@available(SwiftStdlib 5.9, *)
+extension Task where Failure == Error {
+    @_spi(MainActorUtilities)
+    @MainActor
+    @available(SwiftStdlib 5.9, *)
+    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
+    @available(SwiftStdlib 5.9, *)
+    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 +918,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

test/Concurrency/Runtime/startOnMainActor.swift

@@ -0,0 +1,88 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -Xfrontend -disable-availability-checking %s -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN:  %target-run %t/a.out
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+// UNSUPPORTED: back_deploy_concurrency
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: freestanding
+
+import StdlibUnittest
+@_spi(MainActorUtilities) import _Concurrency
+
+func doStuffAsync() async {
+    await Task.sleep(500)
+}
+
+let tests = TestSuite("StartOnMainActor")
+
+tests.test("startOnMainActor") {
+    // "global" variables for this test
+    struct Globals {
+        @MainActor
+        static var ran = false
+    }
+
+    @MainActor
+    func run() async {
+        Globals.ran = true
+        await doStuffAsync()
+    }
+
+    // enqueue item on the MainActor
+    let t1 = Task { @MainActor in
+        await Task.sleep(1000)
+    }
+
+    expectFalse(Globals.ran)
+
+    // Run something with side-effects on the main actor
+    let t2 = Task.startOnMainActor {
+        return await run()
+    }
+
+    expectTrue(Globals.ran)
+    await t1.value
+    await t2.value
+}
+
+tests.test("throwing startOnMainActor") {
+    // "global" variables for this test
+    struct Globals {
+        @MainActor
+        static var ran = false
+    }
+
+    struct StringError: Error {
+        let message: String
+    }
+
+    @MainActor
+    func run() async throws {
+        Globals.ran = true
+        await doStuffAsync()
+        throw StringError(message: "kablamo!")
+    }
+
+    // enqueue item on the MainActor
+    let t1 = Task { @MainActor in
+        await Task.sleep(1000)
+    }
+
+    expectFalse(Globals.ran)
+
+    // Run something with side-effects on the main actor
+    let t2 = Task.startOnMainActor {
+        return try await run()
+    }
+
+    expectTrue(Globals.ran)
+    await t1.value
+    expectNil(try? await t2.value)
+}
+
+await runAllTestsAsync()

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) {