[Distributed] Fix _executeDistributedFunction resume and result buffer (#40814)

* [Distributed] Adjust interface of `swift_distributed_execute_target`

Since this is a special function, `calleeContext` doesn't point to
a direct parent but instead both parent context (uninitialized)
and resume function are passed as last arguments which means that
`callContext` has to act as an intermediate context in call to accessor.

* [Distributed] Drop optionality from result buffer in `_executeDistributedTarget`

`RawPointer?` is lowered into a two arguments since it's a struct,
to make it easy let's just allocate an empty pointer for `Void` result.

* [Distributed] NFC: Update _remoteCall test-case to check multiple different result types

Author: xedin

stdlib/public/Concurrency/Actor.cpp

@@ -1714,7 +1714,12 @@ findDistributedAccessor(const char *targetNameStart, size_t targetNameLength) {
 ///    argumentBuffer: Builtin.RawPointer,
 ///    resultBuffer: Builtin.RawPointer) async throws
 using TargetExecutorSignature =
-    AsyncSignature<void(DefaultActor *, const char *, size_t, void *, void *),
+    AsyncSignature<void(/*on=*/DefaultActor *,
+                        /*targetName=*/const char *, /*targetNameSize=*/size_t,
+                        /*argumentBuffer=*/void *,
+                        /*resultBuffer=*/void *,
+                        /*resumeFunc=*/TaskContinuationFunction *,
+                        /*callContext=*/AsyncContext *),
                    /*throws=*/true>;
 
 SWIFT_CC(swiftasync)
@@ -1743,7 +1748,9 @@ static void ::swift_distributed_execute_target_resume(
       reinterpret_cast<TargetExecutorSignature::ContinuationType *>(
           parentCtx->ResumeParent);
   swift_task_dealloc(context);
-  return resumeInParent(parentCtx, error);
+  // See `swift_distributed_execute_target` - `parentCtx` in this case
+  // is `callContext` which should be completely transparent on resume.
+  return resumeInParent(parentCtx->Parent, error);
 }
 
 SWIFT_CC(swiftasync)
@@ -1752,7 +1759,9 @@ void ::swift_distributed_execute_target(
     DefaultActor *actor,
     const char *targetNameStart, size_t targetNameLength,
     void *argumentBuffer,
-    void *resultBuffer) {
+    void *resultBuffer,
+    TaskContinuationFunction *resumeFunc,
+    AsyncContext *callContext) {
   auto *accessor = findDistributedAccessor(targetNameStart, targetNameLength);
   if (!accessor) {
     assert(false && "no distributed accessor accessor");
@@ -1770,7 +1779,17 @@ void ::swift_distributed_execute_target(
   AsyncContext *calleeContext = reinterpret_cast<AsyncContext *>(
       swift_task_alloc(asyncFnPtr->ExpectedContextSize));
 
-  calleeContext->Parent = callerContext;
+  // TODO(concurrency): Special functions like this one are currently set-up
+  // to pass "caller" context and resume function as extra parameters due to
+  // how they are declared in C. But this particular function behaves exactly
+  // like a regular `async throws`, which means that we need to initialize
+  // intermediate `callContext` using parent `callerContext`. A better fix for
+  // this situation would be to adjust IRGen and handle function like this
+  // like regular `async` functions even though they are classified as special.
+  callContext->Parent = callerContext;
+  callContext->ResumeParent = resumeFunc;
+
+  calleeContext->Parent = callContext;
   calleeContext->ResumeParent = reinterpret_cast<TaskContinuationFunction *>(
       swift_distributed_execute_target_resume);
 

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -222,18 +222,20 @@ extension DistributedActorSystem {
 
     // Decode the return type
     func allocateReturnTypeBuffer<R>(_: R.Type) -> UnsafeRawPointer? {
-      if R.self == Void.self {
-        return nil
-      }
       return UnsafeRawPointer(UnsafeMutablePointer<R>.allocate(capacity: 1))
     }
     guard let returnType: Any.Type = _getReturnTypeInfo(mangledMethodName: mangledTargetName) else {
       throw ExecuteDistributedTargetError(
           message: "Failed to decode distributed target return type")
     }
-    let resultBuffer = _openExistential(returnType, do: allocateReturnTypeBuffer)
+
+    guard let resultBuffer = _openExistential(returnType, do: allocateReturnTypeBuffer) else {
+      throw ExecuteDistributedTargetError(
+          message: "Failed to allocate buffer for distributed target return type")
+    }
+
     func destroyReturnTypeBuffer<R>(_: R.Type) {
-      resultBuffer?.assumingMemoryBound(to: R.self).deallocate()
+      resultBuffer.assumingMemoryBound(to: R.self).deallocate()
     }
     defer {
       _openExistential(returnType, do: destroyReturnTypeBuffer)
@@ -252,16 +254,11 @@ extension DistributedActorSystem {
           on: actor,
           mangledTargetName, UInt(mangledTargetName.count),
           argumentBuffer: hargs.buffer._rawValue,
-          resultBuffer: resultBuffer?._rawValue
+          resultBuffer: resultBuffer._rawValue
       )
 
-      // Get the result out of the buffer and invoke onReturn with the right type
-      guard let resultBuffer = resultBuffer else {
-        try await handler.onReturn(value: ())
-        return
-      }
-      func onReturn<R>(_: R.Type) async throws {
-        try await handler.onReturn/*<R>*/(value: resultBuffer)
+      func onReturn<R>(_ resultTy: R.Type) async throws {
+        try await handler.onReturn/*<R>*/(value: resultBuffer.load(as: resultTy))
       }
 
       try await _openExistential(returnType, do: onReturn)
@@ -277,7 +274,7 @@ func _executeDistributedTarget(
   on actor: AnyObject, // DistributedActor
   _ targetName: UnsafePointer<UInt8>, _ targetNameLength: UInt,
   argumentBuffer: Builtin.RawPointer, // HeterogeneousBuffer of arguments
-  resultBuffer: Builtin.RawPointer?
+  resultBuffer: Builtin.RawPointer
 ) async throws
 
 // ==== ----------------------------------------------------------------------------------------------------------------

test/Distributed/Runtime/distributed_actor_remoteCall.swift

@@ -18,12 +18,37 @@
 import _Distributed
 
 final class Obj: @unchecked Sendable, Codable  {}
+
 struct LargeStruct: Sendable, Codable {
+  var q: String
+  var a: Int
+  var b: Int64
+  var c: Double
+  var d: String
+}
+
+enum E : Sendable, Codable {
+  case foo, bar
 }
 
 distributed actor Greeter {
-  distributed func hello() {
-    print("EXECUTING HELLO")
+  distributed func empty() {
+  }
+
+  distributed func hello() -> String {
+    return "Hello, World!"
+  }
+
+  distributed func answer() -> Int {
+    return 42
+  }
+
+  distributed func largeResult() -> LargeStruct {
+    .init(q: "question", a: 42, b: 1, c: 2.0, d: "Lorum ipsum")
+  }
+
+  distributed func enumResult() -> E {
+    .bar
   }
 }
 
@@ -102,7 +127,11 @@ struct FakeResultHandler: DistributedTargetInvocationResultHandler {
 typealias DefaultDistributedActorSystem = FakeActorSystem
 
 // actual mangled name:
-let helloName = "$s4main7GreeterC5helloyyFTE"
+let emptyName = "$s4main7GreeterC5emptyyyFTE"
+let helloName = "$s4main7GreeterC5helloSSyFTE"
+let answerName = "$s4main7GreeterC6answerSiyFTE"
+let largeResultName = "$s4main7GreeterC11largeResultAA11LargeStructVyFTE"
+let enumResult = "$s4main7GreeterC10enumResultAA1EOyFTE"
 
 func test() async throws {
   let system = FakeActorSystem()
@@ -112,14 +141,51 @@ func test() async throws {
   // act as if we decoded an Invocation:
   let invocation = FakeInvocation()
 
+  try await system.executeDistributedTarget(
+    on: local,
+    mangledTargetName: emptyName,
+    invocation: invocation,
+    handler: FakeResultHandler()
+  )
+
+  // CHECK: RETURN: ()
+
   try await system.executeDistributedTarget(
       on: local,
       mangledTargetName: helloName,
       invocation: invocation,
       handler: FakeResultHandler()
   )
 
-  // CHECK: done
+  // CHECK: RETURN: Hello, World!
+
+  try await system.executeDistributedTarget(
+      on: local,
+      mangledTargetName: answerName,
+      invocation: invocation,
+      handler: FakeResultHandler()
+  )
+
+  // CHECK: RETURN: 42
+
+  try await system.executeDistributedTarget(
+      on: local,
+      mangledTargetName: largeResultName,
+      invocation: invocation,
+      handler: FakeResultHandler()
+  )
+
+  // CHECK: RETURN: LargeStruct(q: "question", a: 42, b: 1, c: 2.0, d: "Lorum ipsum")
+
+  try await system.executeDistributedTarget(
+      on: local,
+      mangledTargetName: enumResult,
+      invocation: invocation,
+      handler: FakeResultHandler()
+  )
+  // CHECK: RETURN: bar
+
+  // CHECK-NEXT: done
   print("done")
 }