Set activity execution context while running local activities (#319)

Author: meiliang86

src/main/java/com/uber/cadence/internal/sync/CurrentActivityExecutionContext.java

@@ -25,8 +25,7 @@
  */
 final class CurrentActivityExecutionContext {
 
-  private static final ThreadLocal<ActivityExecutionContext> CURRENT =
-      new ThreadLocal<ActivityExecutionContext>();
+  private static final ThreadLocal<ActivityExecutionContext> CURRENT = new ThreadLocal<>();
 
   /**
    * This is used by activity implementation to get access to the current ActivityExecutionContext
@@ -39,7 +38,7 @@ public static ActivityExecutionContext get() {
               + "implementation methods and in the same thread that invoked an activity.");
     }
     return result;
-  };
+  }
 
   public static boolean isSet() {
     return CURRENT.get() != null;

src/main/java/com/uber/cadence/internal/sync/LocalActivityExecutionContextImpl.java

@@ -0,0 +1,86 @@
+/*
+ *  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.sync;
+
+import com.uber.cadence.WorkflowExecution;
+import com.uber.cadence.activity.ActivityTask;
+import com.uber.cadence.client.ActivityCompletionException;
+import com.uber.cadence.serviceclient.IWorkflowService;
+import java.lang.reflect.Type;
+import java.util.Optional;
+
+class LocalActivityExecutionContextImpl implements ActivityExecutionContext {
+  private final IWorkflowService service;
+  private final String domain;
+  private final ActivityTask task;
+
+  LocalActivityExecutionContextImpl(IWorkflowService service, String domain, ActivityTask task) {
+    this.domain = domain;
+    this.service = service;
+    this.task = task;
+  }
+
+  @Override
+  public byte[] getTaskToken() {
+    throw new UnsupportedOperationException("getTaskToken is not supported for local activities");
+  }
+
+  @Override
+  public WorkflowExecution getWorkflowExecution() {
+    return task.getWorkflowExecution();
+  }
+
+  @Override
+  public ActivityTask getTask() {
+    return task;
+  }
+
+  @Override
+  public <V> void recordActivityHeartbeat(V details) throws ActivityCompletionException {
+    throw new UnsupportedOperationException(
+        "recordActivityHeartbeat is not supported for local activities");
+  }
+
+  @Override
+  public <V> Optional<V> getHeartbeatDetails(Class<V> detailsClass, Type detailsType) {
+    throw new UnsupportedOperationException(
+        "getHeartbeatDetails is not supported for local activities");
+  }
+
+  @Override
+  public void doNotCompleteOnReturn() {
+    throw new UnsupportedOperationException(
+        "doNotCompleteOnReturn is not supported for local activities");
+  }
+
+  @Override
+  public boolean isDoNotCompleteOnReturn() {
+    throw new UnsupportedOperationException(
+        "isDoNotCompleteOnReturn is not supported for local activities");
+  }
+
+  @Override
+  public IWorkflowService getService() {
+    return service;
+  }
+
+  @Override
+  public String getDomain() {
+    return domain;
+  }
+}

src/main/java/com/uber/cadence/internal/sync/POJOActivityTaskHandler.java

@@ -235,7 +235,9 @@ private class POJOLocalActivityImplementation implements ActivityTaskExecutor {
 
     @Override
     public ActivityTaskHandler.Result execute(ActivityTaskImpl task, Scope metricsScope) {
-
+      ActivityExecutionContext context =
+          new LocalActivityExecutionContextImpl(service, domain, task);
+      CurrentActivityExecutionContext.set(context);
       byte[] input = task.getInput();
       Object[] args = dataConverter.fromDataArray(input, method.getGenericParameterTypes());
       try {
@@ -249,6 +251,8 @@ public ActivityTaskHandler.Result execute(ActivityTaskImpl task, Scope metricsSc
         return mapToActivityFailure(task.getActivityType(), e, metricsScope);
       } catch (InvocationTargetException e) {
         return mapToActivityFailure(task.getActivityType(), e.getTargetException(), metricsScope);
+      } finally {
+        CurrentActivityExecutionContext.unset();
       }
     }
   }

src/test/java/com/uber/cadence/workflow/WorkflowTest.java

@@ -542,6 +542,7 @@ public void testLocalActivityRetry() {
       assertTrue(e.getCause().getCause() instanceof IOException);
     }
     assertEquals(activitiesImpl.toString(), 5, activitiesImpl.invocations.size());
+    assertEquals("last attempt", 5, activitiesImpl.getLastAttempt());
   }
 
   public static class TestActivityRetryOnTimeout implements TestWorkflow1 {
@@ -3292,6 +3293,7 @@ private static class TestActivitiesImpl implements TestActivities {
     final AtomicInteger heartbeatCounter = new AtomicInteger();
     private final ThreadPoolExecutor executor =
         new ThreadPoolExecutor(0, 100, 1, TimeUnit.SECONDS, new LinkedBlockingQueue<>());
+    int lastAttempt;
 
     private TestActivitiesImpl(ActivityCompletionClient completionClient) {
       this.completionClient = completionClient;
@@ -3459,6 +3461,7 @@ public void heartbeatAndThrowIO() {
 
     @Override
     public void throwIO() {
+      lastAttempt = Activity.getTask().getAttempt();
       invocations.add("throwIO");
       try {
         throw new IOException("simulated IO problem");
@@ -3487,6 +3490,10 @@ public void throwIOAnnotated() {
     public List<UUID> activityUUIDList(List<UUID> arg) {
       return arg;
     }
+
+    public int getLastAttempt() {
+      return lastAttempt;
+    }
   }
 
   public interface ProcInvocationQueryable {