Add getMetricsScope interceptor (#2224)

Add getMetricsScope interceptor

Author: Quinn-With-Two-Ns

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

@@ -20,6 +20,7 @@
 
 package io.temporal.common.interceptors;
 
+import com.uber.m3.tally.Scope;
 import io.temporal.activity.ActivityOptions;
 import io.temporal.activity.LocalActivityOptions;
 import io.temporal.api.common.v1.WorkflowExecution;
@@ -621,6 +622,9 @@ <R> R mutableSideEffect(
 
   void upsertMemo(Map<String, Object> memo);
 
+  /** Intercepts call to get the metric scope in a workflow. */
+  Scope getMetricsScope();
+
   /**
    * Intercepts creation of a workflow child thread.
    *

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

@@ -20,6 +20,7 @@
 
 package io.temporal.common.interceptors;
 
+import com.uber.m3.tally.Scope;
 import io.temporal.common.SearchAttributeUpdate;
 import io.temporal.workflow.Functions.Func;
 import io.temporal.workflow.Promise;
@@ -167,6 +168,11 @@ public void upsertMemo(Map<String, Object> memo) {
     next.upsertMemo(memo);
   }
 
+  @Override
+  public Scope getMetricsScope() {
+    return next.getMetricsScope();
+  }
+
   @Override
   public Object newChildThread(Runnable runnable, boolean detached, String name) {
     return next.newChildThread(runnable, detached, name);

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

@@ -1210,6 +1210,7 @@ public CancelWorkflowOutput cancelWorkflow(CancelWorkflowInput input) {
     return new CancelWorkflowOutput(result);
   }
 
+  @Override
   public Scope getMetricsScope() {
     return replayContext.getMetricsScope();
   }

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

@@ -661,7 +661,7 @@ public static Optional<UpdateInfo> getCurrentUpdateInfo() {
   }
 
   public static Scope getMetricsScope() {
-    return getRootWorkflowContext().getMetricsScope();
+    return getWorkflowOutboundInterceptor().getMetricsScope();
   }
 
   private static boolean isLoggingEnabledInReplay() {

temporal-sdk/src/test/java/io/temporal/workflow/MetricsTest.java

@@ -61,6 +61,7 @@
 import io.temporal.workflow.shared.TestWorkflows.ReceiveSignalObjectWorkflow;
 import io.temporal.workflow.shared.TestWorkflows.TestWorkflowReturnString;
 import java.time.Duration;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
@@ -77,6 +78,7 @@ public class MetricsTest {
 
   private static final long REPORTING_FLUSH_TIME = 600;
   private static final String TASK_QUEUE = "metrics_test";
+  private static final String TEST_TAG = "test_tag";
   private TestWorkflowEnvironment testEnvironment;
   private TestStatsReporter reporter;
 
@@ -243,6 +245,44 @@ public void testWorkflowMetrics() throws InterruptedException {
     reporter.assertCounter(TEMPORAL_REQUEST, workflowTaskCompletionTags, 4);
   }
 
+  @Test
+  public void testWorkflowMetricsInterceptor() throws InterruptedException {
+    setUp(
+        WorkerFactoryOptions.getDefaultInstance().toBuilder()
+            .setWorkerInterceptors(new WorkerInterceptor())
+            .build());
+
+    Worker worker = testEnvironment.newWorker(TASK_QUEUE);
+    worker.registerWorkflowImplementationTypes(
+        TestCustomMetricsInWorkflow.class, TestMetricsInChildWorkflow.class);
+    worker.registerActivitiesImplementations(new TestActivityImpl());
+    testEnvironment.start();
+
+    WorkflowClient workflowClient = testEnvironment.getWorkflowClient();
+    WorkflowOptions options =
+        WorkflowOptions.newBuilder()
+            .setWorkflowRunTimeout(Duration.ofSeconds(1000))
+            .setTaskQueue(TASK_QUEUE)
+            .build();
+    NoArgsWorkflow workflow = workflowClient.newWorkflowStub(NoArgsWorkflow.class, options);
+    workflow.execute();
+
+    Thread.sleep(REPORTING_FLUSH_TIME);
+
+    Map<String, String> workflowTags = new LinkedHashMap<>(TAGS_TASK_QUEUE);
+    // Assert the interceptor added the extra tag
+    workflowTags.put(TEST_TAG, NAMESPACE);
+
+    workflowTags.put(MetricsTag.WORKFLOW_TYPE, "NoArgsWorkflow");
+    reporter.assertCounter("test_started", workflowTags, 1);
+    reporter.assertCounter("test_done", workflowTags, 1);
+
+    workflowTags.put(MetricsTag.WORKFLOW_TYPE, "TestChildWorkflow");
+    reporter.assertCounter("test_child_started", workflowTags, 1);
+    reporter.assertCounter("test_child_done", workflowTags, 1);
+    reporter.assertTimerMinDuration("test_timer", workflowTags, Duration.ofSeconds(3));
+  }
+
   @Test
   public void testCorruptedSignalMetrics() throws InterruptedException {
     setUp(
@@ -600,4 +640,33 @@ public SignalExternalOutput signalExternalWorkflow(SignalExternalInput input) {
               overrideArgs.apply(args)));
     }
   }
+
+  private static class WorkerInterceptor extends WorkerInterceptorBase {
+    @Override
+    public WorkflowInboundCallsInterceptor interceptWorkflow(WorkflowInboundCallsInterceptor next) {
+      return new WorkflowInboundCallsInterceptorBase(next) {
+        @Override
+        public void init(WorkflowOutboundCallsInterceptor outboundCalls) {
+          next.init(new OutboundCallsInterceptor(outboundCalls));
+        }
+      };
+    }
+  }
+
+  private static class OutboundCallsInterceptor extends WorkflowOutboundCallsInterceptorBase {
+    WorkflowOutboundCallsInterceptor next;
+
+    public OutboundCallsInterceptor(WorkflowOutboundCallsInterceptor next) {
+      super(next);
+      this.next = next;
+    }
+
+    @Override
+    public Scope getMetricsScope() {
+      return next.getMetricsScope()
+          .tagged(
+              Collections.singletonMap(
+                  TEST_TAG, String.valueOf(Workflow.getInfo().getNamespace())));
+    }
+  }
 }

temporal-testing/src/main/java/io/temporal/testing/TestActivityEnvironmentInternal.java

@@ -481,6 +481,11 @@ public void upsertMemo(Map<String, Object> memo) {
       throw new UnsupportedOperationException("not implemented");
     }
 
+    @Override
+    public Scope getMetricsScope() {
+      throw new UnsupportedOperationException("not implemented");
+    }
+
     @Override
     public Object newChildThread(Runnable runnable, boolean detached, String name) {
       throw new UnsupportedOperationException("not implemented");

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

@@ -23,6 +23,7 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
+import com.uber.m3.tally.Scope;
 import io.temporal.activity.ActivityExecutionContext;
 import io.temporal.client.ActivityCompletionException;
 import io.temporal.common.SearchAttributeUpdate;
@@ -398,6 +399,14 @@ public void upsertMemo(Map<String, Object> memo) {
       next.upsertMemo(memo);
     }
 
+    @Override
+    public Scope getMetricsScope() {
+      if (!WorkflowUnsafe.isReplaying()) {
+        trace.add("getMetricsScope");
+      }
+      return next.getMetricsScope();
+    }
+
     @Override
     public Object newChildThread(Runnable runnable, boolean detached, String name) {
       if (!WorkflowUnsafe.isReplaying()) {