[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.

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,107 @@ 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: remove this special case when MainExecutor implements checkIsolated
+      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(
         currentExecutor.getIdentity(),
-        executor.getIdentity(),
+        expectedExecutor.getIdentity(),
         swift_getObjectType(currentExecutor.getIdentity()),
-        executor.getSerialExecutorWitnessTable());
+        expectedExecutor.getSerialExecutorWitnessTable());
   }
 
-  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

@@ -136,6 +136,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,6 +18,7 @@
 ///   swift_task_enqueueGlobalImpl
 ///   swift_task_enqueueGlobalWithDelayImpl
 ///   swift_task_enqueueMainExecutorImpl
+///   swift_task_checkIsolated
 /// as well as any Dispatch-specific functions for the runtime.
 ///
 ///===------------------------------------------------------------------===///
@@ -233,7 +234,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 +384,39 @@ 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 (auto queue = getAsDispatchSerialQueue(executor)) {
+    dispatch_assert_queue(queue);
+  } else if (executor.hasSerialExecutorWitnessTable()) {
+    _task_serialExecutor_checkIsolated(
+        executor.getIdentity(), swift_getObjectType(executor.getIdentity()),
+        executor.getSerialExecutorWitnessTable());
+  }
+}

stdlib/public/Concurrency/Executor.swift

@@ -100,6 +100,43 @@ public protocol SerialExecutor: Executor {
   ///            perspective–to execute code assuming one on the other.
   @available(SwiftStdlib 5.9, *)
   func isSameExclusiveExecutionContext(other: Self) -> Bool
+
+  /// Last resort "fallback" isolation check, called when the concurrency runtime
+  /// is comparing executors e.g. during ``assumeIsolated()`` and is unable to prove
+  /// serial equivalence between the expected (this object), and the current executor.
+  ///
+  /// During executor comparison, the Swift concurrency runtime attempts to compare
+  /// current and expected executors in a few ways (including "complex" equality
+  /// between executors (see ``isSameExclusiveExecutionContext(other:)``), and if all
+  /// those checks fail, this method is invoked on the expected executor.
+  ///
+  /// This method MUST crash if it is unable to prove that the current execution
+  /// context belongs to this executor. At this point usual executor comparison would
+  /// have already failed, though the executor may have some external tracking of
+  /// threads it owns, and may be able to prove isolation nevertheless.
+  ///
+  /// A default implementation is provided that unconditionally crashes the
+  /// program, and prevents calling code from proceeding with potentially
+  /// not thread-safe execution.
+  ///
+  /// - Warning: This method must crash and halt program execution if unable
+  ///     to prove the isolation of the calling context.
+  @available(SwiftStdlib 6.0, *)
+  func checkIsolated()
+
+}
+
+@available(SwiftStdlib 6.0, *)
+extension SerialExecutor {
+
+  @available(SwiftStdlib 6.0, *)
+  public func checkIsolated() {
+    #if !$Embedded
+    fatalError("Unexpected isolation context, expected to be executing on \(Self.self)")
+    #else
+    Builtin.int_trap()
+    #endif
+  }
 }
 
 /// An executor that may be used as preferred executor by a task.
