Merge pull request #65952 from ktoso/wip-destroy-dist-args-reverse-order

[Distributed] Ensure to swift_release during destroying decoded distributed call arguments

Author: swift-ci

lib/IRGen/GenDistributed.cpp

@@ -581,6 +581,15 @@ void DistributedAccessor::emitLoadOfWitnessTables(llvm::Value *witnessTables,
 }
 
 void DistributedAccessor::emitReturn(llvm::Value *errorValue) {
+  // Destroy loaded arguments.
+  // This MUST be done before deallocating, as otherwise we'd try to
+  // swift_release freed memory, which will be a no-op, however that also would
+  // mean we never drop retain counts to 0 and miss to run deinitializers of
+  // classes!
+  llvm::for_each(LoadedArguments, [&](const auto &argInfo) {
+    emitDestroyCall(IGF, argInfo.second, argInfo.first);
+  });
+
   // Deallocate all of the copied arguments. Since allocations happened
   // on stack they have to be deallocated in reverse order.
   {
@@ -590,11 +599,6 @@ void DistributedAccessor::emitReturn(llvm::Value *errorValue) {
     }
   }
 
-  // Destroy loaded arguments.
-  llvm::for_each(LoadedArguments, [&](const auto &argInfo) {
-    emitDestroyCall(IGF, argInfo.second, argInfo.first);
-  });
-
   Explosion voidResult;
 
   Explosion error;

test/Distributed/Inputs/FakeDistributedActorSystems.swift

@@ -226,6 +226,10 @@ public final class FakeRoundtripActorSystem: DistributedActorSystem, @unchecked
 
   public init() {}
 
+  public func shutdown() {
+    self.activeActors = [:]
+  }
+
   public func resolve<Act>(id: ActorID, as actorType: Act.Type)
     throws -> Act? where Act: DistributedActor {
     print("| resolve \(id) as remote // this system always resolves as remote")
@@ -394,7 +398,14 @@ public struct FakeInvocationEncoder : DistributedTargetInvocationEncoder {
     print(" > done recording")
   }
 
-  public func makeDecoder() -> FakeInvocationDecoder {
+  public mutating func makeDecoder() -> FakeInvocationDecoder {
+    defer {
+      // reset the decoder; we don't want to keep these values retained by accident here
+      genericSubs = []
+      arguments = []
+      returnType = nil
+      errorType = nil
+    }
     return .init(
       args: arguments,
       substitutions: genericSubs,

test/Distributed/Runtime/distributed_actor_func_calls_decoded_args_deinit.swift

@@ -0,0 +1,66 @@
+// 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 -O -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 --color --dump-input=always
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+// FIXME(distributed): Distributed actors currently have some issues on windows, isRemote always returns false. rdar://82593574
+// UNSUPPORTED: OS=windows-msvc
+
+import Distributed
+import FakeDistributedActorSystems
+
+typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
+
+final class SomeClass<T>: Sendable, Codable {
+  let file: String
+  let line: UInt
+  init(file: String = #fileID, line: UInt = #line) {
+    self.file = file
+    self.line = line
+    print("SomeClass: init @ \(file):\(line)")
+  }
+  deinit {
+    print("SomeClass: deinit @ \(file):\(line)")
+  }
+}
+
+struct S<T> : Codable {
+  var data: SomeClass<T>
+}
+
+distributed actor Greeter {
+  distributed func test1<T>(_ value: SomeClass<T>) {
+  }
+  distributed func test2<T>(_ value: S<T>) {}
+}
+
+func test() async throws {
+  let system = DefaultDistributedActorSystem()
+  defer { system.shutdown() }
+
+  let local = Greeter(actorSystem: system)
+  let ref = try Greeter.resolve(id: local.id, using: system)
+
+  try await ref.test1(SomeClass<Int>())
+  // CHECK: SomeClass: init
+  // CHECK: SomeClass: deinit
+
+  try await ref.test2(S(data: SomeClass<Int>()))
+  // CHECK: SomeClass: init
+  // CHECK: SomeClass: deinit
+}
+
+@main struct Main {
+  static func main() async {
+    try! await test()
+  }
+}

test/Distributed/distributed_actor_thunk_doesnt_leak_class_arguments.swift

@@ -25,18 +25,20 @@ struct S<T> : Codable {
 
 distributed actor Greeter {
   // CHECK-LABEL: define linkonce_odr hidden swifttailcc void @"$s15no_to_arg_leaks7GreeterC5test1yyAA9SomeClassCyxGYaKlFTETF"
-  // CHECK: [[ARG_ADDR:%.*]] = bitcast i8* {{.*}} to %T15no_to_arg_leaks9SomeClassC**
+  // CHECK: [[ARG_ADDR:%.*]] = bitcast i8* [[PARAM:%.*]] to %T15no_to_arg_leaks9SomeClassC**
   // CHECK: %destroy = bitcast i8* {{.*}} to void (%swift.opaque*, %swift.type*)*
   // CHECK-NEXT: [[OPAQUE_ARG_ADDR:%.*]] = bitcast %T15no_to_arg_leaks9SomeClassC** [[ARG_ADDR]] to %swift.opaque*
   // CHECK-NEXT: call void %destroy(%swift.opaque* noalias [[OPAQUE_ARG_ADDR]], %swift.type* %arg_type)
+  // CHECK-NEXT: call swiftcc void @swift_task_dealloc(i8* [[PARAM]])
   distributed func test1<T>(_: SomeClass<T>) {
   }
 
   // CHECK-LABEL: define linkonce_odr hidden swifttailcc void @"$s15no_to_arg_leaks7GreeterC5test2yyAA1SVyxGYaKlFTETF"
-  // CHECK: [[ARG_ADDR:%.*]] = bitcast i8* {{.*}} to %T15no_to_arg_leaks1SV*
+  // CHECK: [[ARG_ADDR:%.*]] = bitcast i8* [[PARAM:%.*]] to %T15no_to_arg_leaks1SV*
   // CHECK: %destroy = bitcast i8* {{.*}} to void (%swift.opaque*, %swift.type*)*
   // CHECK-NEXT: [[OPAQUE_ARG_ADDR:%.*]] = bitcast %T15no_to_arg_leaks1SV* [[ARG_ADDR]] to %swift.opaque*
   // CHECK-NEXT: call void %destroy(%swift.opaque* noalias [[OPAQUE_ARG_ADDR]], %swift.type* %arg_type)
+  // CHECK-NEXT: call swiftcc void @swift_task_dealloc(i8* [[PARAM]])
   distributed func test2<T>(_: S<T>) {}
 }