Emit supervisorId dimension with streaming task metrics (#18803)

kfaraz · web-flow · commit 71fbba6812c4 · 2025-12-03T19:10:07.000+05:30
Changes:
- Add dimension `supervisorId` to all metrics emitted by streaming tasks
- Simplify constructor of `TaskRealtimeMetricsMonitor`
- Remove `TaskRealtimeMetricsMonitorBuilder`
diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/KafkaClusterMetricsTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/KafkaClusterMetricsTest.java
@@ -51,7 +51,6 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -154,7 +153,8 @@ public void test_ingest10kRows_ofSelfClusterMetrics_andVerifyValues()
     // Wait for segments to be handed off
     indexer.latchableEmitter().waitForEventAggregate(
         event -> event.hasMetricName("ingest/handoff/count")
-                      .hasDimension(DruidMetrics.DATASOURCE, List.of(dataSource)),
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource)
+                      .hasDimension(DruidMetrics.SUPERVISOR_ID, supervisorId),
         agg -> agg.hasSumAtLeast(expectedSegmentsHandedOff)
     );
 
diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java
@@ -145,7 +145,7 @@ protected void setupLookups() throws Exception
     // Wait for it to be loaded.
     indexer.latchableEmitter().waitForEventAggregate(
         event -> event.hasMetricName("ingest/events/processed")
-                      .hasDimension(DruidMetrics.DATASOURCE, Collections.singletonList(dataSource)),
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource),
         agg -> agg.hasSumAtLeast(totalRows)
     );
     broker.latchableEmitter().waitForEvent(
diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeUnnestQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeUnnestQueryTest.java
@@ -37,7 +37,6 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
-import java.util.Collections;
 import java.util.List;
 
 public class EmbeddedMSQRealtimeUnnestQueryTest extends BaseRealtimeQueryTest
@@ -98,7 +97,7 @@ void setupAll()
     // Wait for it to be loaded.
     indexer.latchableEmitter().waitForEventAggregate(
         event -> event.hasMetricName("ingest/events/processed")
-                      .hasDimension(DruidMetrics.DATASOURCE, Collections.singletonList(dataSource)),
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource),
         agg -> agg.hasSumAtLeast(totalRows)
     );
   }
diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java
@@ -127,7 +127,7 @@ public void test_runKafkaSupervisor()
     Assertions.assertTrue(supervisorStatus.isSuspended());
     indexer.latchableEmitter().waitForEventAggregate(
         event -> event.hasMetricName("ingest/handoff/count")
-                      .hasDimension(DruidMetrics.DATASOURCE, List.of(dataSource)),
+                      .hasDimension(DruidMetrics.DATASOURCE, dataSource),
         agg -> agg.hasSumAtLeast(expectedSegments)
     );
     overlord.latchableEmitter().waitForEventAggregate(
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/TaskRealtimeMetricsMonitorBuilder.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/TaskRealtimeMetricsMonitorBuilder.java
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/stats/TaskRealtimeMetricsMonitor.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/stats/TaskRealtimeMetricsMonitor.java
@@ -19,20 +19,14 @@
 
 package org.apache.druid.indexing.common.stats;
 
-import com.google.common.collect.ImmutableMap;
 import org.apache.druid.java.util.emitter.EmittingLogger;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
 import org.apache.druid.java.util.metrics.AbstractMonitor;
-import org.apache.druid.java.util.metrics.MonitorUtils;
-import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.segment.incremental.RowIngestionMeters;
 import org.apache.druid.segment.incremental.RowIngestionMetersTotals;
 import org.apache.druid.segment.realtime.SegmentGenerationMetrics;
 
-import javax.annotation.Nullable;
-import java.util.Map;
-
 /**
  * Emits metrics from {@link SegmentGenerationMetrics} and {@link RowIngestionMeters}.
  */
@@ -42,24 +36,20 @@ public class TaskRealtimeMetricsMonitor extends AbstractMonitor
 
   private final SegmentGenerationMetrics segmentGenerationMetrics;
   private final RowIngestionMeters rowIngestionMeters;
-  private final Map<String, String[]> dimensions;
-  @Nullable
-  private final Map<String, Object> metricTags;
+  private final ServiceMetricEvent.Builder builder;
 
   private SegmentGenerationMetrics previousSegmentGenerationMetrics;
   private RowIngestionMetersTotals previousRowIngestionMetersTotals;
 
   public TaskRealtimeMetricsMonitor(
       SegmentGenerationMetrics segmentGenerationMetrics,
       RowIngestionMeters rowIngestionMeters,
-      Map<String, String[]> dimensions,
-      @Nullable Map<String, Object> metricTags
+      ServiceMetricEvent.Builder metricEventBuilder
   )
   {
     this.segmentGenerationMetrics = segmentGenerationMetrics;
     this.rowIngestionMeters = rowIngestionMeters;
-    this.dimensions = ImmutableMap.copyOf(dimensions);
-    this.metricTags = metricTags;
+    this.builder = metricEventBuilder;
     previousSegmentGenerationMetrics = new SegmentGenerationMetrics();
     previousRowIngestionMetersTotals = new RowIngestionMetersTotals(0, 0, 0, 0, 0);
   }
@@ -70,17 +60,13 @@ public boolean doMonitor(ServiceEmitter emitter)
     SegmentGenerationMetrics metrics = segmentGenerationMetrics.snapshot();
     RowIngestionMetersTotals rowIngestionMetersTotals = rowIngestionMeters.getTotals();
 
-    final ServiceMetricEvent.Builder builder = new ServiceMetricEvent.Builder();
-    MonitorUtils.addDimensionsToBuilder(builder, dimensions);
-
     final long thrownAway = rowIngestionMetersTotals.getThrownAway() - previousRowIngestionMetersTotals.getThrownAway();
     if (thrownAway > 0) {
       log.warn(
           "[%,d] events thrown away. Possible causes: null events, events filtered out by transformSpec, or events outside earlyMessageRejectionPeriod / lateMessageRejectionPeriod.",
           thrownAway
       );
     }
-    builder.setDimensionIfNotNull(DruidMetrics.TAGS, metricTags);
     emitter.emit(builder.setMetric("ingest/events/thrownAway", thrownAway));
 
     final long unparseable = rowIngestionMetersTotals.getUnparseable()
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/task/AbstractTask.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/task/AbstractTask.java
@@ -101,8 +101,6 @@ public static IngestionMode fromString(String name)
   private File reportsFile;
   private File statusFile;
 
-  private final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();
-
   private volatile CountDownLatch cleanupCompletionLatch;
 
   protected AbstractTask(String id, String dataSource, Map<String, Object> context, IngestionMode ingestionMode)
@@ -131,7 +129,6 @@ protected AbstractTask(
     // Copy the given context into a new mutable map because the Druid indexing service can add some internal contexts.
     this.context = context == null ? new HashMap<>() : new HashMap<>(context);
     this.ingestionMode = ingestionMode;
-    IndexTaskUtils.setTaskDimensions(metricBuilder, this);
   }
 
   protected AbstractTask(
@@ -385,8 +382,14 @@ public Map<String, Object> getContext()
     return context;
   }
 
-  protected ServiceMetricEvent.Builder getMetricBuilder()
+  /**
+   * @return A fresh instance of {@link ServiceMetricEvent.Builder} that can be
+   * used to emit metrics for this task.
+   */
+  public ServiceMetricEvent.Builder getMetricBuilder()
   {
+    final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();
+    IndexTaskUtils.setTaskDimensions(metricBuilder, this);
     return metricBuilder;
   }
 
@@ -426,13 +429,15 @@ private static IngestionMode computeIngestionMode(boolean isAppendToExisting, bo
         + "Either dropExisting or appendToExisting should be set to false");
   }
 
+  /**
+   * Emits a metric for this task using the {@link #getMetricBuilder() metric builder}.
+   */
   public void emitMetric(
       ServiceEmitter emitter,
       String metric,
       Number value
   )
   {
-
     if (emitter == null || metric == null || value == null) {
       return;
     }
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/task/IndexTask.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/task/IndexTask.java
@@ -48,7 +48,6 @@
 import org.apache.druid.indexer.partitions.PartitionsSpec;
 import org.apache.druid.indexer.partitions.SecondaryPartitionType;
 import org.apache.druid.indexer.report.TaskReport;
-import org.apache.druid.indexing.common.TaskRealtimeMetricsMonitorBuilder;
 import org.apache.druid.indexing.common.TaskToolbox;
 import org.apache.druid.indexing.common.actions.TaskActionClient;
 import org.apache.druid.indexing.common.stats.TaskRealtimeMetricsMonitor;
@@ -818,7 +817,7 @@ private TaskStatus generateAndPublishSegments(
   {
     final SegmentGenerationMetrics buildSegmentsSegmentGenerationMetrics = new SegmentGenerationMetrics();
     final TaskRealtimeMetricsMonitor metricsMonitor =
-        TaskRealtimeMetricsMonitorBuilder.build(this, buildSegmentsSegmentGenerationMetrics, buildSegmentsMeters);
+        new TaskRealtimeMetricsMonitor(buildSegmentsSegmentGenerationMetrics, buildSegmentsMeters, getMetricBuilder());
     toolbox.addMonitor(metricsMonitor);
 
     final PartitionsSpec partitionsSpec = partitionAnalysis.getPartitionsSpec();
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/PartialSegmentGenerateTask.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/PartialSegmentGenerateTask.java
@@ -25,7 +25,6 @@
 import org.apache.druid.indexer.TaskStatus;
 import org.apache.druid.indexer.partitions.PartitionsSpec;
 import org.apache.druid.indexer.report.TaskReport;
-import org.apache.druid.indexing.common.TaskRealtimeMetricsMonitorBuilder;
 import org.apache.druid.indexing.common.TaskToolbox;
 import org.apache.druid.indexing.common.stats.TaskRealtimeMetricsMonitor;
 import org.apache.druid.indexing.common.task.BatchAppenderators;
@@ -169,7 +168,7 @@ private List<DataSegment> generateSegments(
     final SegmentGenerationMetrics segmentGenerationMetrics = new SegmentGenerationMetrics();
     buildSegmentsMeters = toolbox.getRowIngestionMetersFactory().createRowIngestionMeters();
     final TaskRealtimeMetricsMonitor metricsMonitor =
-        TaskRealtimeMetricsMonitorBuilder.build(this, segmentGenerationMetrics, buildSegmentsMeters);
+        new TaskRealtimeMetricsMonitor(segmentGenerationMetrics, buildSegmentsMeters, getMetricBuilder());
     toolbox.addMonitor(metricsMonitor);
 
     final ParallelIndexTuningConfig tuningConfig = ingestionSchema.getTuningConfig();
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/SinglePhaseSubTask.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/SinglePhaseSubTask.java
@@ -34,7 +34,6 @@
 import org.apache.druid.indexer.granularity.GranularitySpec;
 import org.apache.druid.indexer.partitions.DynamicPartitionsSpec;
 import org.apache.druid.indexer.report.TaskReport;
-import org.apache.druid.indexing.common.TaskRealtimeMetricsMonitorBuilder;
 import org.apache.druid.indexing.common.TaskToolbox;
 import org.apache.druid.indexing.common.actions.SurrogateTaskActionClient;
 import org.apache.druid.indexing.common.actions.TaskActionClient;
@@ -362,7 +361,7 @@ private DataSegmentsWithSchemas generateAndPushSegments(
     final GranularitySpec granularitySpec = dataSchema.getGranularitySpec();
     final SegmentGenerationMetrics segmentGenerationMetrics = new SegmentGenerationMetrics();
     final TaskRealtimeMetricsMonitor metricsMonitor =
-        TaskRealtimeMetricsMonitorBuilder.build(this, segmentGenerationMetrics, rowIngestionMeters);
+        new TaskRealtimeMetricsMonitor(segmentGenerationMetrics, rowIngestionMeters, getMetricBuilder());
     toolbox.addMonitor(metricsMonitor);
 
     final ParallelIndexTuningConfig tuningConfig = ingestionSchema.getTuningConfig();
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTask.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTask.java
@@ -43,6 +43,8 @@
 import org.apache.druid.indexing.common.task.Tasks;
 import org.apache.druid.indexing.seekablestream.common.RecordSupplier;
 import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.query.NoopQueryRunner;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryRunner;
@@ -197,6 +199,13 @@ public String getCurrentRunnerStatus()
     return (status != null) ? status.toString() : null;
   }
 
+  @Override
+  public ServiceMetricEvent.Builder getMetricBuilder()
+  {
+    return super.getMetricBuilder()
+                .setDimensionIfNotNull(DruidMetrics.SUPERVISOR_ID, supervisorId);
+  }
+
   public Appenderator newAppenderator(
       TaskToolbox toolbox,
       SegmentGenerationMetrics metrics,
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java
@@ -59,7 +59,6 @@
 import org.apache.druid.indexing.common.LockGranularity;
 import org.apache.druid.indexing.common.TaskLock;
 import org.apache.druid.indexing.common.TaskLockType;
-import org.apache.druid.indexing.common.TaskRealtimeMetricsMonitorBuilder;
 import org.apache.druid.indexing.common.TaskToolbox;
 import org.apache.druid.indexing.common.actions.CheckPointDataSourceMetadataAction;
 import org.apache.druid.indexing.common.actions.ResetDataSourceMetadataAction;
@@ -431,7 +430,7 @@ private TaskStatus runInternal(TaskToolbox toolbox) throws Exception
     // Set up SegmentGenerationMetrics
     this.segmentGenerationMetrics = new SegmentGenerationMetrics();
     final TaskRealtimeMetricsMonitor metricsMonitor =
-        TaskRealtimeMetricsMonitorBuilder.build(task, segmentGenerationMetrics, rowIngestionMeters);
+        new TaskRealtimeMetricsMonitor(segmentGenerationMetrics, rowIngestionMeters, task.getMetricBuilder());
     toolbox.addMonitor(metricsMonitor);
 
     final String lookupTier = task.getContextValue(CTX_KEY_LOOKUP_TIER);
diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/TaskRealtimeMetricsMonitorTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/TaskRealtimeMetricsMonitorTest.java
@@ -25,6 +25,7 @@
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 import org.apache.druid.java.util.emitter.service.ServiceEventBuilder;
 import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.java.util.metrics.MonitorUtils;
 import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.segment.incremental.RowIngestionMeters;
 import org.apache.druid.segment.realtime.SegmentGenerationMetrics;
@@ -74,7 +75,11 @@ public void setUp()
           return null;
         })
         .when(emitter).emit(ArgumentMatchers.any(Event.class));
-    target = new TaskRealtimeMetricsMonitor(segmentGenerationMetrics, rowIngestionMeters, DIMENSIONS, TAGS);
+    target = new TaskRealtimeMetricsMonitor(
+        segmentGenerationMetrics,
+        rowIngestionMeters,
+        createMetricEventBuilder()
+    );
   }
 
   @Test
@@ -89,7 +94,11 @@ public void testdoMonitorShouldEmitUserProvidedTags()
   @Test
   public void testdoMonitorWithoutTagsShouldNotEmitTags()
   {
-    target = new TaskRealtimeMetricsMonitor(segmentGenerationMetrics, rowIngestionMeters, DIMENSIONS, null);
+    target = new TaskRealtimeMetricsMonitor(
+        segmentGenerationMetrics,
+        rowIngestionMeters,
+        createMetricEventBuilder()
+    );
     for (ServiceMetricEvent sme : emittedEvents.values()) {
       Assert.assertFalse(sme.getUserDims().containsKey(DruidMetrics.TAGS));
     }
@@ -98,7 +107,11 @@ public void testdoMonitorWithoutTagsShouldNotEmitTags()
   @Test
   public void testMessageGapAggStats()
   {
-    target = new TaskRealtimeMetricsMonitor(segmentGenerationMetrics, rowIngestionMeters, DIMENSIONS, null);
+    target = new TaskRealtimeMetricsMonitor(
+        segmentGenerationMetrics,
+        rowIngestionMeters,
+        createMetricEventBuilder()
+    );
 
     target.doMonitor(emitter);
     Assert.assertFalse(emittedEvents.containsKey("ingest/events/minMessageGap"));
@@ -113,4 +126,12 @@ public void testMessageGapAggStats()
     Assert.assertTrue(emittedEvents.containsKey("ingest/events/maxMessageGap"));
     Assert.assertTrue(emittedEvents.containsKey("ingest/events/avgMessageGap"));
   }
+
+  private ServiceMetricEvent.Builder createMetricEventBuilder()
+  {
+    final ServiceMetricEvent.Builder builder = new ServiceMetricEvent.Builder();
+    MonitorUtils.addDimensionsToBuilder(builder, DIMENSIONS);
+    builder.setDimensionIfNotNull(DruidMetrics.TAGS, TAGS);
+    return builder;
+  }
 }
diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/AbstractTaskTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/AbstractTaskTest.java

Original file line number	Diff line number	Diff line change
`@@ -101,8 +101,6 @@ public static IngestionMode fromString(String name)`
`101`	`101`	`private File reportsFile;`
`102`	`102`	`private File statusFile;`
`103`	`103`
`104`		`- private final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();`
`105`		`-`
`106`	`104`	`private volatile CountDownLatch cleanupCompletionLatch;`
`107`	`105`
`108`	`106`	`protected AbstractTask(String id, String dataSource, Map<String, Object> context, IngestionMode ingestionMode)`
`@@ -131,7 +129,6 @@ protected AbstractTask(`
`131`	`129`	`// Copy the given context into a new mutable map because the Druid indexing service can add some internal contexts.`
`132`	`130`	`this.context = context == null ? new HashMap<>() : new HashMap<>(context);`
`133`	`131`	`this.ingestionMode = ingestionMode;`
`134`		`- IndexTaskUtils.setTaskDimensions(metricBuilder, this);`
`135`	`132`	`}`
`136`	`133`
`137`	`134`	`protected AbstractTask(`
`@@ -385,8 +382,14 @@ public Map<String, Object> getContext()`
`385`	`382`	`return context;`
`386`	`383`	`}`
`387`	`384`
`388`		`- protected ServiceMetricEvent.Builder getMetricBuilder()`
	`385`	`+ /**`
	`386`	`+ * @return A fresh instance of {@link ServiceMetricEvent.Builder} that can be`
	`387`	`+ * used to emit metrics for this task.`
	`388`	`+ */`
	`389`	`+ public ServiceMetricEvent.Builder getMetricBuilder()`
`389`	`390`	`{`
	`391`	`+ final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();`
	`392`	`+ IndexTaskUtils.setTaskDimensions(metricBuilder, this);`
`390`	`393`	`return metricBuilder;`
`391`	`394`	`}`
`392`	`395`
`@@ -426,13 +429,15 @@ private static IngestionMode computeIngestionMode(boolean isAppendToExisting, bo`
`426`	`429`	`+ "Either dropExisting or appendToExisting should be set to false");`
`427`	`430`	`}`
`428`	`431`
	`432`	`+ /**`
	`433`	`+ * Emits a metric for this task using the {@link #getMetricBuilder() metric builder}.`
	`434`	`+ */`
`429`	`435`	`public void emitMetric(`
`430`	`436`	`ServiceEmitter emitter,`
`431`	`437`	`String metric,`
`432`	`438`	`Number value`
`433`	`439`	`)`
`434`	`440`	`{`
`435`		`-`
`436`	`441`	`if (emitter == null \|\| metric == null \|\| value == null) {`
`437`	`442`	`return;`
`438`	`443`	`}`