Merge pull request swiftlang#63491 from ktoso/wip-custom-executor-not-directly-on-actor

[CustomExecutors] unownedExecutor can be declared not directly on actor

Author: swift-ci

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

@@ -9365,38 +9365,26 @@ Type TypeBase::getSwiftNewtypeUnderlyingType() {
 }
 
 const VarDecl *ClassDecl::getUnownedExecutorProperty() const {
-  auto &ctx = getASTContext();
+  auto &C = getASTContext();
 
-  auto hasUnownedSerialExecutorType = [&](VarDecl *property) {
-    if (auto type = property->getInterfaceType())
-      if (auto td = type->getAnyNominal())
-        if (td == ctx.getUnownedSerialExecutorDecl())
-          return true;
-    return false;
-  };
+  if (!isAnyActor())
+    return nullptr;
 
-  VarDecl *candidate = nullptr;
-  for (auto member: getMembers()) {
-    // Instance properties called unownedExecutor.
-    if (auto property = dyn_cast<VarDecl>(member)) {
-      if (property->getName() == ctx.Id_unownedExecutor &&
-          !property->isStatic()) {
-        if (!candidate) {
-          candidate = property;
-          continue;
-        }
+  llvm::SmallVector<ValueDecl *, 2> results;
+  this->lookupQualified(getSelfNominalTypeDecl(),
+                        DeclNameRef(C.Id_unownedExecutor),
+                        NL_ProtocolMembers,
+                        results);
 
-        bool oldHasRightType = hasUnownedSerialExecutorType(candidate);
-        if (oldHasRightType == hasUnownedSerialExecutorType(property)) {
-          // just ignore the new property, we should diagnose this eventually
-        } else if (!oldHasRightType) {
-          candidate = property;
-        }
-      }
-    }
+  for (auto candidate: results) {
+    if (isa<ProtocolDecl>(candidate->getDeclContext()))
+      continue;
+
+    if (VarDecl *var = dyn_cast<VarDecl>(candidate))
+      return var;
   }
 
-  return candidate;
+  return nullptr;
 }
 
 bool ClassDecl::isRootDefaultActor() const {

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,28 @@ 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  *******************************/
+/*****************************************************************************/
+
 void swift::swift_task_enqueueMainExecutor(Job *job) {
   concurrency::trace::job_enqueue_main_executor(job);
   if (swift_task_enqueueMainExecutor_hook)

test/Concurrency/Runtime/custom_executors_default.swift

@@ -0,0 +1,50 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking %import-libdispatch -parse-as-library) | %FileCheck %s
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// UNSUPPORTED: freestanding
+
+// UNSUPPORTED: back_deployment_runtime
+// REQUIRES: concurrency_runtime
+
+
+actor Simple {
+  var count = 0
+  func report() {
+    print("simple.count == \(count)")
+    count += 1
+  }
+}
+
+actor Custom {
+  var count = 0
+  let simple = Simple()
+
+  func report() async {
+    print("custom.count == \(count)")
+    count += 1
+
+    await simple.report()
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+@main struct Main {
+  static func main() async {
+    print("begin")
+    let actor = Custom()
+    await actor.report()
+    await actor.report()
+    await actor.report()
+    print("end")
+  }
+}
+
+// CHECK:      begin
+// CHECK-NEXT: custom.count == 0
+// CHECK-NEXT: simple.count == 0
+// CHECK-NEXT: custom.count == 1
+// CHECK-NEXT: simple.count == 1
+// CHECK-NEXT: custom.count == 2
+// CHECK-NEXT: simple.count == 2
+// CHECK-NEXT: end

test/Concurrency/Runtime/custom_executors_protocol.swift

@@ -0,0 +1,90 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking %import-libdispatch -parse-as-library) | %FileCheck %s
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// REQUIRES: libdispatch
+// UNSUPPORTED: freestanding
+
+// UNSUPPORTED: back_deployment_runtime
+// REQUIRES: concurrency_runtime
+
+import Dispatch
+
+func checkIfMainQueue(expectedAnswer expected: Bool) {
+  dispatchPrecondition(condition: expected ? .onQueue(DispatchQueue.main)
+      : .notOnQueue(DispatchQueue.main))
+}
+
+protocol WithSpecifiedExecutor: Actor {
+  nonisolated var executor: SpecifiedExecutor { get }
+}
+
+protocol SpecifiedExecutor: SerialExecutor {}
+
+extension WithSpecifiedExecutor {
+  /// Establishes the WithSpecifiedExecutorExecutor as the serial
+  /// executor that will coordinate execution for the actor.
+  nonisolated var unownedExecutor: UnownedSerialExecutor {
+    executor.asUnownedSerialExecutor()
+  }
+}
+
+final class InlineExecutor: SpecifiedExecutor, Swift.CustomStringConvertible {
+  let name: String
+
+  init(_ name: String) {
+    self.name = name
+  }
+
+  public func enqueue(_ job: UnownedJob) {
+    print("\(self): enqueue")
+    job._runSynchronously(on: self.asUnownedSerialExecutor())
+    print("\(self): after run")
+  }
+
+  public func asUnownedSerialExecutor() -> UnownedSerialExecutor {
+    return UnownedSerialExecutor(ordinary: self)
+  }
+
+  var description: Swift.String {
+    "InlineExecutor(\(name))"
+  }
+}
+
+actor MyActor: WithSpecifiedExecutor {
+
+  nonisolated let executor: SpecifiedExecutor
+
+  // Note that we don't have to provide the unownedExecutor in the actor itself.
+  // We obtain it from the extension on `WithSpecifiedExecutor`.
+
+  init(executor: SpecifiedExecutor) {
+    self.executor = executor
+  }
+
+  func test(expectedExecutor: some SerialExecutor) {
+    precondition(_taskIsOnExecutor(expectedExecutor), "Expected to be on: \(expectedExecutor)")
+    checkIfMainQueue(expectedAnswer: true)
+    print("\(Self.self): on executor \(expectedExecutor)")
+  }
+}
+
+@main struct Main {
+  static func main() async {
+    print("begin")
+    let one = InlineExecutor("one")
+    let actor = MyActor(executor: one)
+    await actor.test(expectedExecutor: one)
+    await actor.test(expectedExecutor: one)
+    await actor.test(expectedExecutor: one)
+    print("end")
+  }
+}
+
+// CHECK:      begin
+// CHECK-NEXT: InlineExecutor(one): enqueue
+// CHECK-NEXT: MyActor: on executor InlineExecutor(one)
+// CHECK-NEXT: MyActor: on executor InlineExecutor(one)
+// CHECK-NEXT: MyActor: on executor InlineExecutor(one)
+// CHECK-NEXT: InlineExecutor(one): after run
+// CHECK-NEXT: end