🍒[5.9][IRGen] Distributed: Destroy loaded arguments after the call

lib/IRGen/GenDistributed.cpp

@@ -134,6 +134,9 @@ class DistributedAccessor {
   /// The list of all arguments that were allocated on the stack.
   SmallVector<StackAddress, 4> AllocatedArguments;
 
+  /// The list of all the arguments that were loaded.
+  SmallVector<std::pair<Address, /*type=*/llvm::Value *>, 4> LoadedArguments;
+
 public:
   DistributedAccessor(IRGenFunction &IGF, SILFunction *target,
                       CanSILFunctionType accessorTy);
@@ -498,13 +501,16 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
         resultValue.getAddress(), IGM.getStorageType(paramTy));
 
     cast<LoadableTypeInfo>(paramInfo).loadAsTake(IGF, eltPtr, arguments);
+    LoadedArguments.push_back(std::make_pair(eltPtr, argumentType));
     break;
   }
 
   case ParameterConvention::Direct_Owned: {
     // Copy the value out at +1.
     cast<LoadableTypeInfo>(paramInfo).loadAsCopy(IGF, resultValue.getAddress(),
                                                  arguments);
+    LoadedArguments.push_back(
+        std::make_pair(resultValue.getAddress(), argumentType));
     break;
   }
   }
@@ -575,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.
   {

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

@@ -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-swift-frontend -module-name no_to_arg_leaks -emit-irgen -disable-availability-checking -I %t 2>&1 %s | %IRGenFileCheck %s -check-prefix CHECK-%target-import-type
+
+// UNSUPPORTED: back_deploy_concurrency
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// REQUIRES: CPU=x86_64 || CPU=arm64
+
+// UNSUPPORTED: OS=windows-msvc
+
+import Distributed
+import FakeDistributedActorSystems
+
+@available(SwiftStdlib 5.5, *)
+typealias DefaultDistributedActorSystem = FakeActorSystem
+
+final class SomeClass<T>: Sendable, Codable {
+  init() {}
+}
+
+struct S<T> : Codable {
+  var data: SomeClass<T>
+}
+
+distributed actor Greeter {
+  // CHECK-LABEL: define linkonce_odr hidden swifttailcc void @"$s15no_to_arg_leaks7GreeterC5test1yyAA9SomeClassCyxGYaKlFTETF"
+  // 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* [[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>) {}
+}
+
+func test() async throws {
+  let system = DefaultDistributedActorSystem()
+
+  let local = Greeter(actorSystem: system)
+  let ref = try Greeter.resolve(id: local.id, using: system)
+
+  try await ref.test1(SomeClass<Int>())
+  try await ref.test2(S(data: SomeClass<Int>()))
+}