[Concurrency] SerialExecutor.checkIsolated

Allow serial executors to perform a "last resort" check before crashing
in calls like preconditionIsolated and assumeIsolated.

This allows custom excecutors to use external information to claim that
while the synchronous code is not executing in a Swift Task, it IS
executing on the apropriate thread or queue the code is expecting.

add swift_task_checkIsolatedImpl to CooperativeGlobalExecutor

Author: ktoso

include/swift/ABI/Executor.h

@@ -155,8 +155,12 @@ class SerialExecutorRef {
     return reinterpret_cast<DefaultActor*>(Identity);
   }
 
+  bool hasSerialExecutorWitnessTable() const {
+    return !isGeneric() && !isDefaultActor();
+  }
+
   const SerialExecutorWitnessTable *getSerialExecutorWitnessTable() const {
-    assert(!isGeneric() && !isDefaultActor());
+    assert(hasSerialExecutorWitnessTable());
     auto table = Implementation & WitnessTableMask;
     return reinterpret_cast<const SerialExecutorWitnessTable*>(table);
   }

include/swift/Runtime/Concurrency.h

@@ -715,6 +715,11 @@ void swift_task_enqueue(Job *job, SerialExecutorRef executor);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_enqueueGlobal(Job *job);
 
+/// Invoke an executor's `checkIsolated` or otherwise equivalent API,
+/// that will crash if the current executor is NOT the passed executor.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_checkIsolated(SerialExecutorRef executor);
+
 /// A count in nanoseconds.
 using JobDelay = unsigned long long;
 
