[Distributed] IRGen: Emit an early return upon argument decoding failure

lib/IRGen/GenDistributed.cpp

@@ -123,6 +123,10 @@ private:
                                unsigned expectedWitnessTables,
                                Explosion &arguments);
 
+  /// Emit an async return from accessor which does cleanup of
+  /// all the argument allocations.
+  void emitReturn(llvm::Value *errorValue);
+
   FunctionPointer getPointerToTarget() const;
 
   Callee getCalleeForDistributedTarget(llvm::Value *self) const;
@@ -131,6 +135,12 @@ private:
   /// could be used to decode argument values to pass to its invocation.
   static ArgumentDecoderInfo findArgumentDecoder(IRGenModule &IGM,
                                                  SILFunction *thunk);
+
+  /// The result type of the accessor.
+  SILType getResultType() const;
+
+  /// The error type of this accessor.
+  SILType getErrorType() const;
 };
 
 } // end namespace
@@ -316,6 +326,9 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
   // substitution Argument -> <argument metadata>
   decodeArgs.add(argumentType);
 
+  Address calleeErrorSlot;
+  llvm::Value *decodeError = nullptr;
+
   emission->begin();
   {
     emission->setArgs(decodeArgs, /*isOutlined=*/false,
@@ -325,14 +338,43 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
     emission->emitToExplosion(result, /*isOutlined=*/false);
     assert(result.empty());
 
-    // TODO: Add error handling a new block that uses `emitAsyncReturn`
+    // Load error from the slot to emit an early return if necessary.
-    // if error slot is non-null.
+    {
+      SILFunctionConventions conv(ArgumentDecoder.Type, IGM.getSILModule());
+      SILType errorType =
+          conv.getSILErrorType(IGM.getMaximalTypeExpansionContext());
+
+      calleeErrorSlot =
+          emission->getCalleeErrorSlot(errorType, /*isCalleeAsync=*/true);
+      decodeError = IGF.Builder.CreateLoad(calleeErrorSlot);
+    }
   }
   emission->end();
 
   // Remember to deallocate later.
   AllocatedArguments.push_back(resultValue);
 
+  // Check whether the error slot has been set and if so
+  // emit an early return from accessor.
+  {
+    auto contBB = IGF.createBasicBlock("");
+    auto errorBB = IGF.createBasicBlock("on-error");
+
+    auto nullError = llvm::Constant::getNullValue(decodeError->getType());
+    auto hasError = IGF.Builder.CreateICmpNE(decodeError, nullError);
+
+    IGF.Builder.CreateCondBr(hasError, errorBB, contBB);
+    {
+      IGF.Builder.emitBlock(errorBB);
+      // Emit an early return if argument decoding failed.
+      emitReturn(decodeError);
+    }
+
+    IGF.Builder.emitBlock(contBB);
+    // Reset value of the slot back to `null`
+    IGF.Builder.CreateStore(nullError, calleeErrorSlot);
+  }
+
   switch (param.getConvention()) {
   case ParameterConvention::Indirect_In:
   case ParameterConvention::Indirect_In_Constant: {
@@ -412,10 +454,28 @@ void DistributedAccessor::emitLoadOfWitnessTables(llvm::Value *witnessTables,
   }
 }
 
+void DistributedAccessor::emitReturn(llvm::Value *errorValue) {
+  // Deallocate all of the copied arguments. Since allocations happened
+  // on stack they have to be deallocated in reverse order.
+  {
+    for (auto alloca = AllocatedArguments.rbegin();
+         alloca != AllocatedArguments.rend(); ++alloca) {
+      IGF.emitDeallocateDynamicAlloca(*alloca);
+    }
+  }
+
+  Explosion voidResult;
+
+  Explosion error;
+  error.add(errorValue);
+
+  emitAsyncReturn(IGF, AsyncLayout, getResultType(), AccessorType, voidResult,
+                  error);
+}
+
 void DistributedAccessor::emit() {
   auto targetTy = Target->getLoweredFunctionType();
   SILFunctionConventions targetConv(targetTy, IGF.getSILModule());
-  SILFunctionConventions accessorConv(AccessorType, IGF.getSILModule());
   TypeExpansionContext expansionContext = IGM.getMaximalTypeExpansionContext();
 
   auto params = IGF.collectParameters();
@@ -510,7 +570,7 @@ void DistributedAccessor::emit() {
   // using computed argument explosion.
   {
     Explosion result;
-    Explosion error;
+    llvm::Value *targetError = nullptr;
 
     auto callee = getCalleeForDistributedTarget(actorSelf);
     auto emission =
@@ -536,27 +596,16 @@ void DistributedAccessor::emit() {
     {
       assert(targetTy->hasErrorResult());
 
-      SILType errorType = accessorConv.getSILErrorType(expansionContext);
       Address calleeErrorSlot =
-          emission->getCalleeErrorSlot(errorType, /*isCalleeAsync=*/true);
+          emission->getCalleeErrorSlot(getErrorType(), /*isCalleeAsync=*/true);
-      error.add(IGF.Builder.CreateLoad(calleeErrorSlot));
+      targetError = IGF.Builder.CreateLoad(calleeErrorSlot);
     }
 
     emission->end();
 
-    // Deallocate all of the copied arguments. Since allocations happened
+    // Emit an async return that does allocation cleanup and propagates error
-    // on stack they have to be deallocated in reverse order.
+    // (if any) back to the caller.
-    {
+    emitReturn(targetError);
-      while (!AllocatedArguments.empty()) {
-        auto argument = AllocatedArguments.pop_back_val();
-        IGF.emitDeallocateDynamicAlloca(argument);
-      }
-    }
-
-    Explosion voidResult;
-    emitAsyncReturn(IGF, AsyncLayout,
-                    accessorConv.getSILResultType(expansionContext),
-                    AccessorType, voidResult, error);
   }
 }
 
@@ -604,6 +653,16 @@ DistributedAccessor::findArgumentDecoder(IRGenModule &IGM, SILFunction *thunk) {
   return {.Type = methodTy, .Fn = methodPtr};
 }
 
+SILType DistributedAccessor::getResultType() const {
+  SILFunctionConventions conv(AccessorType, IGF.getSILModule());
+  return conv.getSILResultType(IGM.getMaximalTypeExpansionContext());
+}
+
+SILType DistributedAccessor::getErrorType() const {
+  SILFunctionConventions conv(AccessorType, IGF.getSILModule());
+  return conv.getSILErrorType(IGM.getMaximalTypeExpansionContext());
+}
+
 Callee ArgumentDecoderInfo::getCallee(llvm::Value *decoder) const {
   CalleeInfo info(Type, Type, SubstitutionMap());
   return {std::move(info), Fn, decoder};

test/Distributed/Runtime/distributed_actor_remoteCall.swift

@@ -96,6 +96,9 @@ distributed actor Greeter {
   distributed func genericOptional<T: Codable>(t: T?) {
     print("---> T = \(t!), type(of:) = \(type(of: t))")
   }
+
+  distributed func expectsDecodeError(v: Int???) {
+  }
 }
 
 
@@ -203,6 +206,11 @@ class FakeInvocation: DistributedTargetInvocationEncoder, DistributedTargetInvoc
       fatalError("Cannot cast argument\(anyArgument) to expected \(Argument.self)")
     }
 
+    if (argumentIndex == 0 && Argument.self == Int???.self) {
+      throw ExecuteDistributedTargetError(message: "Failed to decode of Int??? (for a test)")
+    }
+
+
     argumentIndex += 1
     return argument
   }
@@ -244,6 +252,7 @@ let generic3Name = "$s4main7GreeterC8generic31a1b1cyx_Sayq_Gq0_tSeRzSERzSeR_SER_
 let generic4Name = "$s4main7GreeterC8generic41a1b1cyx_AA1SVyq_GSayq0_GtSeRzSERzSeR_SER_SeR0_SER0_r1_lFTE"
 let generic5Name = "$s4main7GreeterC8generic51a1b1c1dyx_AA1SVyq_Gq0_q1_tSeRzSERzSeR_SER_SeR0_SER0_SeR1_SER1_r2_lFTE"
 let genericOptionalName = "$s4main7GreeterC15genericOptional1tyxSg_tSeRzSERzlFTE"
+let expectsDecodeErrorName = "$s4main7GreeterC18expectsDecodeError1vySiSgSgSg_tFTE"
 
 func test() async throws {
   let system = DefaultDistributedActorSystem()
@@ -419,6 +428,19 @@ func test() async throws {
   // CHECK: ---> T = [0.0, 3737844653.0], type(of:) = Optional<Array<Double>>
   // CHECK-NEXT: RETURN: ()
 
+  let decodeErrInvocation = system.makeInvocationEncoder()
+
+  try decodeErrInvocation.recordArgument(42)
+  try decodeErrInvocation.doneRecording()
+
+  try await system.executeDistributedTarget(
+    on: local,
+    mangledTargetName: expectsDecodeErrorName,
+    invocationDecoder: decodeErrInvocation,
+    handler: FakeResultHandler()
+  )
+  // CHECK: ERROR: ExecuteDistributedTargetError(message: "Failed to decode of Int??? (for a test)")
+
   print("done")
   // CHECK-NEXT: done
 }