Add Upsert Search Attributes (#406)

vancexu · web-flow · commit d05514c659d6 · 2019-10-15T13:36:21.000-07:00
Support upsert search attributes inside workflow.
diff --git a/.gitignore b/.gitignore
@@ -7,3 +7,4 @@ target
 .gradle
 /build
 /out
+dummy
diff --git a/src/main/java/com/uber/cadence/internal/common/WorkflowExecutionUtils.java b/src/main/java/com/uber/cadence/internal/common/WorkflowExecutionUtils.java
@@ -972,7 +972,8 @@ public static boolean isDecisionEvent(HistoryEvent event) {
                 || eventType == EventType.CancelTimerFailed
                 || eventType == EventType.RequestCancelExternalWorkflowExecutionInitiated
                 || eventType == EventType.MarkerRecorded
-                || eventType == EventType.SignalExternalWorkflowExecutionInitiated));
+                || eventType == EventType.SignalExternalWorkflowExecutionInitiated
+                || eventType == EventType.UpsertWorkflowSearchAttributes));
     return result;
   }
 
diff --git a/src/main/java/com/uber/cadence/internal/replay/ClockDecisionContext.java b/src/main/java/com/uber/cadence/internal/replay/ClockDecisionContext.java
@@ -21,6 +21,7 @@
 import com.uber.cadence.ActivityType;
 import com.uber.cadence.HistoryEvent;
 import com.uber.cadence.MarkerRecordedEventAttributes;
+import com.uber.cadence.SearchAttributes;
 import com.uber.cadence.StartTimerDecisionAttributes;
 import com.uber.cadence.TimerCanceledEventAttributes;
 import com.uber.cadence.TimerFiredEventAttributes;
@@ -217,6 +218,14 @@ Optional<byte[]> mutableSideEffect(
     return mutableSideEffectHandler.handle(id, converter, func);
   }
 
