[Concurrenncy] add internal way to assert we're on expected executor

Author: ktoso

include/swift/Runtime/Concurrency.h

@@ -729,6 +729,13 @@ void swift_task_enqueueGlobalWithDeadline(long long sec, long long nsec,
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_enqueueMainExecutor(Job *job);
 
+/// Return true if the caller is running in a Task on the passed Executor.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+bool swift_task_isOnExecutor(
+    HeapObject * executor,
+    const Metadata *selfType,
+    const SerialExecutorWitnessTable *wtable);
+
 #if SWIFT_CONCURRENCY_ENABLE_DISPATCH
 
 /// Enqueue the given job on the main executor.
@@ -766,6 +773,17 @@ SWIFT_CC(swift) void (*swift_task_enqueueGlobalWithDeadline_hook)(
     int clock, Job *job,
     swift_task_enqueueGlobalWithDeadline_original original);
 
+typedef SWIFT_CC(swift) bool (*swift_task_isOnExecutor_original)(
+    HeapObject *executor,
+    const Metadata *selfType,
+    const SerialExecutorWitnessTable *wtable);
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift) bool (*swift_task_isOnExecutor_hook)(
+    HeapObject *executor,
+    const Metadata *selfType,
+    const SerialExecutorWitnessTable *wtable,
+    swift_task_isOnExecutor_original original);
+
 /// A hook to take over main executor enqueueing.
 typedef SWIFT_CC(swift) void (*swift_task_enqueueMainExecutor_original)(
     Job *job);

lib/AST/Decl.cpp

@@ -9357,7 +9357,6 @@ Type TypeBase::getSwiftNewtypeUnderlyingType() {
 
 const VarDecl *ClassDecl::getUnownedExecutorProperty() const {
   auto &C = getASTContext();
-  auto module = getParentModule();
 
   if (!isAnyActor())
     return nullptr;

stdlib/public/Concurrency/Executor.swift

@@ -69,6 +69,18 @@ public struct UnownedSerialExecutor: Sendable {
   }
 }
 
+/// Checks if the current task is running on the expected executor.
+///
+/// Generally, Swift programs should be constructed such that it is statically
+/// known that a specific executor is used, for example by using global actors or
+/// custom executors. However, in some APIs it may be useful to provide an
+/// additional runtime check for this, especially when moving towards Swift
+/// concurrency from other runtimes which frequently use such assertions.
+/// - Parameter executor: The expected executor.
+@available(SwiftStdlib 5.9, *)
+@_silgen_name("swift_task_isOnExecutor")
+public func _taskIsOnExecutor<Executor: SerialExecutor>(_ executor: Executor) -> Bool
+
 // Used by the concurrency runtime
 @available(SwiftStdlib 5.1, *)
 @_silgen_name("_swift_task_enqueueOnExecutor")

stdlib/public/Concurrency/GlobalExecutor.cpp

@@ -79,6 +79,13 @@ void (*swift::swift_task_enqueueGlobalWithDeadline_hook)(
     int clock, Job *job,
     swift_task_enqueueGlobalWithDeadline_original original) = nullptr;
 
+SWIFT_CC(swift)
+bool (*swift::swift_task_isOnExecutor_hook)(
+    HeapObject *executor,
+    const Metadata *selfType,
+    const SerialExecutorWitnessTable *wtable,
+    swift_task_isOnExecutor_original original) = nullptr;
+
 SWIFT_CC(swift)
 void (*swift::swift_task_enqueueMainExecutor_hook)(
     Job *job, swift_task_enqueueMainExecutor_original original) = nullptr;
@@ -125,6 +132,24 @@ void swift::swift_task_enqueueGlobalWithDeadline(
     swift_task_enqueueGlobalWithDeadlineImpl(sec, nsec, tsec, tnsec, clock, job);
 }
 
+SWIFT_CC(swift)
+static bool swift_task_isOnExecutorImpl(HeapObject *executor,
+                                        const Metadata *selfType,
+                                        const SerialExecutorWitnessTable *wtable) {
+  auto executorRef = ExecutorRef::forOrdinary(executor, wtable);
+  return swift_task_isCurrentExecutor(executorRef);
+}
+
+bool swift::swift_task_isOnExecutor(HeapObject *executor,
+                                    const Metadata *selfType,
+                                    const SerialExecutorWitnessTable *wtable) {
+  if (swift_task_isOnExecutor_hook)
+    return swift_task_isOnExecutor_hook(
+        executor, selfType, wtable, swift_task_isOnExecutorImpl);
+  else
+    return swift_task_isOnExecutorImpl(executor, selfType, wtable);
+}
+
 /*****************************************************************************/
 /****************************** MAIN EXECUTOR  *******************************/
 /*****************************************************************************/

test/Concurrency/Runtime/custom_executors_protocol.swift

@@ -63,6 +63,7 @@ actor MyActor: WithSpecifiedExecutor {
   }
 
   func test(expectedExecutor: some SerialExecutor) {
+    precondition(_taskIsOnExecutor(expectedExecutor), "Expected to be on: \(expectedExecutor)")
     checkIfMainQueue(expectedAnswer: true)
     print("\(Self.self): on executor \(expectedExecutor)")
   }