[Distributed] Avoid infinite recursion in distributed thunk on protocol extensions (#73032)

Author: ktoso

lib/SILGen/SILGenApply.cpp

@@ -682,20 +682,9 @@ class Callee {
     }
     case Kind::WitnessMethod: {
       if (auto func = constant->getFuncDecl()) {
-        auto isDistributedFuncOrAccessor =
-            func->isDistributed();
-        if (auto acc = dyn_cast<AccessorDecl>(func)) {
-          isDistributedFuncOrAccessor =
-              acc->getStorage()->isDistributed();
-        }
-        if (isa<ProtocolDecl>(func->getDeclContext()) && isDistributedFuncOrAccessor) {
-          // If we're calling cross-actor, we must always use a distributed thunk
-          if (!isSameActorIsolated(func, SGF.FunctionDC)) {
-            // the protocol witness must always be a distributed thunk, as we
-            // may be crossing a remote boundary here.
-            auto thunk = func->getDistributedThunk();
-            constant = SILDeclRef(thunk).asDistributed();
-          }
+        if (SGF.shouldReplaceConstantForApplyWithDistributedThunk(func)) {
+          auto thunk = func->getDistributedThunk();
+          constant = SILDeclRef(thunk).asDistributed();
         }
       }
 
@@ -783,12 +772,8 @@ class Callee {
     }
     case Kind::WitnessMethod: {
       if (auto func = constant->getFuncDecl()) {
-        if (func->isDistributed() && isa<ProtocolDecl>(func->getDeclContext())) {
-          // If we're calling cross-actor, we must always use a distributed thunk
-          if (!isSameActorIsolated(func, SGF.FunctionDC)) {
-            /// We must adjust the constant to use a distributed thunk.
-            constant = constant->asDistributed();
-          }
+        if (SGF.shouldReplaceConstantForApplyWithDistributedThunk(func)) {
+          constant = constant->asDistributed();
         }
       }
 

lib/SILGen/SILGenDistributed.cpp

@@ -274,6 +274,34 @@ void InitializeDistActorIdentity::dump(SILGenFunction &) const {
 #endif
 }
 
+bool SILGenFunction::shouldReplaceConstantForApplyWithDistributedThunk(
+    FuncDecl *func) const {
+  auto isDistributedFuncOrAccessor =
+      func->isDistributed();
+  if (auto acc = dyn_cast<AccessorDecl>(func)) {
+    isDistributedFuncOrAccessor =
+        acc->getStorage()->isDistributed();
+  }
+
+  if (!isDistributedFuncOrAccessor)
+    return false;
+
+  // If we are inside a distributed thunk, we want to call the "real" method,
+  // in order to avoid infinitely recursively calling the thunk from itself.
+  if (F.isDistributed() && F.isThunk())
+    return false;
+
+  // If caller and called func are isolated to the same (distributed) actor,
+  // (i.e. we are "inside the distributed actor"), there is no need to call
+  // the thunk.
+  if (isSameActorIsolated(func, FunctionDC))
+    return false;
+
+  // In all other situations, we may have to replace the called function,
+  // depending on isolation (to be checked in SILGenApply).
+  return true;
+}
+
 void SILGenFunction::emitDistributedActorImplicitPropertyInits(
     ConstructorDecl *ctor, ManagedValue selfArg) {
   // Only designated initializers should perform this initialization.

lib/SILGen/SILGenFunction.h

@@ -2452,6 +2452,34 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   // Distributed Actors
   //===---------------------------------------------------------------------===//
 
+  /// Determine if the target `func` should be replaced with a
+  /// 'distributed thunk'.
+  ///
+  /// This only applies to distributed functions when calls are made cross-actor
+  /// isolation. One notable exception is a distributed thunk calling the "real
+  /// underlying method", in which case (to avoid the thunk calling into itself,
+  /// the real method must be called).
+  ///
+  /// Witness calls which may need to be replaced with a distributed thunk call
+  /// happen either when the target type is generic, or if we are inside an
+  /// extension on a protocol. This method checks if we are in a context
+  /// where we should be calling the distributed thunk of the `func` or not.
+  /// Notably, if we are inside a distributed thunk already and are trying to
+  /// apply distributed method calls, all those must be to the "real" method,
+  /// because the thunks' responsibility is to call the real method, so this
+  /// replacement cannot be applied (or we'd recursively keep calling the same
+  /// thunk via witness).
+  ///
+  /// In situations which do not use a witness call, distributed methods are always
+  /// invoked Direct, and never ClassMethod, because distributed are effectively
+  /// final.
+  ///
+  /// \param func the target func that we are trying to "apply"
+  /// \return true when the function should be considered for replacement
+  ///         with distributed thunk when applying it
+  bool
+  shouldReplaceConstantForApplyWithDistributedThunk(FuncDecl *func) const;
+
   /// Initializes the implicit stored properties of a distributed actor that correspond to
   /// its transport and identity.
   void emitDistributedActorImplicitPropertyInits(

test/Distributed/Runtime/distributed_actor_func_calls_remoteCall_extension.swift

@@ -19,20 +19,28 @@ import FakeDistributedActorSystems
 
 typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
 
-extension FakeRoundtripActorSystemDistributedActor {
+protocol KappaProtocol : DistributedActor where ActorSystem == FakeRoundtripActorSystem {
+  distributed func echo(name: String) -> String
+}
+
+extension KappaProtocol {
   distributed func echo(name: String) -> String {
     return "Echo: \(name)"
   }
 }
 
+distributed actor KappaProtocolImpl: KappaProtocol {
+  // empty, gets default impl from extension on protocol
+}
+
 func test() async throws {
   let system = DefaultDistributedActorSystem()
 
-  let local = FakeRoundtripActorSystemDistributedActor(actorSystem: system)
-  let ref = try FakeRoundtripActorSystemDistributedActor.resolve(id: local.id, using: system)
+  let local = KappaProtocolImpl(actorSystem: system)
+  let ref = try KappaProtocolImpl.resolve(id: local.id, using: system)
 
   let reply = try await ref.echo(name: "Caplin")
-  // CHECK: >> remoteCall: on:main.FakeRoundtripActorSystemDistributedActor, target:main.FakeRoundtripActorSystemDistributedActor.echo(name:), invocation:FakeInvocationEncoder(genericSubs: [], arguments: ["Caplin"], returnType: Optional(Swift.String), errorType: nil), throwing:Swift.Never, returning:Swift.String
+  // CHECK: >> remoteCall: on:main.KappaProtocolImpl, target:main.$KappaProtocol.echo(name:), invocation:FakeInvocationEncoder(genericSubs: [main.KappaProtocolImpl], arguments: ["Caplin"], returnType: Optional(Swift.String), errorType: nil), throwing:Swift.Never, returning:Swift.String
 
   // CHECK: << remoteCall return: Echo: Caplin
   print("reply: \(reply)")

test/Distributed/Runtime/distributed_actor_func_calls_remoteCall_extension_concrete.swift

@@ -0,0 +1,54 @@
+// 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 -module-name main -Xfrontend -disable-availability-checking -j2 -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 | %FileCheck %s --dump-input=always
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+// UNSUPPORTED: OS=windows-msvc
+
+import Distributed
+import FakeDistributedActorSystems
+
+typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
+
+protocol KappaProtocol : DistributedActor where ActorSystem == FakeRoundtripActorSystem {
+  distributed func echo(name: String) -> String
+}
+
+distributed actor KappaProtocolImpl: KappaProtocol {
+  // empty, gets default impl from extension on this actor
+}
+
+extension KappaProtocolImpl {
+  distributed func echo(name: String) -> String {
+    return "Echo: \(name)"
+  }
+}
+
+func test() async throws {
+  let system = DefaultDistributedActorSystem()
+
+  let local = KappaProtocolImpl(actorSystem: system)
+  let ref = try KappaProtocolImpl.resolve(id: local.id, using: system)
+
+  let reply = try await ref.echo(name: "Caplin")
+  // CHECK: >> remoteCall: on:main.KappaProtocolImpl, target:main.KappaProtocolImpl.echo(name:), invocation:FakeInvocationEncoder(genericSubs: [], arguments: ["Caplin"], returnType: Optional(Swift.String), errorType: nil), throwing:Swift.Never, returning:Swift.String
+
+  // CHECK: << remoteCall return: Echo: Caplin
+  print("reply: \(reply)")
+  // CHECK: reply: Echo: Caplin
+}
+
+@main struct Main {
+  static func main() async {
+    try! await test()
+  }
+}