@@ -729,6 +734,9 @@ void swift_task_enqueueGlobalWithDeadline(long long sec, long long nsec,
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_enqueueMainExecutor(Job *job);
 
+/// WARNING: This method is expected to CRASH when caller is not on the
+/// expected executor.
+///
 /// 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(
@@ -781,6 +789,12 @@ SWIFT_CC(swift) void (*swift_task_enqueueGlobalWithDeadline_hook)(
     int clock, Job *job,
     swift_task_enqueueGlobalWithDeadline_original original);
 
+typedef SWIFT_CC(swift) void (*swift_task_checkIsolated_original)(SerialExecutorRef executor);
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift) void (*swift_task_checkIsolated_hook)(
+    SerialExecutorRef executor, swift_task_checkIsolated_original original);
+
+
 typedef SWIFT_CC(swift) bool (*swift_task_isOnExecutor_original)(
     HeapObject *executor,
     const Metadata *selfType,

stdlib/public/Concurrency/Actor.cpp

@@ -22,6 +22,7 @@
 #include "../CompatibilityOverride/CompatibilityOverride.h"
 #include "swift/ABI/Actor.h"
 #include "swift/ABI/Task.h"
+#include "TaskPrivate.h"
 #include "swift/Basic/ListMerger.h"
 #include "swift/Concurrency/Actor.h"
 #include "swift/Runtime/AccessibleFunction.h"
@@ -314,39 +315,109 @@ bool _task_serialExecutor_isSameExclusiveExecutionContext(
     const SerialExecutorWitnessTable *wtable);
 
 SWIFT_CC(swift)
-static bool swift_task_isCurrentExecutorImpl(SerialExecutorRef executor) {
+static bool swift_task_isCurrentExecutorImpl(SerialExecutorRef expectedExecutor) {
   auto current = ExecutorTrackingInfo::current();
 
   if (!current) {
-    // TODO(ktoso): checking the "is main thread" is not correct, main executor can be not main thread, relates to rdar://106188692
-    return executor.isMainExecutor() && isExecutingOnMainThread();
+    // We have no current executor, i.e. we are running "outside" of Swift
+    // Concurrency. We could still be running on a thread/queue owned by
+    // the expected executor however, so we need to try a bit harder before
+    // we fail.
+
+    // Are we expecting the main executor and are using the main thread?
+    if (expectedExecutor.isMainExecutor() && isExecutingOnMainThread()) {
+      // TODO(concurrency): consider removing this special case, as checkIsolated will compare against mainQueue already
+      return true;
+    }
+
+    // Otherwise, as last resort, let the expected executor check using
+    // external means, as it may "know" this thread is managed by it etc.
+    swift_task_checkIsolated(expectedExecutor);
+
+    // checkIsolated did not crash, so we are on the right executor, after all!
+    return true;
   }
 
-  auto currentExecutor = current->getActiveExecutor();
-  if (currentExecutor == executor) {
+  SerialExecutorRef currentExecutor = current->getActiveExecutor();
+
+  // Fast-path: the executor is exactly the same memory address;
+  // We assume executors do not come-and-go appearing under the same address,
+  // and treat pointer equality of executors as good enough to assume the executor.
+  if (currentExecutor == expectedExecutor) {
     return true;
   }
 
-  if (executor.isComplexEquality()) {
+  // Fast-path, specialize the common case of comparing two main executors.
+  if (currentExecutor.isMainExecutor() && expectedExecutor.isMainExecutor()) {
+    return true;
+  }
+
+  // If the expected executor is "default" then we should have matched
+  // by pointer equality already with the current executor.
+  if (expectedExecutor.isDefaultActor()) {
+    // If the expected executor is a default actor, it makes no sense to try
+    // the 'checkIsolated' call, it must be equal to the other actor, or it is
+    // not the same isolation domain.
+    swift_Concurrency_fatalError(0, "Incorrect actor executor assumption");
+    return false;
+  }
+
+  if (expectedExecutor.isMainExecutor() && !currentExecutor.isMainExecutor()) {
+    // TODO: Invoke checkIsolated() on "main" SerialQueue once it implements `checkIsolated`, otherwise messages will be sub-par and hard to address
+    swift_Concurrency_fatalError(0, "Incorrect actor executor assumption; Expected MainActor executor");
+    return false;
+  } else if (!expectedExecutor.isMainExecutor() && currentExecutor.isMainExecutor()) {
+    // TODO: Invoke checkIsolated() on "main" SerialQueue once it implements `checkIsolated`, otherwise messages will be sub-par and hard to address
+    swift_Concurrency_fatalError(0, "Incorrect actor executor assumption; Expected not-MainActor executor");
+    return false;
+  }
+
+  if (expectedExecutor.isComplexEquality()) {
     if (!swift_compareWitnessTables(
         reinterpret_cast<const WitnessTable*>(currentExecutor.getSerialExecutorWitnessTable()),
-        reinterpret_cast<const  WitnessTable*>(executor.getSerialExecutorWitnessTable()))) {
+        reinterpret_cast<const  WitnessTable*>(expectedExecutor.getSerialExecutorWitnessTable()))) {
       // different witness table, we cannot invoke complex equality call
       return false;
     }
+
     // Avoid passing nulls to Swift for the isSame check:
-    if (!currentExecutor.getIdentity() || !executor.getIdentity()) {
+    if (!currentExecutor.getIdentity() || !expectedExecutor.getIdentity()) {
       return false;
     }
 
-    return _task_serialExecutor_isSameExclusiveExecutionContext(
+    auto result = _task_serialExecutor_isSameExclusiveExecutionContext(
         currentExecutor.getIdentity(),
-        executor.getIdentity(),
+        expectedExecutor.getIdentity(),
         swift_getObjectType(currentExecutor.getIdentity()),
-        executor.getSerialExecutorWitnessTable());
+        expectedExecutor.getSerialExecutorWitnessTable());
+
+    return result;
   }
 
-  return false;
+  // This provides a last-resort check by giving the expected SerialExecutor the
+  // chance to perform a check using some external knowledge if perhaps we are,
+  // after all, on this executor, but the Swift concurrency runtime was just not
+  // aware.
+  //
+  // Unless handled in `swift_task_checkIsolated` directly, this should call
+  // through to the executor's `SerialExecutor.checkIsolated`.
+  //
+  // This call is expected to CRASH, unless it has some way of proving that
+  // we're actually indeed running on this executor.
+  //
+  // For example, when running outside of Swift concurrency tasks, but trying to
+  // `MainActor.assumeIsolated` while executing DIRECTLY on the main dispatch
+  // queue, this allows Dispatch to check for this using its own tracking
+  // mechanism, and thus allow the assumeIsolated to work correctly, even though
+  // the code executing is not even running inside a Task.
+  //
+  // Note that this only works because the closure in assumeIsolated is
+  // synchronous, and will not cause suspensions, as that would require the
+  // presence of a Task.
+  swift_task_checkIsolated(expectedExecutor);
+
+  // The checkIsolated call did not crash, so we are on the right executor.
+  return true;
 }
 
 /// Logging level for unexpected executors:

stdlib/public/Concurrency/CooperativeGlobalExecutor.inc

@@ -18,6 +18,7 @@
 ///   swift_task_enqueueGlobalImpl
 ///   swift_task_enqueueGlobalWithDelayImpl
 ///   swift_task_enqueueMainExecutorImpl
+///   swift_task_checkIsolatedImpl
 /// as well as any cooperative-executor-specific functions in the runtime.
 ///
 ///===------------------------------------------------------------------===///
@@ -136,6 +137,13 @@ static void insertDelayedJob(Job *newJob, JobDeadline deadline) {
   *position = newJob;
 }
 
+SWIFT_CC(swift)
+static void swift_task_checkIsolatedImpl(SerialExecutorRef executor) {
+  _task_serialExecutor_checkIsolated(
+      executor.getIdentity(), swift_getObjectType(executor.getIdentity()),
+      executor.getSerialExecutorWitnessTable());
+}
+
 /// Insert a job into the cooperative global queue with a delay.
 SWIFT_CC(swift)
 static void swift_task_enqueueGlobalWithDelayImpl(JobDelay delay,

stdlib/public/Concurrency/DispatchGlobalExecutor.inc

@@ -18,11 +18,13 @@
 ///   swift_task_enqueueGlobalImpl
 ///   swift_task_enqueueGlobalWithDelayImpl
 ///   swift_task_enqueueMainExecutorImpl
+///   swift_task_checkIsolated
 /// as well as any Dispatch-specific functions for the runtime.
 ///
 ///===------------------------------------------------------------------===///
 
 #if SWIFT_CONCURRENCY_ENABLE_DISPATCH
+#include "swift/Runtime/HeapObject.h"
 #include <dispatch/dispatch.h>
 #if defined(_WIN32)
 #include <Windows.h>
@@ -233,7 +235,7 @@ static void swift_task_enqueueGlobalImpl(Job *job) {
 
 SWIFT_CC(swift)
 static void swift_task_enqueueGlobalWithDelayImpl(JobDelay delay,
-                                                  Job *job) {
+                                              Job *job) {
   assert(job && "no job provided");
 
   dispatch_function_t dispatchFunction = &__swift_run_job;
@@ -383,3 +385,56 @@ void swift::swift_task_enqueueOnDispatchQueue(Job *job,
   auto queue = reinterpret_cast<dispatch_queue_t>(_queue);
   dispatchEnqueue(queue, job, (dispatch_qos_class_t)priority, queue);
 }
+
+/// Recognize if the SerialExecutor is specifically a `DispatchSerialQueue`
+/// by comparing witness tables and return it if true.
+static dispatch_queue_s *getAsDispatchSerialQueue(SerialExecutorRef executor) {
+  if (!executor.hasSerialExecutorWitnessTable()) {
+    return nullptr;
+  }
+
+  auto executorWitnessTable = reinterpret_cast<const WitnessTable *>(
+      executor.getSerialExecutorWitnessTable());
+  auto serialQueueWitnessTable = reinterpret_cast<const WitnessTable *>(
+      _swift_task_getDispatchQueueSerialExecutorWitnessTable());
+
+  if (swift_compareWitnessTables(executorWitnessTable,
+                                 serialQueueWitnessTable)) {
+    return reinterpret_cast<dispatch_queue_s *>(executor.getIdentity());
+  } else {
+    return nullptr;
+  }
+}
+
+/// If the executor is a `DispatchSerialQueue` we're able to invoke the
+/// dispatch's precondition API directly -- this is more efficient than going
+/// through the runtime call to end up calling the same API, and also allows us
+/// to perform this assertion on earlier platforms, where the `checkIsolated`
+/// requirement/witness was not shipping yet.
+SWIFT_CC(swift)
+static void swift_task_checkIsolatedImpl(SerialExecutorRef executor) {
+  // If it is the main executor, compare with the Main queue
+  if (executor.isMainExecutor()) {
+    dispatch_assert_queue(dispatch_get_main_queue());
+    return;
+  }
+
+  // if able to, use the checkIsolated implementation in Swift
+  if (executor.hasSerialExecutorWitnessTable()) {
+    _task_serialExecutor_checkIsolated(
+        executor.getIdentity(), swift_getObjectType(executor.getIdentity()),
+        executor.getSerialExecutorWitnessTable());
+    return;
+  }
+
+  if (auto queue = getAsDispatchSerialQueue(executor)) {
+    // if the executor was not SerialExecutor for some reason but we're able
+    // to get a queue from it anyway, use the assert directly on it.
+    dispatch_assert_queue(queue); // TODO(concurrency): could we report a better message here somehow?
+    return;
+  }
+
+  // otherwise, we have no way to check, so report an error
+  // TODO: can we swift_getTypeName(swift_getObjectType(executor.getIdentity()), false).data safely in the message here?
+  swift_Concurrency_fatalError(0, "Incorrect actor executor assumption");
+}