+  void upsertSearchAttributes(SearchAttributes searchAttributes) {
+    decisions.upsertSearchAttributes(searchAttributes);
+  }
+
+  void handleUpsertSearchAttributesEvent(HistoryEvent event) {
+    // todo: update workflow info
+  }
+
   void handleMarkerRecorded(HistoryEvent event) {
     MarkerRecordedEventAttributes attributes = event.getMarkerRecordedEventAttributes();
     String name = attributes.getMarkerName();
diff --git a/src/main/java/com/uber/cadence/internal/replay/DecisionContext.java b/src/main/java/com/uber/cadence/internal/replay/DecisionContext.java
@@ -18,6 +18,7 @@
 package com.uber.cadence.internal.replay;
 
 import com.uber.cadence.ChildPolicy;
+import com.uber.cadence.SearchAttributes;
 import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.WorkflowType;
 import com.uber.cadence.converter.DataConverter;
@@ -179,4 +180,6 @@ Optional<byte[]> mutableSideEffect(
 
   /** @return replay safe UUID */
   UUID randomUUID();
+
+  void upsertSearchAttributes(SearchAttributes searchAttributes);
 }
diff --git a/src/main/java/com/uber/cadence/internal/replay/DecisionContextImpl.java b/src/main/java/com/uber/cadence/internal/replay/DecisionContextImpl.java
@@ -22,6 +22,7 @@
 import com.uber.cadence.DecisionTaskFailedEventAttributes;
 import com.uber.cadence.HistoryEvent;
 import com.uber.cadence.PollForDecisionTaskResponse;
+import com.uber.cadence.SearchAttributes;
 import com.uber.cadence.TimerFiredEventAttributes;
 import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.WorkflowExecutionStartedEventAttributes;
@@ -368,4 +369,14 @@ int numPendingLaTasks() {
   void awaitTaskCompletion(Duration duration) throws InterruptedException {
     workflowClock.awaitTaskCompletion(duration);
   }
+
+  @Override
+  public void upsertSearchAttributes(SearchAttributes searchAttributes) {
+    workflowClock.upsertSearchAttributes(searchAttributes);
+  }
+
+  @Override
+  public void handleUpsertSearchAttributes(HistoryEvent event) {
+    workflowClock.handleUpsertSearchAttributesEvent(event);
+  }
 }
diff --git a/src/main/java/com/uber/cadence/internal/replay/DecisionTarget.java b/src/main/java/com/uber/cadence/internal/replay/DecisionTarget.java
@@ -24,5 +24,6 @@ enum DecisionTarget {
   SIGNAL_EXTERNAL_WORKFLOW,
   TIMER,
   MARKER,
+  UPSERT_SEARCH_ATTRIBUTES,
   SELF
 }
diff --git a/src/main/java/com/uber/cadence/internal/replay/DecisionsHelper.java b/src/main/java/com/uber/cadence/internal/replay/DecisionsHelper.java
@@ -44,6 +44,7 @@
 import com.uber.cadence.RequestCancelExternalWorkflowExecutionDecisionAttributes;
 import com.uber.cadence.RequestCancelExternalWorkflowExecutionFailedEventAttributes;
 import com.uber.cadence.ScheduleActivityTaskDecisionAttributes;
+import com.uber.cadence.SearchAttributes;
 import com.uber.cadence.SignalExternalWorkflowExecutionDecisionAttributes;
 import com.uber.cadence.StartChildWorkflowExecutionDecisionAttributes;
 import com.uber.cadence.StartChildWorkflowExecutionFailedEventAttributes;
@@ -52,6 +53,7 @@
 import com.uber.cadence.TaskList;
 import com.uber.cadence.TimerCanceledEventAttributes;
 import com.uber.cadence.TimerFiredEventAttributes;
+import com.uber.cadence.UpsertWorkflowSearchAttributesDecisionAttributes;
 import com.uber.cadence.WorkflowExecutionStartedEventAttributes;
 import com.uber.cadence.WorkflowType;
 import com.uber.cadence.internal.common.WorkflowExecutionUtils;
@@ -545,6 +547,20 @@ void recordMarker(String markerName, Header header, byte[] details) {
     addDecision(decisionId, new MarkerDecisionStateMachine(decisionId, decision));
   }
 
+  void upsertSearchAttributes(SearchAttributes searchAttributes) {
+    UpsertWorkflowSearchAttributesDecisionAttributes decisionAttr =
+        new UpsertWorkflowSearchAttributesDecisionAttributes()
+            .setSearchAttributes(searchAttributes);
+    Decision decision =
+        new Decision()
+            .setDecisionType(DecisionType.UpsertWorkflowSearchAttributes)
+            .setUpsertWorkflowSearchAttributesDecisionAttributes(decisionAttr);
+    long nextDecisionEventId = getNextDecisionEventId();
+    DecisionId decisionId =
+        new DecisionId(DecisionTarget.UPSERT_SEARCH_ATTRIBUTES, nextDecisionEventId);
+    addDecision(decisionId, new UpsertSearchAttributesDecisionStateMachine(decisionId, decision));
+  }
+
   List<Decision> getDecisions() {
     List<Decision> result = new ArrayList<>(MAXIMUM_DECISIONS_PER_COMPLETION + 1);
     for (DecisionStateMachine decisionStateMachine : decisions.values()) {
diff --git a/src/main/java/com/uber/cadence/internal/replay/HistoryEventHandler.java b/src/main/java/com/uber/cadence/internal/replay/HistoryEventHandler.java
@@ -53,4 +53,6 @@ interface HistoryEventHandler {
   void handleExternalWorkflowExecutionSignaled(HistoryEvent event);
 
   void handleMarkerRecorded(HistoryEvent event);
+
+  void handleUpsertSearchAttributes(HistoryEvent event);
 }
diff --git a/src/main/java/com/uber/cadence/internal/replay/ReplayDecider.java b/src/main/java/com/uber/cadence/internal/replay/ReplayDecider.java
@@ -229,7 +229,7 @@ private void processEvent(HistoryEvent event) {
         context.handleDecisionTaskFailed(event);
         break;
       case UpsertWorkflowSearchAttributes:
-        // TODO: https://github.com/uber/cadence-java-client/issues/360
+        context.handleUpsertSearchAttributes(event);
         break;
     }
   }
diff --git a/src/main/java/com/uber/cadence/internal/replay/UpsertSearchAttributesDecisionStateMachine.java b/src/main/java/com/uber/cadence/internal/replay/UpsertSearchAttributesDecisionStateMachine.java
@@ -0,0 +1,38 @@
+/*
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package com.uber.cadence.internal.replay;
+
+import com.uber.cadence.Decision;
+
+public class UpsertSearchAttributesDecisionStateMachine extends DecisionStateMachineBase {
+
+  private final Decision decision;
+
+  UpsertSearchAttributesDecisionStateMachine(DecisionId id, Decision decision) {
+    super(id);
+    this.decision = decision;
+  }
+
+  @Override
+  public Decision getDecision() {
+    if (state == DecisionState.CREATED) {
+      return decision;
+    }
+    return null;
+  }
+}
diff --git a/src/main/java/com/uber/cadence/internal/sync/DeterministicRunnerImpl.java b/src/main/java/com/uber/cadence/internal/sync/DeterministicRunnerImpl.java
@@ -18,6 +18,7 @@
 package com.uber.cadence.internal.sync;
 
 import com.uber.cadence.ChildPolicy;
+import com.uber.cadence.SearchAttributes;
 import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.WorkflowType;
 import com.uber.cadence.converter.DataConverter;
@@ -653,5 +654,10 @@ public boolean getEnableLoggingInReplay() {
     public UUID randomUUID() {
       return UUID.randomUUID();
     }
+
+    @Override
+    public void upsertSearchAttributes(SearchAttributes searchAttributes) {
+      throw new UnsupportedOperationException("not implemented");
+    }
   }
 }
diff --git a/src/main/java/com/uber/cadence/internal/sync/SyncDecisionContext.java b/src/main/java/com/uber/cadence/internal/sync/SyncDecisionContext.java
@@ -20,12 +20,14 @@
 import static com.uber.cadence.internal.common.OptionsUtils.roundUpToSeconds;
 
 import com.uber.cadence.ActivityType;
+import com.uber.cadence.SearchAttributes;
 import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.WorkflowType;
 import com.uber.cadence.activity.ActivityOptions;
 import com.uber.cadence.activity.LocalActivityOptions;
 import com.uber.cadence.common.RetryOptions;
 import com.uber.cadence.converter.DataConverter;
+import com.uber.cadence.converter.JsonDataConverter;
 import com.uber.cadence.internal.common.RetryParameters;
 import com.uber.cadence.internal.replay.ActivityTaskFailedException;
 import com.uber.cadence.internal.replay.ActivityTaskTimeoutException;
@@ -54,6 +56,7 @@
 import com.uber.cadence.workflow.WorkflowInterceptor;
 import com.uber.m3.tally.Scope;
 import java.lang.reflect.Type;
+import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
@@ -693,4 +696,24 @@ public <R> R getLastCompletionResult(Class<R> resultClass, Type resultType) {
     DataConverter dataConverter = getDataConverter();
     return dataConverter.fromData(lastCompletionResult, resultClass, resultType);
   }
+
+  @Override
+  public void upsertSearchAttributes(Map<String, Object> searchAttributes) {
+    if (searchAttributes.isEmpty()) {
+      throw new IllegalArgumentException("Empty search attributes");
+    }
+
+    SearchAttributes attr = convertMapToSearchAttributes(searchAttributes);
+    context.upsertSearchAttributes(attr);
+  }
+
+  protected SearchAttributes convertMapToSearchAttributes(Map<String, Object> searchAttributes) {
+    DataConverter converter = JsonDataConverter.getInstance();
+    Map<String, ByteBuffer> mapOfByteBuffer = new HashMap<>();
+    searchAttributes.forEach(
+        (key, value) -> {
+          mapOfByteBuffer.put(key, ByteBuffer.wrap(converter.toData(value)));
+        });
+    return new SearchAttributes().setIndexedFields(mapOfByteBuffer);
+  }
 }
diff --git a/src/main/java/com/uber/cadence/internal/sync/TestActivityEnvironmentInternal.java b/src/main/java/com/uber/cadence/internal/sync/TestActivityEnvironmentInternal.java
@@ -104,6 +104,7 @@
 import java.lang.reflect.Type;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Random;
@@ -313,6 +314,11 @@ public UUID randomUUID() {
       throw new UnsupportedOperationException("not implemented");
     }
 
+    @Override
+    public void upsertSearchAttributes(Map<String, Object> searchAttributes) {
+      throw new UnsupportedOperationException("not implemented");
+    }
+
     private <T> T getReply(
         PollForActivityTaskResponse task,
         ActivityTaskHandler.Result response,
diff --git a/src/main/java/com/uber/cadence/internal/sync/WorkflowInternal.java b/src/main/java/com/uber/cadence/internal/sync/WorkflowInternal.java
@@ -50,6 +50,7 @@
 import java.time.Duration;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Random;
 import java.util.UUID;
@@ -384,4 +385,8 @@ public static Logger getLogger(String name) {
   public static <R> R getLastCompletionResult(Class<R> resultClass, Type resultType) {
     return getRootDecisionContext().getLastCompletionResult(resultClass, resultType);
   }
+
+  public static void upsertSearchAttributes(Map<String, Object> searchAttributes) {
+    getWorkflowInterceptor().upsertSearchAttributes(searchAttributes);
+  }
 }
diff --git a/src/main/java/com/uber/cadence/workflow/Workflow.java b/src/main/java/com/uber/cadence/workflow/Workflow.java
@@ -30,6 +30,7 @@
 import com.uber.m3.tally.Scope;
 import java.lang.reflect.Type;
 import java.time.Duration;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Random;
 import java.util.UUID;
@@ -864,7 +865,7 @@ public static boolean isReplaying() {
    * failure. The decision task after timeout is rescheduled and re-executed giving SideEffect
    * another chance to succeed.
    *
-   * <p>Caution: do not use sideEffect function to modify any workflow sate. Only use the
+   * <p>Caution: do not use sideEffect function to modify any workflow state. Only use the
    * SideEffect's return value. For example this code is BROKEN:
    *
    * <pre><code>
@@ -880,7 +881,7 @@ public static boolean isReplaying() {
    *  } else {
    *         ....
    *  }
-   *  </code></pre>
+   * </code></pre>
    *
    * On replay the provided function is not executed, the random will always be 0, and the workflow
    * could takes a different path breaking the determinism.
@@ -895,7 +896,7 @@ public static boolean isReplaying() {
    *  } else {
    *         ....
    *  }
-   *  </code></pre>
+   * </code></pre>
    *
    * If function throws any exception it is not delivered to the workflow code. It is wrapped in
    * {@link Error} causing failure of the current decision.
@@ -918,7 +919,7 @@ public static <R> R sideEffect(Class<R> resultClass, Func<R> func) {
    * failure. The decision task after timeout is rescheduled and re-executed giving SideEffect
    * another chance to succeed.
    *
-   * <p>Caution: do not use sideEffect function to modify any workflow sate. Only use the
+   * <p>Caution: do not use sideEffect function to modify any workflow state. Only use the
    * SideEffect's return value. For example this code is BROKEN:
    *
    * <pre><code>
@@ -934,7 +935,7 @@ public static <R> R sideEffect(Class<R> resultClass, Func<R> func) {
    *  } else {
    *         ....
    *  }
-   *  </code></pre>
+   * </code></pre>
    *
    * On replay the provided function is not executed, the random will always be 0, and the workflow
    * could takes a different path breaking the determinism.
@@ -949,7 +950,7 @@ public static <R> R sideEffect(Class<R> resultClass, Func<R> func) {
    *  } else {
    *         ....
    *  }
-   *  </code></pre>
+   * </code></pre>
    *
    * If function throws any exception it is not delivered to the workflow code. It is wrapped in
    * {@link Error} causing failure of the current decision.
@@ -1182,6 +1183,39 @@ public static <R> R getLastCompletionResult(Class<R> resultClass, Type resultTyp
     return WorkflowInternal.getLastCompletionResult(resultClass, resultType);
   }
 
+  /**
+   * {@code upsertSearchAttributes} is used to add or update workflow search attributes. The search
+   * attributes can be used in query of List/Scan/Count workflow APIs. The key and value type must
+   * be registered on cadence server side; The value has to be Json serializable.
+   * UpsertSearchAttributes will merge attributes to existing map in workflow, for example workflow
+   * code:
+   *
+   * <pre><code>
+   *     Map<String, Object> attr1 = new HashMap<>();
+   *     attr1.put("CustomIntField", 1);
+   *     attr1.put("CustomBoolField", true);
+   *     Workflow.upsertSearchAttributes(attr1);
+   *
+   *     Map<String, Object> attr2 = new HashMap<>();
+   *     attr2.put("CustomIntField", 2);
+   *     attr2.put("CustomKeywordField", "Seattle");
+   *     Workflow.upsertSearchAttributes(attr2);
+   * </pre></code> will eventually have search attributes as:
+   *
+   * <pre><code>
+   *     {
+   *       "CustomIntField": 2,
+   *       "CustomBoolField": true,
+   *       "CustomKeywordField": "Seattle",
+   *     }
+   * </pre></code>
+   *
+   * @param searchAttributes map of String to Object value that can be used to search in list APIs
+   */
+  public static void upsertSearchAttributes(Map<String, Object> searchAttributes) {
+    WorkflowInternal.upsertSearchAttributes(searchAttributes);
+  }
+
   /** Prohibit instantiation. */
   private Workflow() {}
 }
diff --git a/src/main/java/com/uber/cadence/workflow/WorkflowInterceptor.java b/src/main/java/com/uber/cadence/workflow/WorkflowInterceptor.java
@@ -23,6 +23,7 @@
 import com.uber.cadence.workflow.Functions.Func;
 import java.lang.reflect.Type;
 import java.time.Duration;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Random;
 import java.util.UUID;
@@ -99,4 +100,6 @@ void continueAsNew(
   void registerQuery(String queryType, Type[] argTypes, Functions.Func1<Object[], Object> callback);
 
   UUID randomUUID();
+
+  void upsertSearchAttributes(Map<String, Object> searchAttributes);
 }
diff --git a/src/main/java/com/uber/cadence/workflow/WorkflowInterceptorBase.java b/src/main/java/com/uber/cadence/workflow/WorkflowInterceptorBase.java
diff --git a/src/test/java/com/uber/cadence/internal/sync/SyncDecisionContextTest.java b/src/test/java/com/uber/cadence/internal/sync/SyncDecisionContextTest.java
diff --git a/src/test/java/com/uber/cadence/workflow/WorkflowTest.java b/src/test/java/com/uber/cadence/workflow/WorkflowTest.java
diff --git a/src/test/java/com/uber/cadence/workflow/interceptors/SignalWorkflowInterceptor.java b/src/test/java/com/uber/cadence/workflow/interceptors/SignalWorkflowInterceptor.java

-Original file line number
+Diff line change
 .gradle
 /build
 /out
 +dummy
Original file line number	Diff line number	Diff line change
`@@ -972,7 +972,8 @@ public static boolean isDecisionEvent(HistoryEvent event) {`
`972`	`972`	`\|\| eventType == EventType.CancelTimerFailed`
`973`	`973`	`\|\| eventType == EventType.RequestCancelExternalWorkflowExecutionInitiated`
`974`	`974`	`\|\| eventType == EventType.MarkerRecorded`
`975`		`- \|\| eventType == EventType.SignalExternalWorkflowExecutionInitiated));`
	`975`	`+ \|\| eventType == EventType.SignalExternalWorkflowExecutionInitiated`
	`976`	`+ \|\| eventType == EventType.UpsertWorkflowSearchAttributes));`
`976`	`977`	`return result;`
`977`	`978`	`}`
`978`	`979`
Original file line number	Diff line number	Diff line change
`@@ -24,5 +24,6 @@ enum DecisionTarget {`
`24`	`24`	`SIGNAL_EXTERNAL_WORKFLOW,`
`25`	`25`	`TIMER,`
`26`	`26`	`MARKER,`
	`27`	`+ UPSERT_SEARCH_ATTRIBUTES,`
`27`	`28`	`SELF`
`28`	`29`	`}`
Original file line number	Diff line number	Diff line change
`@@ -53,4 +53,6 @@ interface HistoryEventHandler {`
`53`	`53`	`void handleExternalWorkflowExecutionSignaled(HistoryEvent event);`
`54`	`54`
`55`	`55`	`void handleMarkerRecorded(HistoryEvent event);`
	`56`	`+`
	`57`	`+ void handleUpsertSearchAttributes(HistoryEvent event);`
`56`	`58`	`}`
Original file line number	Diff line number	Diff line change
`@@ -229,7 +229,7 @@ private void processEvent(HistoryEvent event) {`
`229`	`229`	`context.handleDecisionTaskFailed(event);`
`230`	`230`	`break;`
`231`	`231`	`case UpsertWorkflowSearchAttributes:`
`232`		`- // TODO: https://github.com/uber/cadence-java-client/issues/360`
	`232`	`+ context.handleUpsertSearchAttributes(event);`
`233`	`233`	`break;`
`234`	`234`	`}`
`235`	`235`	`}`