From fc3df01911d5d1a27c8e81cd47c2b939fd3fd42c Mon Sep 17 00:00:00 2001
From: Maxim Fateev <mfateev@gmail.com>
Date: Fri, 29 Dec 2023 10:51:02 -0800
Subject: [PATCH] Copy retryOptions when continue as new

---
 .../OpenTracingWorkflowOutboundCallsInterceptor.java   |  8 ++++----
 .../temporal/internal/replay/BasicWorkflowContext.java |  1 +
 .../io/temporal/internal/sync/SyncWorkflowContext.java | 10 +++++++++-
 .../java/io/temporal/workflow/ContinueAsNewTest.java   |  9 +++++++--
 4 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/temporal-opentracing/src/main/java/io/temporal/opentracing/internal/OpenTracingWorkflowOutboundCallsInterceptor.java b/temporal-opentracing/src/main/java/io/temporal/opentracing/internal/OpenTracingWorkflowOutboundCallsInterceptor.java
index 5d6661d53..7b6add447 100644
--- a/temporal-opentracing/src/main/java/io/temporal/opentracing/internal/OpenTracingWorkflowOutboundCallsInterceptor.java
+++ b/temporal-opentracing/src/main/java/io/temporal/opentracing/internal/OpenTracingWorkflowOutboundCallsInterceptor.java
@@ -177,11 +177,11 @@ private <R> Tracer.SpanBuilder createChildWorkflowStartSpanBuilder(ChildWorkflow
   }
 
   private Tracer.SpanBuilder createContinueAsNewWorkflowStartSpanBuilder(ContinueAsNewInput input) {
-    WorkflowInfo parentWorkflowInfo = Workflow.getInfo();
+    WorkflowInfo continuedWorkflowInfo = Workflow.getInfo();
     return spanFactory.createContinueAsNewWorkflowStartSpan(
         tracer,
-        MoreObjects.firstNonNull(input.getWorkflowType(), parentWorkflowInfo.getWorkflowType()),
-        parentWorkflowInfo.getWorkflowId(),
-        parentWorkflowInfo.getRunId());
+        MoreObjects.firstNonNull(input.getWorkflowType(), continuedWorkflowInfo.getWorkflowType()),
+        continuedWorkflowInfo.getWorkflowId(),
+        continuedWorkflowInfo.getRunId());
   }
 }
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/BasicWorkflowContext.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/BasicWorkflowContext.java
index 997b4d72a..3e0932b0b 100644
--- a/temporal-sdk/src/main/java/io/temporal/internal/replay/BasicWorkflowContext.java
+++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/BasicWorkflowContext.java
@@ -150,6 +150,7 @@ public Failure getPreviousRunFailure() {
     return previousRunFailure;
   }
 
+  @Nullable
   public RetryOptions getRetryOptions() {
     if (!startedAttributes.hasRetryPolicy()) {
       return null;
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java b/temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java
index 14797e09c..4d3768ca4 100644
--- a/temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java
+++ b/temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java
@@ -34,7 +34,13 @@
 import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes;
 import io.temporal.api.command.v1.SignalExternalWorkflowExecutionCommandAttributes;
 import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributes;
-import io.temporal.api.common.v1.*;
+import io.temporal.api.common.v1.ActivityType;
+import io.temporal.api.common.v1.Memo;
+import io.temporal.api.common.v1.Payload;
+import io.temporal.api.common.v1.Payloads;
+import io.temporal.api.common.v1.SearchAttributes;
+import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.api.common.v1.WorkflowType;
 import io.temporal.api.enums.v1.ParentClosePolicy;
 import io.temporal.api.failure.v1.Failure;
 import io.temporal.api.history.v1.HistoryEvent;
@@ -1104,6 +1110,8 @@ public void continueAsNew(ContinueAsNewInput input) {
       }
       if (options.getRetryOptions() != null) {
         attributes.setRetryPolicy(toRetryPolicy(options.getRetryOptions()));
+      } else if (replayContext.getRetryOptions() != null) {
+        attributes.setRetryPolicy(toRetryPolicy(replayContext.getRetryOptions()));
       }
       Map<String, Object> searchAttributes = options.getSearchAttributes();
       if (searchAttributes != null && !searchAttributes.isEmpty()) {
diff --git a/temporal-sdk/src/test/java/io/temporal/workflow/ContinueAsNewTest.java b/temporal-sdk/src/test/java/io/temporal/workflow/ContinueAsNewTest.java
index eabddb4ea..5f6f9f5aa 100644
--- a/temporal-sdk/src/test/java/io/temporal/workflow/ContinueAsNewTest.java
+++ b/temporal-sdk/src/test/java/io/temporal/workflow/ContinueAsNewTest.java
@@ -86,7 +86,7 @@ public static class TestContinueAsNewImpl implements TestContinueAsNew {
     @Override
     public int execute(int count, String continueAsNewTaskQueue) {
       String taskQueue = Workflow.getInfo().getTaskQueue();
-      if (count == INITIAL_COUNT) {
+      if (count >= INITIAL_COUNT - 2) {
         assertEquals(10, Workflow.getInfo().getRetryOptions().getMaximumAttempts());
       } else {
         assertEquals(5, Workflow.getInfo().getRetryOptions().getMaximumAttempts());
@@ -97,7 +97,12 @@ public int execute(int count, String continueAsNewTaskQueue) {
       }
       Map<String, Object> memo = new HashMap<>();
       memo.put("myKey", "MyValue");
-      RetryOptions retryOptions = RetryOptions.newBuilder().setMaximumAttempts(5).build();
+      RetryOptions retryOptions = null;
+      // don't specify retryOptions on the first continue-as-new to test that they are copied from
+      // the previous run.
+      if (count < INITIAL_COUNT - 1) {
+        retryOptions = RetryOptions.newBuilder().setMaximumAttempts(5).build();
+      }
       SearchAttributes searchAttributes =
           SearchAttributes.newBuilder()
               .set(SearchAttributeKey.forKeyword("CustomKeywordField"), "foo1")