[Executors][Distributed] assumeOnLocalDistributedActorExecutor

lib/Sema/DerivedConformanceDistributedActor.cpp

@@ -787,7 +787,6 @@ std::pair<Type, TypeDecl *> DerivedConformance::deriveDistributedActor(
   if (!canDeriveDistributedActor(Nominal, cast<DeclContext>(ConformanceDecl)))
     return std::make_pair(Type(), nullptr);
 
-
   if (assocType->getName() == Context.Id_ActorSystem) {
     return std::make_pair(deriveDistributedActorType_ActorSystem(*this),
                           nullptr);

stdlib/public/Concurrency/Actor.cpp

@@ -1993,7 +1993,7 @@ void swift::swift_nonDefaultDistributedActor_initialize(NonDefaultDistributedAct
 
 OpaqueValue*
 swift::swift_distributedActor_remote_initialize(const Metadata *actorType) {
-  auto *metadata = actorType->getClassObject();
+  const ClassMetadata *metadata = actorType->getClassObject();
 
   // TODO(distributed): make this allocation smaller
   // ==== Allocate the memory for the remote instance
@@ -2014,21 +2014,22 @@ swift::swift_distributedActor_remote_initialize(const Metadata *actorType) {
   if (isDefaultActorClass(metadata)) {
     auto actor = asImpl(reinterpret_cast<DefaultActor *>(alloc));
     actor->initialize(/*remote*/true);
+    assert(swift_distributed_actor_is_remote(alloc));
     return reinterpret_cast<OpaqueValue*>(actor);
   } else {
     auto actor = asImpl(reinterpret_cast<NonDefaultDistributedActor *>(alloc));
     actor->initialize(/*remote*/true);
+    assert(swift_distributed_actor_is_remote(alloc));
+    return reinterpret_cast<OpaqueValue*>(actor);
   }
-  assert(swift_distributed_actor_is_remote(alloc));
-
 }
 
 bool swift::swift_distributed_actor_is_remote(HeapObject *_actor) {
-  auto metadata = cast<ClassMetadata>(_actor->metadata);
+  const ClassMetadata *metadata = cast<ClassMetadata>(_actor->metadata);
   if (isDefaultActorClass(metadata)) {
-    return asImpl((DefaultActor *) _actor)->isDistributedRemote();
+    return asImpl(reinterpret_cast<DefaultActor *>(_actor))->isDistributedRemote();
   } else {
-    return asImpl((NonDefaultDistributedActor *) _actor)->isDistributedRemote(); // NEW
+    return asImpl(reinterpret_cast<NonDefaultDistributedActor *>(_actor))->isDistributedRemote();
   }
 }
 

stdlib/public/Concurrency/Executor.swift

@@ -50,6 +50,12 @@ public struct UnownedSerialExecutor: Sendable {
   #if compiler(>=5.5) && $BuiltinExecutor
   @usableFromInline
   internal var executor: Builtin.Executor
+
+  @_spi(ConcurrencyExecutors)
+  @available(SwiftStdlib 5.9, *)
+  public var _executor: Builtin.Executor {
+    self.executor
+  }
   #endif
 
   @inlinable
@@ -67,6 +73,7 @@ public struct UnownedSerialExecutor: Sendable {
     fatalError("Swift compiler is incompatible with this SDK version")
     #endif
   }
+
 }
 
 /// Checks if the current task is running on the expected executor.

stdlib/public/Concurrency/ExecutorAssertions.swift

@@ -218,7 +218,7 @@ func assumeOnMainActorExecutor<T>(
 /// if another actor uses the same serial executor--by using that actor's ``Actor/unownedExecutor``
 /// as its own ``Actor/unownedExecutor``--this check will succeed, as from a concurrency safety
 /// perspective, the serial executor guarantees mutual exclusion of those two actors.
-@available(SwiftStdlib 5.9, *) // FIXME: use @backDeploy(before: SwiftStdlib 5.9)
+@available(SwiftStdlib 5.9, *)
 @_unavailableFromAsync(message: "express the closure as an explicit function declared on the specified 'actor' instead")
 public
 func assumeOnActorExecutor<Act: Actor, T>(
@@ -246,4 +246,5 @@ func assumeOnActorExecutor<Act: Actor, T>(
 }
 
 // TODO(ktoso): implement assume for distributed actors as well
+
 #endif // not SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY

stdlib/public/Distributed/CMakeLists.txt

@@ -18,6 +18,7 @@ add_swift_target_library(swiftDistributed ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} IS
   DistributedActor.cpp
   DistributedActor.swift
   DistributedActorSystem.swift
+  DistributedAssertions.swift
   DistributedMetadata.swift
   LocalTestingDistributedActorSystem.swift
 

stdlib/public/Distributed/DistributedAssertions.swift

@@ -0,0 +1,50 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Swift
+@_spi(ConcurrencyExecutors) import _Concurrency
+
+@available(SwiftStdlib 5.9, *)
+@_unavailableFromAsync(message: "express the closure as an explicit function declared on the specified 'distributed actor' instead")
+public
+func assumeOnLocalDistributedActorExecutor<Act: DistributedActor, T>(
+    _ actor: Act,
+    _ operation: (isolated Act) throws -> T,
+    file: StaticString = #fileID, line: UInt = #line
+) rethrows -> T {
+  typealias YesActor = (isolated Act) throws -> T
+  typealias NoActor = (Act) throws -> T
+
+  guard __isLocalActor(actor) else {
+    fatalError("Cannot assume to be 'isolated \(Act.self)' since distributed actor '\(actor)' is remote.")
+  }
+
+  /// This is guaranteed to be fatal if the check fails,
+  /// as this is our "safe" version of this API.
+  let executor: Builtin.Executor = actor.unownedExecutor._executor
+  guard _taskIsCurrentExecutor(executor) else {
+    // TODO: offer information which executor we actually got when
+    fatalError("Incorrect actor executor assumption; Expected same executor as \(actor).", file: file, line: line)
+  }
+
+  // To do the unsafe cast, we have to pretend it's @escaping.
+  return try withoutActuallyEscaping(operation) {
+    (_ fn: @escaping YesActor) throws -> T in
+    let rawFn = unsafeBitCast(fn, to: NoActor.self)
+    return try rawFn(actor)
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+@usableFromInline
+@_silgen_name("swift_task_isCurrentExecutor")
+func _taskIsCurrentExecutor(_ executor: Builtin.Executor) -> Bool

test/Distributed/Runtime/distributed_actor_assume_executor.swift

@@ -0,0 +1,107 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/../Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-build-swift -Xfrontend -disable-availability-checking -parse-as-library -I %t %s %S/../Inputs/FakeDistributedActorSystems.swift -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN:  %target-run %t/a.out
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+
+// UNSUPPORTED: back_deploy_concurrency
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: freestanding
+
+import StdlibUnittest
+import Distributed
+import FakeDistributedActorSystems
+
+typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
+
+func checkAssumeLocalDistributedActor(actor: MainFriend) /* synchronous! */ -> String {
+  assumeOnLocalDistributedActorExecutor(actor) { dist in
+    print("gained access to: \(dist.isolatedProperty)")
+    return dist.isolatedProperty
+  }
+}
+
+func checkAssumeMainActor(actor: MainFriend) /* synchronous! */ {
+  assumeOnMainActorExecutor {
+    print("yay")
+  }
+}
+
+@MainActor
+func check(actor: MainFriend) {
+  _ = checkAssumeLocalDistributedActor(actor: actor)
+  checkAssumeMainActor(actor: actor)
+}
+
+distributed actor MainFriend {
+  nonisolated var unownedExecutor: UnownedSerialExecutor {
+    print("get unowned executor")
+    return MainActor.sharedUnownedExecutor
+  }
+
+  let isolatedProperty: String = "Hello there!"
+
+  distributed func test(x: Int) async throws {
+    print("executed: \(#function)")
+    defer {
+      print("done executed: \(#function)")
+    }
+    return checkAssumeMainActor(actor: self)
+  }
+
+}
+
+actor OtherMain {
+  nonisolated var unownedExecutor: UnownedSerialExecutor {
+    return MainActor.sharedUnownedExecutor
+  }
+
+  func checkAssumeLocalDistributedActor(actor: MainFriend) /* synchronous! */ {
+    _ = assumeOnLocalDistributedActorExecutor(actor) { dist in
+      print("gained access to: \(dist.isolatedProperty)")
+      return dist.isolatedProperty
+    }
+  }
+}
+
+@main struct Main {
+  static func main() async {
+    let tests = TestSuite("AssumeLocalDistributedActorExecutor")
+
+    let system = FakeRoundtripActorSystem()
+    let distLocal = MainFriend(actorSystem: system)
+
+    if #available(SwiftStdlib 5.9, *) {
+
+      tests.test("assumeOnLocalDistributedActorExecutor: assume the main executor, inside the DistributedMainFriend local actor") {
+        _ = checkAssumeLocalDistributedActor(actor: distLocal)
+        try! await distLocal.test(x: 42)
+      }
+
+      tests.test("assumeOnLocalDistributedActorExecutor: assume same actor as the DistributedMainFriend") {
+        await OtherMain().checkAssumeLocalDistributedActor(actor: distLocal)
+        try! await distLocal.test(x: 42)
+      }
+
+      tests.test("assumeOnLocalDistributedActorExecutor: wrongly assume the same actor as the DistributedmainFriend") {
+        await OtherMain().checkAssumeLocalDistributedActor(actor: distLocal)
+      }
+
+      tests.test("assumeOnLocalDistributedActorExecutor: on remote actor reference") {
+        expectCrashLater(withMessage: "Cannot assume to be 'isolated MainFriend' since distributed actor 'a.MainFriend' is remote.")
+        let remoteRef = try! MainFriend.resolve(id: distLocal.id, using: system)
+        await OtherMain().checkAssumeLocalDistributedActor(actor: remoteRef)
+      }
+
+
+    }
+
+    await runAllTestsAsync()
+  }
+}