@@ -120,7 +157,7 @@ public protocol SerialExecutor: Executor {
 ///
 /// Unstructured tasks do not inherit the task executor.
 @_unavailableInEmbedded
-@available(SwiftStdlib 9999, *)
+@available(SwiftStdlib 6.0, *)
 public protocol TaskExecutor: Executor {
   // This requirement is repeated here as a non-override so that we
   // get a redundant witness-table entry for it.  This allows us to
@@ -152,7 +189,7 @@ public protocol TaskExecutor: Executor {
 }
 
 @_unavailableInEmbedded
-@available(SwiftStdlib 9999, *)
+@available(SwiftStdlib 6.0, *)
 extension TaskExecutor {
   public func asUnownedTaskExecutor() -> UnownedTaskExecutor {
     UnownedTaskExecutor(ordinary: self)
@@ -270,7 +307,7 @@ public struct UnownedSerialExecutor: Sendable {
 
 
 @_unavailableInEmbedded
-@available(SwiftStdlib 9999, *)
+@available(SwiftStdlib 6.0, *)
 @frozen
 public struct UnownedTaskExecutor: Sendable {
   #if $BuiltinExecutor
@@ -279,7 +316,7 @@ public struct UnownedTaskExecutor: Sendable {
 
   /// SPI: Do not use. Cannot be marked @_spi, since we need to use it from Distributed module
   /// which needs to reach for this from an @_transparent function which prevents @_spi use.
-  @available(SwiftStdlib 9999, *)
+  @available(SwiftStdlib 6.0, *)
   public var _executor: Builtin.Executor {
     self.executor
   }
@@ -303,24 +340,34 @@ public struct UnownedTaskExecutor: Sendable {
 }
 
 @_unavailableInEmbedded
-@available(SwiftStdlib 9999, *)
+@available(SwiftStdlib 6.0, *)
 extension UnownedTaskExecutor: Equatable {
   @inlinable
   public static func == (_ lhs: UnownedTaskExecutor, _ rhs: UnownedTaskExecutor) -> Bool {
     unsafeBitCast(lhs.executor, to: (Int, Int).self) == unsafeBitCast(rhs.executor, to: (Int, Int).self)
   }
 }
 
-/// Checks if the current task is running on the expected executor.
+/// Returns either `true` or will CRASH if called from a different executor
+/// than the passed `executor`.
+///
+/// This method will attempt to verify the current executor against `executor`,
+/// and as a last-resort call through to `SerialExecutor.checkIsolated`.
+///
+/// This method will never return `false`. It either can verify we're on the
+/// correct executor, or will crash the program. It should be used in
+/// isolation correctness guaranteeing APIs.
 ///
 /// 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.
+@_spi(ConcurrencyExecutors)
 @available(SwiftStdlib 5.9, *)
-@_silgen_name("swift_task_isOnExecutor")
+@_silgen_name("swift_task_isOnExecutor") // This function will CRASH rather than return `false`!
 public func _taskIsOnExecutor<Executor: SerialExecutor>(_ executor: Executor) -> Bool
 
 @_spi(ConcurrencyExecutors)
@@ -361,6 +408,13 @@ internal func _task_serialExecutor_isSameExclusiveExecutionContext<E>(current cu
   currentExecutor.isSameExclusiveExecutionContext(other: executor)
 }
 
+@available(SwiftStdlib 6.0, *)
+@_silgen_name("_task_serialExecutor_checkIsolated")
+internal func _task_serialExecutor_checkIsolated<E>(executor: E)
+    where E: SerialExecutor {
+  executor.checkIsolated()
+}
+
 /// Obtain the executor ref by calling the executor's `asUnownedSerialExecutor()`.
 /// The obtained executor ref will have all the user-defined flags set on the executor.
 @available(SwiftStdlib 5.9, *)
@@ -373,7 +427,7 @@ internal func _task_serialExecutor_getExecutorRef<E>(_ executor: E) -> Builtin.E
 /// Obtain the executor ref by calling the executor's `asUnownedTaskExecutor()`.
 /// The obtained executor ref will have all the user-defined flags set on the executor.
 @_unavailableInEmbedded
-@available(SwiftStdlib 9999, *)
+@available(SwiftStdlib 6.0, *)
 @_silgen_name("_task_executor_getTaskExecutorRef")
 internal func _task_executor_getTaskExecutorRef(_ taskExecutor: any TaskExecutor) -> Builtin.Executor {
   return taskExecutor.asUnownedTaskExecutor().executor
@@ -396,7 +450,7 @@ where E: SerialExecutor {
 }
 
 @_unavailableInEmbedded
-@available(SwiftStdlib 9999, *)
+@available(SwiftStdlib 6.0, *)
 @_silgen_name("_swift_task_enqueueOnTaskExecutor")
 internal func _enqueueOnTaskExecutor<E>(job unownedJob: UnownedJob, executor: E) where E: TaskExecutor {
   #if !SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY