Add num_pollers metric (#2514)

Quinn-With-Two-Ns · web-flow · commit b7c72a28f56d · 2025-05-13T07:49:41.000-07:00
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/ActivityPollTask.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/ActivityPollTask.java
@@ -11,10 +11,13 @@
 import io.temporal.api.workflowservice.v1.PollActivityTaskQueueRequest;
 import io.temporal.api.workflowservice.v1.PollActivityTaskQueueResponse;
 import io.temporal.internal.common.ProtobufTimeUtils;
+import io.temporal.serviceclient.MetricsTag;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.worker.MetricsType;
+import io.temporal.worker.PollerTypeMetricsTag;
 import io.temporal.worker.tuning.*;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -28,6 +31,7 @@ final class ActivityPollTask implements Poller.PollTask<ActivityTask> {
   private final TrackingSlotSupplier<ActivitySlotInfo> slotSupplier;
   private final Scope metricsScope;
   private final PollActivityTaskQueueRequest pollRequest;
+  private final AtomicInteger pollGauge = new AtomicInteger();
 
   @SuppressWarnings("deprecation")
   public ActivityPollTask(
@@ -91,6 +95,10 @@ public ActivityTask poll() {
     permit = Poller.getSlotPermitAndHandleInterrupts(future, slotSupplier);
     if (permit == null) return null;
 
+    MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.ACTIVITY_TASK)
+        .gauge(MetricsType.NUM_POLLERS)
+        .update(pollGauge.incrementAndGet());
+
     try {
       response =
           service
@@ -113,6 +121,10 @@ public ActivityTask poll() {
           permit,
           () -> slotSupplier.releaseSlot(SlotReleaseReason.taskComplete(), permit));
     } finally {
+      MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.ACTIVITY_TASK)
+          .gauge(MetricsType.NUM_POLLERS)
+          .update(pollGauge.decrementAndGet());
+
       if (!isSuccessful) slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), permit);
     }
   }
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusPollTask.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusPollTask.java
@@ -8,10 +8,13 @@
 import io.temporal.api.taskqueue.v1.TaskQueue;
 import io.temporal.api.workflowservice.v1.*;
 import io.temporal.internal.common.ProtobufTimeUtils;
+import io.temporal.serviceclient.MetricsTag;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.worker.MetricsType;
+import io.temporal.worker.PollerTypeMetricsTag;
 import io.temporal.worker.tuning.*;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -25,6 +28,7 @@ final class NexusPollTask implements Poller.PollTask<NexusTask> {
   private final TrackingSlotSupplier<NexusSlotInfo> slotSupplier;
   private final Scope metricsScope;
   private final PollNexusTaskQueueRequest pollRequest;
+  private final AtomicInteger pollGauge = new AtomicInteger();
 
   @SuppressWarnings("deprecation")
   public NexusPollTask(
@@ -81,6 +85,10 @@ public NexusTask poll() {
     permit = Poller.getSlotPermitAndHandleInterrupts(future, slotSupplier);
     if (permit == null) return null;
 
+    MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.NEXUS_TASK)
+        .gauge(MetricsType.NUM_POLLERS)
+        .update(pollGauge.incrementAndGet());
+
     try {
       response =
           service
@@ -106,6 +114,10 @@ public NexusTask poll() {
           permit,
           () -> slotSupplier.releaseSlot(SlotReleaseReason.taskComplete(), permit));
     } finally {
+      MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.NEXUS_TASK)
+          .gauge(MetricsType.NUM_POLLERS)
+          .update(pollGauge.decrementAndGet());
+
       if (!isSuccessful) slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), permit);
     }
   }
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java
@@ -15,11 +15,13 @@
 import io.temporal.serviceclient.MetricsTag;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.worker.MetricsType;
