Make sure task failures use a serialization context (#2022)

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java

@@ -40,9 +40,11 @@
 import io.temporal.api.taskqueue.v1.StickyExecutionAttributes;
 import io.temporal.api.taskqueue.v1.TaskQueue;
 import io.temporal.api.workflowservice.v1.*;
+import io.temporal.common.converter.DataConverter;
 import io.temporal.internal.common.ProtobufTimeUtils;
 import io.temporal.internal.common.WorkflowExecutionUtils;
 import io.temporal.internal.worker.*;
+import io.temporal.payload.context.WorkflowSerializationContext;
 import io.temporal.serviceclient.MetricsTag;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.worker.NonDeterministicException;
@@ -173,7 +175,12 @@ private Result handleWorkflowTaskWithQuery(
         return createDirectQueryResult(workflowTask, null, e);
       } else {
         // this call rethrows an exception in some scenarios
-        return failureToWFTResult(workflowTask, e);
+        DataConverter dataConverterWithWorkflowContext =
+            options
+                .getDataConverter()
+                .withContext(
+                    new WorkflowSerializationContext(namespace, execution.getWorkflowId()));
+        return failureToWFTResult(workflowTask, e, dataConverterWithWorkflowContext);
       }
     } finally {
       if (!useCache && workflowRunTaskHandler != null) {
@@ -254,7 +261,8 @@ private Result createCompletedWFTRequest(
   }
 
   private Result failureToWFTResult(
-      PollWorkflowTaskQueueResponseOrBuilder workflowTask, Throwable e) throws Exception {
+      PollWorkflowTaskQueueResponseOrBuilder workflowTask, Throwable e, DataConverter dc)
+      throws Exception {
     String workflowType = workflowTask.getWorkflowType().getName();
     if (e instanceof WorkflowExecutionException) {
       RespondWorkflowTaskCompletedRequest response =
@@ -299,7 +307,7 @@ private Result failureToWFTResult(
       throw (Exception) e;
     }
 
-    Failure failure = options.getDataConverter().exceptionToFailure(e);
+    Failure failure = dc.exceptionToFailure(e);
     RespondWorkflowTaskFailedRequest.Builder failedRequest =
         RespondWorkflowTaskFailedRequest.newBuilder()
             .setTaskToken(workflowTask.getTaskToken())

temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflow.java

@@ -167,7 +167,10 @@ public void handleUpdate(
               // Rethrow instead on rejecting the update to fail the WFT
               throw r;
             } catch (Exception e) {
-              callbacks.reject(this.dataConverter.exceptionToFailure(e));
+              callbacks.reject(
+                  workflowContext
+                      .getDataConverterWithCurrentWorkflowContext()
+                      .exceptionToFailure(e));
               return;
             } finally {
               workflowContext.setReadOnly(false);

temporal-sdk/src/test/java/io/temporal/functional/serialization/WorkflowIdSignedPayloadsTest.java

@@ -29,8 +29,8 @@
 import io.temporal.activity.*;
 import io.temporal.api.common.v1.Payload;
 import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.api.enums.v1.EventType;
 import io.temporal.client.*;
-import io.temporal.client.schedules.*;
 import io.temporal.common.converter.*;
 import io.temporal.failure.CanceledFailure;
 import io.temporal.payload.codec.PayloadCodec;
@@ -46,6 +46,7 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.util.*;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -71,15 +72,17 @@ public class WorkflowIdSignedPayloadsTest {
   private static final DataConverter codecDataConverter =
       new CodecDataConverter(
           DefaultDataConverter.STANDARD_INSTANCE,
-          Collections.singletonList(new PayloadEncoderWithWorkflowIdSignature()));
+          Collections.singletonList(new PayloadEncoderWithWorkflowIdSignature()),
+          true);
 
   @Rule
   public SDKTestWorkflowRule testWorkflowRule =
       SDKTestWorkflowRule.newBuilder()
           .setWorkflowTypes(
               SimpleWorkflowWithAnActivity.class,
               TestWorkflowWithCronScheduleImpl.class,
-              DynamicWorkflowImpl.class)
+              DynamicWorkflowImpl.class,
+              WorkflowWithQuery.class)
           .setWorkflowClientOptions(
               WorkflowClientOptions.newBuilder().setDataConverter(codecDataConverter).build())
           .setActivityImplementations(
@@ -173,6 +176,26 @@ public void testDynamicWorkflow() {
     assertEquals("Hello World", workflow.getResult(String.class));
   }
 
+  @Test
+  public void testWorkflowWithTaskFailure() {
+    WorkflowOptions options =
+        SDKTestOptions.newWorkflowOptionsWithTimeouts(testWorkflowRule.getTaskQueue());
+    TestWorkflows.TestWorkflowReturnString workflow =
+        testWorkflowRule
+            .getWorkflowClient()
+            .newWorkflowStub(TestWorkflows.TestWorkflowReturnString.class, options);
+    assertEquals("Hello World", workflow.execute());
+    // Test that the task failure is recorded in the history
+    // if the serialization fails the workflow task will timeout.
+    assertEquals(
+        1,
+        testWorkflowRule
+            .getHistoryEvents(
+                WorkflowStub.fromTyped(workflow).getExecution().getWorkflowId(),
+                EventType.EVENT_TYPE_WORKFLOW_TASK_FAILED)
+            .size());
+  }
+
   @ActivityInterface
   public interface SimpleActivity {
     @ActivityMethod(name = "simple")
@@ -213,6 +236,18 @@ public String execute(String input) {
     }
   }
 
+  public static class WorkflowWithQuery implements TestWorkflows.TestWorkflowReturnString {
+    static AtomicInteger taskRetryCount = new AtomicInteger();
+
+    @Override
+    public String execute() {
+      if (taskRetryCount.incrementAndGet() == 1) {
+        throw new RuntimeException("test");
+      }
+      return "Hello World";
+    }
+  }
+
   public static class SimpleWorkflowWithAnActivity implements TestWorkflows.TestWorkflow1 {
 
     private final SimpleActivity activity =