InboundActivityInterceptor now has an access to header (temporalio#375)

Author: Spikhalskiy

temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityInboundCallsInterceptor.java

@@ -22,7 +22,42 @@
 import io.temporal.activity.ActivityExecutionContext;
 
 public interface ActivityInboundCallsInterceptor {
+  final class ActivityInput {
+    private final Header header;
+    private final Object[] arguments;
+
+    public ActivityInput(Header header, Object[] arguments) {
+      this.header = header;
+      this.arguments = arguments;
+    }
+
+    public Header getHeader() {
+      return header;
+    }
+
+    public Object[] getArguments() {
+      return arguments;
+    }
+  }
+
+  final class ActivityOutput {
+    private final Object result;
+
+    public ActivityOutput(Object result) {
+      this.result = result;
+    }
+
+    public Object getResult() {
+      return result;
+    }
+  }
+
   void init(ActivityExecutionContext context);
 
-  Object execute(Object[] arguments);
+  /**
+   * Intercepts a call to the main activity entry method.
+   *
+   * @return result of the activity execution.
+   */
+  ActivityOutput execute(ActivityInput input);
 }

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

@@ -20,7 +20,7 @@
 package io.temporal.internal.sync;
 
 import com.google.protobuf.util.Timestamps;
-import io.temporal.activity.ActivityInfo;
+import io.temporal.api.common.v1.Header;
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.workflowservice.v1.PollActivityTaskQueueResponse;
 import io.temporal.internal.common.ProtobufTimeUtils;
@@ -29,7 +29,7 @@
 import java.util.Objects;
 import java.util.Optional;
 
-final class ActivityInfoImpl implements ActivityInfo {
+final class ActivityInfoImpl implements ActivityInfoInternal {
   private final PollActivityTaskQueueResponse response;
   private final String activityNamespace;
   private final boolean local;
@@ -46,6 +46,7 @@ final class ActivityInfoImpl implements ActivityInfo {
     this.completionHandle = completionHandle;
   }
 
+  @Override
   public byte[] getTaskToken() {
     return response.getTaskToken().toByteArray();
   }
@@ -124,14 +125,21 @@ public boolean isLocal() {
     return local;
   }
 
+  @Override
   public Functions.Proc getCompletionHandle() {
     return completionHandle;
   }
 
+  @Override
   public Optional<Payloads> getInput() {
     if (response.hasInput()) {
       return Optional.of(response.getInput());
     }
     return Optional.empty();
   }
+
+  @Override
+  public Header getHeader() {
+    return response.getHeader();
+  }
 }

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

@@ -0,0 +1,45 @@
+/*
+ *  Copyright (C) 2020 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ *  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 io.temporal.internal.sync;
+
+import io.temporal.activity.ActivityInfo;
+import io.temporal.api.common.v1.Header;
+import io.temporal.api.common.v1.Payloads;
+import io.temporal.workflow.Functions;
+import java.util.Optional;
+
+/**
+ * An extension for {@link ActivityInfo} with information about the activity task that the current
+ * activity is handling that should be available for Temporal SDK, but shouldn't be available or
+ * exposed to Activity implementation code.
+ */
+public interface ActivityInfoInternal extends ActivityInfo {
+  /**
+   * @return function shat should be triggered after activity completion with any outcome (success,
+   *     failure, cancelling)
+   */
+  Functions.Proc getCompletionHandle();
+
+  /** @return input parameters of the activity execution */
+  Optional<Payloads> getInput();
+
+  /** @return header that is passed with the activity execution */
+  Header getHeader();
+}

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

@@ -36,6 +36,7 @@
 import io.temporal.common.converter.DataConverter;
 import io.temporal.common.converter.EncodedValues;
 import io.temporal.common.interceptors.ActivityInboundCallsInterceptor;
+import io.temporal.common.interceptors.Header;
 import io.temporal.common.interceptors.WorkerInterceptor;
 import io.temporal.common.metadata.POJOActivityImplMetadata;
 import io.temporal.common.metadata.POJOActivityInterfaceMetadata;
@@ -189,7 +190,7 @@ void registerLocalActivityImplementations(Object[] activitiesImplementation) {
   public Result handle(ActivityTask activityTask, Scope metricsScope, boolean localActivity) {
     PollActivityTaskQueueResponse pollResponse = activityTask.getResponse();
     String activityType = pollResponse.getActivityType().getName();
-    ActivityInfoImpl activityInfo =
+    ActivityInfoInternal activityInfo =
         new ActivityInfoImpl(
             pollResponse, this.namespace, localActivity, activityTask.getCompletionHandle());
     ActivityTaskExecutor activity = activities.get(activityType);
@@ -212,7 +213,7 @@ public Result handle(ActivityTask activityTask, Scope metricsScope, boolean loca
   }
 
   private interface ActivityTaskExecutor {
-    ActivityTaskHandler.Result execute(ActivityInfoImpl task, Scope metricsScope);
+    ActivityTaskHandler.Result execute(ActivityInfoInternal task, Scope metricsScope);
   }
 
   private class POJOActivityImplementation implements ActivityTaskExecutor {
@@ -225,7 +226,7 @@ private class POJOActivityImplementation implements ActivityTaskExecutor {
     }
 
     @Override
-    public ActivityTaskHandler.Result execute(ActivityInfoImpl info, Scope metricsScope) {
+    public ActivityTaskHandler.Result execute(ActivityInfoInternal info, Scope metricsScope) {
       ActivityExecutionContext context =
           new ActivityExecutionContextImpl(
               service,
@@ -249,15 +250,18 @@ public ActivityTaskHandler.Result execute(ActivityInfoImpl info, Scope metricsSc
                 input,
                 method.getParameterTypes(),
                 method.getGenericParameterTypes());
-        Object result = inboundCallsInterceptor.execute(args);
+        ActivityInboundCallsInterceptor.ActivityOutput result =
+            inboundCallsInterceptor.execute(
+                new ActivityInboundCallsInterceptor.ActivityInput(
+                    new Header(info.getHeader()), args));
         if (context.isDoNotCompleteOnReturn()) {
           return new ActivityTaskHandler.Result(
               info.getActivityId(), null, null, null, null, context.isUseLocalManualCompletion());
         }
         RespondActivityTaskCompletedRequest.Builder request =
             RespondActivityTaskCompletedRequest.newBuilder();
         if (method.getReturnType() != Void.TYPE) {
-          Optional<Payloads> serialized = dataConverter.toPayloads(result);
+          Optional<Payloads> serialized = dataConverter.toPayloads(result.getResult());
           if (serialized.isPresent()) {
             request.setResult(serialized.get());
           }
@@ -287,10 +291,11 @@ public void init(ActivityExecutionContext context) {
     }
 
     @Override
-    public Object execute(Object[] arguments) {
+    public ActivityOutput execute(ActivityInput input) {
       CurrentActivityExecutionContext.set(context);
       try {
-        return method.invoke(activity, arguments);
+        Object result = method.invoke(activity, input.getArguments());
+        return new ActivityOutput(result);
       } catch (InvocationTargetException e) {
         throw Activity.wrap(e.getTargetException());
       } catch (Exception e) {
@@ -310,7 +315,7 @@ public DynamicActivityImplementation(DynamicActivity activity) {
     }
 
     @Override
-    public ActivityTaskHandler.Result execute(ActivityInfoImpl info, Scope metricsScope) {
+    public ActivityTaskHandler.Result execute(ActivityInfoInternal info, Scope metricsScope) {
       ActivityExecutionContext context =
           new ActivityExecutionContextImpl(
               service,
@@ -328,15 +333,20 @@ public ActivityTaskHandler.Result execute(ActivityInfoImpl info, Scope metricsSc
       }
       inboundCallsInterceptor.init(context);
       try {
-        EncodedValues args = new EncodedValues(input, dataConverter);
-        Object result = inboundCallsInterceptor.execute(new Object[] {args});
+        EncodedValues encodedValues = new EncodedValues(input, dataConverter);
+        Object[] args = new Object[] {encodedValues};
+
+        ActivityInboundCallsInterceptor.ActivityOutput result =
+            inboundCallsInterceptor.execute(
+                new ActivityInboundCallsInterceptor.ActivityInput(
+                    new Header(info.getHeader()), args));
         if (context.isDoNotCompleteOnReturn()) {
           return new ActivityTaskHandler.Result(
               info.getActivityId(), null, null, null, null, context.isUseLocalManualCompletion());
         }
         RespondActivityTaskCompletedRequest.Builder request =
             RespondActivityTaskCompletedRequest.newBuilder();
-        Optional<Payloads> serialized = dataConverter.toPayloads(result);
+        Optional<Payloads> serialized = dataConverter.toPayloads(result.getResult());
         if (serialized.isPresent()) {
           request.setResult(serialized.get());
         }
@@ -348,7 +358,8 @@ public ActivityTaskHandler.Result execute(ActivityInfoImpl info, Scope metricsSc
     }
   }
 
-  private Result activityFailureToResult(ActivityInfoImpl info, Scope metricsScope, Throwable e) {
+  private Result activityFailureToResult(
+      ActivityInfoInternal info, Scope metricsScope, Throwable e) {
     e = CheckedExceptionWrapper.unwrap(e);
     if (e instanceof ActivityCanceledException) {
       if (log.isInfoEnabled()) {
@@ -388,10 +399,11 @@ public void init(ActivityExecutionContext context) {
     }
 
     @Override
-    public Object execute(Object[] arguments) {
+    public ActivityOutput execute(ActivityInput input) {
       CurrentActivityExecutionContext.set(context);
       try {
-        return activity.execute((EncodedValues) arguments[0]);
+        Object result = activity.execute((EncodedValues) input.getArguments()[0]);
+        return new ActivityOutput(result);
       } catch (Exception e) {
         throw Activity.wrap(e);
       } finally {
@@ -410,7 +422,7 @@ private class POJOLocalActivityImplementation implements ActivityTaskExecutor {
     }
 
     @Override
-    public ActivityTaskHandler.Result execute(ActivityInfoImpl info, Scope metricsScope) {
+    public ActivityTaskHandler.Result execute(ActivityInfoInternal info, Scope metricsScope) {
       ActivityExecutionContext context = new LocalActivityExecutionContextImpl(info, metricsScope);
       Optional<Payloads> input = info.getInput();
       ActivityInboundCallsInterceptor inboundCallsInterceptor =
@@ -426,11 +438,14 @@ public ActivityTaskHandler.Result execute(ActivityInfoImpl info, Scope metricsSc
                 input,
                 method.getParameterTypes(),
                 method.getGenericParameterTypes());
-        Object result = inboundCallsInterceptor.execute(args);
+        ActivityInboundCallsInterceptor.ActivityOutput result =
+            inboundCallsInterceptor.execute(
+                new ActivityInboundCallsInterceptor.ActivityInput(
+                    new Header(info.getHeader()), args));
         RespondActivityTaskCompletedRequest.Builder request =
             RespondActivityTaskCompletedRequest.newBuilder();
         if (method.getReturnType() != Void.TYPE) {
-          Optional<Payloads> serialized = dataConverter.toPayloads(result);
+          Optional<Payloads> serialized = dataConverter.toPayloads(result.getResult());
           if (serialized.isPresent()) {
             request.setResult(serialized.get());
           }

temporal-testing-junit4/src/main/java/io/temporal/testing/TracingWorkerInterceptor.java

@@ -382,9 +382,9 @@ public <V> void heartbeat(V details) throws ActivityCompletionException {
     }
 
     @Override
-    public Object execute(Object[] arguments) {
+    public ActivityOutput execute(ActivityInput input) {
       trace.add((local ? "local " : "") + "activity " + type);
-      return next.execute(arguments);
+      return next.execute(input);
     }
   }
 }