+import io.temporal.worker.PollerTypeMetricsTag;
 import io.temporal.worker.tuning.SlotPermit;
 import io.temporal.worker.tuning.SlotReleaseReason;
 import io.temporal.worker.tuning.SlotSupplierFuture;
 import io.temporal.worker.tuning.WorkflowSlotInfo;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -36,6 +38,8 @@ final class WorkflowPollTask implements Poller.PollTask<WorkflowTask> {
   private final WorkflowServiceGrpc.WorkflowServiceBlockingStub serviceStub;
   private final PollWorkflowTaskQueueRequest pollRequest;
   private final PollWorkflowTaskQueueRequest stickyPollRequest;
+  private final AtomicInteger normalPollGauge = new AtomicInteger();
+  private final AtomicInteger stickyPollGauge = new AtomicInteger();
 
   @SuppressWarnings("deprecation")
   public WorkflowPollTask(
@@ -131,6 +135,16 @@ public WorkflowTask poll() {
     Scope scope = isSticky ? stickyMetricsScope : metricsScope;
 
     log.trace("poll request begin: {}", request);
+    if (isSticky) {
+      MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.WORKFLOW_STICKY_TASK)
+          .gauge(MetricsType.NUM_POLLERS)
+          .update(stickyPollGauge.incrementAndGet());
+    } else {
+      MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.WORKFLOW_TASK)
+          .gauge(MetricsType.NUM_POLLERS)
+          .update(normalPollGauge.incrementAndGet());
+    }
+
     try {
       PollWorkflowTaskQueueResponse response = doPoll(request, scope);
       if (response == null) {
@@ -141,6 +155,17 @@ public WorkflowTask poll() {
       slotSupplier.markSlotUsed(new WorkflowSlotInfo(response, pollRequest), permit);
       return new WorkflowTask(response, (rr) -> slotSupplier.releaseSlot(rr, permit));
     } finally {
+
+      if (isSticky) {
+        MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.WORKFLOW_STICKY_TASK)
+            .gauge(MetricsType.NUM_POLLERS)
+            .update(stickyPollGauge.decrementAndGet());
+      } else {
+        MetricsTag.tagged(metricsScope, PollerTypeMetricsTag.PollerType.WORKFLOW_TASK)
+            .gauge(MetricsType.NUM_POLLERS)
+            .update(normalPollGauge.decrementAndGet());
+      }
+
       if (!isSuccessful) {
         slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), permit);
         stickyQueueBalancer.finishPoll(taskQueueKind, 0);
diff --git a/temporal-sdk/src/main/java/io/temporal/worker/MetricsType.java b/temporal-sdk/src/main/java/io/temporal/worker/MetricsType.java
@@ -120,6 +120,8 @@ private MetricsType() {}
   public static final String WORKER_START_COUNTER = TEMPORAL_METRICS_PREFIX + "worker_start";
   public static final String POLLER_START_COUNTER = TEMPORAL_METRICS_PREFIX + "poller_start";
   // gauge
+  public static final String NUM_POLLERS = TEMPORAL_METRICS_PREFIX + "num_pollers";
+
   public static final String WORKER_TASK_SLOTS_AVAILABLE =
       TEMPORAL_METRICS_PREFIX + "worker_task_slots_available";
 
diff --git a/temporal-sdk/src/main/java/io/temporal/worker/PollerTypeMetricsTag.java b/temporal-sdk/src/main/java/io/temporal/worker/PollerTypeMetricsTag.java
@@ -0,0 +1,28 @@
+package io.temporal.worker;
+
+import io.temporal.serviceclient.MetricsTag;
+
+public class PollerTypeMetricsTag {
+  public enum PollerType implements MetricsTag.TagValue {
+    WORKFLOW_TASK("workflow_task"),
+    WORKFLOW_STICKY_TASK("workflow_sticky_task"),
+    ACTIVITY_TASK("activity_task"),
+    NEXUS_TASK("nexus_task"),
+    ;
+
+    PollerType(String value) {
+      this.value = value;
+    }
+
+    private final String value;
+
+    @Override
+    public String getTag() {
+      return MetricsTag.POLLER_TYPE;
+    }
+
+    public String getValue() {
+      return value;
+    }
+  }
+}
diff --git a/temporal-sdk/src/test/java/io/temporal/workflow/MetricsTest.java b/temporal-sdk/src/test/java/io/temporal/workflow/MetricsTest.java
@@ -31,6 +31,7 @@
 import io.temporal.testing.TestEnvironmentOptions;
 import io.temporal.testing.TestWorkflowEnvironment;
 import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.worker.PollerTypeMetricsTag;
 import io.temporal.worker.Worker;
 import io.temporal.worker.WorkerFactoryOptions;
 import io.temporal.worker.WorkerMetricsTag;
@@ -89,6 +90,26 @@ public class MetricsTest {
           .put(MetricsTag.WORKER_TYPE, WorkerMetricsTag.WorkerType.WORKFLOW_WORKER.getValue())
           .build();
 
+  private static final Map<String, String> TAGS_WORKFLOW_NORMAL_POLLER =
+      new ImmutableMap.Builder<String, String>()
+          .putAll(TAGS_WORKFLOW_WORKER)
+          .put(MetricsTag.POLLER_TYPE, PollerTypeMetricsTag.PollerType.WORKFLOW_TASK.getValue())
+          .build();
+
+  private static final Map<String, String> TAGS_WORKFLOW_STICKY_POLLER =
+      new ImmutableMap.Builder<String, String>()
+          .putAll(TAGS_WORKFLOW_WORKER)
+          .put(
+              MetricsTag.POLLER_TYPE,
+              PollerTypeMetricsTag.PollerType.WORKFLOW_STICKY_TASK.getValue())
+          .build();
+
+  private static final Map<String, String> TAGS_ACTIVITY_POLLER =
+      new ImmutableMap.Builder<String, String>()
+          .putAll(TAGS_ACTIVITY_WORKER)
+          .put(MetricsTag.POLLER_TYPE, PollerTypeMetricsTag.PollerType.ACTIVITY_TASK.getValue())
+          .build();
+
   @Rule
   public TestWatcher watchman =
       new TestWatcher() {
@@ -150,7 +171,10 @@ public void testWorkerMetrics() throws InterruptedException {
     reporter.assertCounter("temporal_worker_start", TAGS_LOCAL_ACTIVITY_WORKER, 1);
 
     reporter.assertCounter("temporal_poller_start", TAGS_WORKFLOW_WORKER, 5);
+    reporter.assertGauge("temporal_num_pollers", TAGS_WORKFLOW_NORMAL_POLLER, 2);
+    reporter.assertGauge("temporal_num_pollers", TAGS_WORKFLOW_STICKY_POLLER, 3);
     reporter.assertCounter("temporal_poller_start", TAGS_ACTIVITY_WORKER, 5);
+    reporter.assertGauge("temporal_num_pollers", TAGS_ACTIVITY_POLLER, 5);
   }
 
   @Test
diff --git a/temporal-serviceclient/src/main/java/io/temporal/serviceclient/MetricsTag.java b/temporal-serviceclient/src/main/java/io/temporal/serviceclient/MetricsTag.java
@@ -24,6 +24,7 @@ public class MetricsTag {
   public static final String EXCEPTION = "exception";
   public static final String OPERATION_NAME = "operation";
   public static final String TASK_FAILURE_TYPE = "failure_reason";
+  public static final String POLLER_TYPE = "poller_type";
 
   /** Used to pass metrics scope to the interceptor */
   public static final CallOptions.Key<Scope> METRICS_TAGS_CALL_OPTIONS_KEY =
