From bb8986c1e0013387ca1cc8e06a03458ff8255d7f Mon Sep 17 00:00:00 2001
From: Zijian Chen <zijian@uber.com>
Date: Mon, 19 Apr 2021 19:49:41 -0700
Subject: [PATCH] add all missing version marker before upsert searchattributes

---
 .../internal/replay/DecisionsHelper.java      |  2 ++
 .../uber/cadence/workflow/WorkflowTest.java   | 31 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/src/main/java/com/uber/cadence/internal/replay/DecisionsHelper.java b/src/main/java/com/uber/cadence/internal/replay/DecisionsHelper.java
index b209d1643..3f3ea787f 100644
--- a/src/main/java/com/uber/cadence/internal/replay/DecisionsHelper.java
+++ b/src/main/java/com/uber/cadence/internal/replay/DecisionsHelper.java
@@ -558,6 +558,8 @@ void recordMarker(String markerName, Header header, byte[] details) {
   }
 
   void upsertSearchAttributes(SearchAttributes searchAttributes) {
+    addAllMissingVersionMarker(false, Optional.empty());
+
     UpsertWorkflowSearchAttributesDecisionAttributes decisionAttr =
         new UpsertWorkflowSearchAttributesDecisionAttributes()
             .setSearchAttributes(searchAttributes);
diff --git a/src/test/java/com/uber/cadence/workflow/WorkflowTest.java b/src/test/java/com/uber/cadence/workflow/WorkflowTest.java
index 69894bf43..843f37786 100644
--- a/src/test/java/com/uber/cadence/workflow/WorkflowTest.java
+++ b/src/test/java/com/uber/cadence/workflow/WorkflowTest.java
@@ -4524,6 +4524,37 @@ public void testGetVersionRemovedInReplay2() {
     assertEquals("activity", workflowStub.query());
   }
 
+  // The following test covers the scenario where getVersion call is removed before
+  // upsertSearchAttributes.
+  public static class TestGetVersionRemovedInReplay3WorkflowImpl implements TestWorkflow1 {
+    @Override
+    public String execute(String taskList) {
+      Map<String, Object> searchAttrMap = new HashMap<>();
+      searchAttrMap.put("CustomKeywordField", "abc");
+      // Test removing a version check in replay code.
+      if (!getVersionExecuted.contains(taskList)) {
+        Workflow.getVersion("test_change", Workflow.DEFAULT_VERSION, 1);
+        Workflow.upsertSearchAttributes(searchAttrMap);
+        getVersionExecuted.add(taskList);
+      } else {
+        Workflow.upsertSearchAttributes(searchAttrMap);
+      }
+
+      return "done";
+    }
+  }
+
+  @Test
+  public void testGetVersionRemovedInReplay3() {
+    startWorkerFor(TestGetVersionRemovedInReplay3WorkflowImpl.class);
+    TestWorkflow1 workflowStub =
+        workflowClient.newWorkflowStub(
+            TestWorkflow1.class, newWorkflowOptionsBuilder(taskList).build());
+    String result = workflowStub.execute(taskList);
+    assertEquals("done", result);
+    tracer.setExpected("getVersion", "upsertSearchAttributes");
+  }
+
   public static class TestVersionNotSupportedWorkflowImpl implements TestWorkflow1 {
 
     @Override