Merge pull request #76862 from ktoso/wip-distributed-leak-fixes

Author: ktoso

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -531,8 +531,8 @@ extension DistributedActorSystem {
     }
 
     // Decode the return type
-    func allocateReturnTypeBuffer<R>(_: R.Type) -> UnsafeRawPointer? {
-      return UnsafeRawPointer(UnsafeMutablePointer<R>.allocate(capacity: 1))
+    func doAllocateReturnTypeBuffer<R>(_: R.Type) -> UnsafeMutableRawPointer? {
+      return UnsafeMutableRawPointer(UnsafeMutablePointer<R>.allocate(capacity: 1))
     }
 
     let maybeReturnTypeFromTypeInfo =
@@ -549,18 +549,23 @@ extension DistributedActorSystem {
         errorCode: .typeDeserializationFailure)
     }
 
-    guard let resultBuffer = _openExistential(returnTypeFromTypeInfo, do: allocateReturnTypeBuffer) else {
+    guard let resultBuffer = _openExistential(returnTypeFromTypeInfo, do: doAllocateReturnTypeBuffer) else {
       throw ExecuteDistributedTargetError(
         message: "Failed to allocate buffer for distributed target return type",
         errorCode: .typeDeserializationFailure)
     }
 
-    func destroyReturnTypeBuffer<R>(_: R.Type) {
-      resultBuffer.assumingMemoryBound(to: R.self).deallocate()
+    var executeDistributedTargetHasThrown = false
+    func doDestroyReturnTypeBuffer<R>(_: R.Type) {
+      let buf = resultBuffer.assumingMemoryBound(to: R.self)
+      if !executeDistributedTargetHasThrown {
+        buf.deinitialize(count: 1)
+      }
+      buf.deallocate()
     }
 
     defer {
-      _openExistential(returnTypeFromTypeInfo, do: destroyReturnTypeBuffer)
+      _openExistential(returnTypeFromTypeInfo, do: doDestroyReturnTypeBuffer)
     }
 
     do {
@@ -590,6 +595,7 @@ extension DistributedActorSystem {
         )
       }
     } catch {
+      executeDistributedTargetHasThrown = true
       try await handler.onThrow(error: error)
     }
   }