Add task dimensions to ServiceEmitter (#18876)

All task metrics are now emitted with the following dimensions: taskId, dataSource, taskType, groupId, and id (for backward compatibility; id will be removed in favor of the taskId dimension in a future release).

This is achieved by wiring in the TaskHolder and getting the task dimensions from it in the ServiceEmitter initialization, so individual monitors no longer need to pass the task holder to retrieve and append task dimensions on every monitor call. Some monitors were previously missing these dimensions. This change ensures that all task metrics are emitted with these dimensions by default and in a consistent format.

Author: abhishekrb19

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java

@@ -39,6 +39,7 @@
 import org.apache.commons.io.FileUtils;
 import org.apache.curator.test.TestingCluster;
 import org.apache.druid.cli.CliPeon;
+import org.apache.druid.cli.CliPeonTest;
 import org.apache.druid.cli.PeonLoadSpecHolder;
 import org.apache.druid.cli.PeonTaskHolder;
 import org.apache.druid.data.input.InputEntity;
@@ -88,6 +89,12 @@
 import org.apache.druid.java.util.common.guava.Sequences;
 import org.apache.druid.java.util.common.parsers.CloseableIterator;
 import org.apache.druid.java.util.common.parsers.ParseException;
+import org.apache.druid.java.util.emitter.core.Emitter;
+import org.apache.druid.java.util.emitter.core.Event;
+import org.apache.druid.java.util.emitter.core.EventMap;
+import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.java.util.metrics.StubServiceEmitter;
+import org.apache.druid.java.util.metrics.TaskHolder;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.query.DefaultGenericQueryMetricsFactory;
 import org.apache.druid.query.DefaultQueryRunnerFactoryConglomerate;
@@ -121,7 +128,6 @@
 import org.apache.druid.server.coordination.BroadcastDatasourceLoadingSpec;
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
 import org.apache.druid.server.metrics.LoadSpecHolder;
-import org.apache.druid.server.metrics.TaskHolder;
 import org.apache.druid.server.security.Action;
 import org.apache.druid.server.security.Resource;
 import org.apache.druid.server.security.ResourceAction;
@@ -3487,6 +3493,64 @@ public void testTaskWithTransformSpecDoesNotCauseCliPeonCyclicDependency()
     verifyPersistAndMergeTimeMetricsArePositive(observedSegmentGenerationMetrics);
   }
 
+  @Test(timeout = 60_000L)
+  public void testKafkaTaskContainsAllTaskDimensions()
+      throws IOException, ExecutionException, InterruptedException
+  {
+    insertData();
+
+    final KafkaIndexTask task = createTask(
+        "index_kafka_test_id1",
+        new KafkaIndexTaskIOConfig(
+            0,
+            "sequence0",
+            new SeekableStreamStartSequenceNumbers<>(topic, ImmutableMap.of(new KafkaTopicPartition(false, topic, 0), 2L), ImmutableSet.of()),
+            new SeekableStreamEndSequenceNumbers<>(topic, ImmutableMap.of(new KafkaTopicPartition(false, topic, 0), 5L)),
+            kafkaServer.consumerProperties(),
+            KafkaSupervisorIOConfig.DEFAULT_POLL_TIMEOUT_MILLIS,
+            true,
+            null,
+            null,
+            INPUT_FORMAT,
+            null,
+            Duration.standardHours(2).getStandardMinutes()
+        )
+    );
+
+    Injector peonInjector = CliPeonTest.makePeonInjectorWithStubEmitter(task, temporaryFolder, OBJECT_MAPPER);
+    Emitter peonEmitter = peonInjector.getInstance(Emitter.class);
+    Assert.assertTrue(peonEmitter instanceof StubServiceEmitter);
+    emitter = (StubServiceEmitter) peonEmitter;
+    emitter.start();
+    makeToolboxFactory();
+
+    final ListenableFuture<TaskStatus> future = runTask(task);
+
+    // Wait for task to exit
+    Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());
+    Assert.assertTrue(emitter.getNumEmittedEvents() > 0);
+
+    // Check published metadata & segments in deep storage
+    final List<SegmentDescriptor> publishedDescriptors = publishedDescriptors();
+    Assert.assertEquals(
+        new KafkaDataSourceMetadata(new SeekableStreamEndSequenceNumbers<>(topic, ImmutableMap.of(new KafkaTopicPartition(false, topic, 0), 5L))),
+        newDataSchemaMetadata()
+    );
+    Assert.assertEquals(ImmutableList.of("c"), readSegmentColumn("dim1", publishedDescriptors.get(0)));
+
+    Assert.assertTrue(emitter.getNumEmittedEvents() > 0);
+    for (Event event : emitter.getEvents()) {
+      if (event instanceof ServiceMetricEvent) {
+        EventMap observedEvent = event.toMap();
+        Assert.assertEquals("test_ds", observedEvent.get("dataSource"));
+        Assert.assertEquals("index_kafka_test_id1", observedEvent.get("id"));
+        Assert.assertEquals("index_kafka_test_id1", observedEvent.get("taskId"));
+        Assert.assertEquals("index_kafka", observedEvent.get("taskType"));
+        Assert.assertEquals("index_kafka_test_ds", observedEvent.get("groupId"));
+      }
+    }
+  }
+
   public static class TestKafkaInputFormat implements InputFormat
   {
     final InputFormat baseInputFormat;

processing/src/main/java/org/apache/druid/java/util/emitter/service/ServiceEmitter.java

@@ -25,39 +25,56 @@
 import org.apache.druid.java.util.common.lifecycle.LifecycleStop;
 import org.apache.druid.java.util.emitter.core.Emitter;
 import org.apache.druid.java.util.emitter.core.Event;
+import org.apache.druid.java.util.metrics.NoopTaskHolder;
+import org.apache.druid.java.util.metrics.TaskHolder;
 
 import java.io.IOException;
 
 public class ServiceEmitter implements Emitter
 {
-  private final ImmutableMap<String, String> serviceDimensions;
   private final Emitter emitter;
+  private final String service;
+  private final ImmutableMap<String, String> otherServiceDimensions;
+  private final String host;
+  private final TaskHolder taskHolder;
+
+  /**
+   * This is initialized in {@link #start()} rather than in the constructor, since calling {@link TaskHolder#getMetricDimensions()}
+   * may introduce cyclic dependencies. So we defer initialization until {@link #start()} which is {@link LifecycleStart} managed.
+   */
+  private ImmutableMap<String, String> serviceDimensions;
 
   public ServiceEmitter(String service, String host, Emitter emitter)
   {
-    this(service, host, emitter, ImmutableMap.of());
+    this(service, host, emitter, ImmutableMap.of(), new NoopTaskHolder());
   }
 
   public ServiceEmitter(
       String service,
       String host,
       Emitter emitter,
-      ImmutableMap<String, String> otherServiceDimensions
+      ImmutableMap<String, String> otherServiceDimensions,
+      TaskHolder taskHolder
   )
   {
-    this.serviceDimensions = ImmutableMap
-        .<String, String>builder()
-        .put("service", Preconditions.checkNotNull(service, "service should be non-null"))
-        .put("host", Preconditions.checkNotNull(host, "host should be non-null"))
-        .putAll(otherServiceDimensions)
-        .build();
+    this.service = Preconditions.checkNotNull(service, "service should be non-null");
+    this.host = Preconditions.checkNotNull(host, "host should be non-null");
+    this.otherServiceDimensions = otherServiceDimensions;
     this.emitter = emitter;
+    this.taskHolder = taskHolder;
   }
 
   @Override
   @LifecycleStart
   public void start()
   {
+    serviceDimensions = ImmutableMap
+        .<String, String>builder()
+        .put(Event.SERVICE, service)
+        .put(Event.HOST, host)
+        .putAll(otherServiceDimensions)
+        .putAll(taskHolder.getMetricDimensions())
+        .build();
     emitter.start();
   }
 

processing/src/main/java/org/apache/druid/java/util/metrics/CgroupCpuMonitor.java

@@ -40,23 +40,21 @@ public class CgroupCpuMonitor extends FeedDefiningMonitor
   private static final Logger LOG = new Logger(CgroupCpuMonitor.class);
   private static final Long DEFAULT_USER_HZ = 100L;
   final CgroupDiscoverer cgroupDiscoverer;
-  final Map<String, String[]> dimensions;
   private Long userHz;
   private final KeyedDiff jiffies = new KeyedDiff();
   private long prevJiffiesSnapshotAt = 0;
   private final boolean isRunningOnCgroupsV2;
   private final CgroupV2CpuMonitor cgroupV2CpuMonitor;
 
-  public CgroupCpuMonitor(CgroupDiscoverer cgroupDiscoverer, final Map<String, String[]> dimensions, String feed)
+  public CgroupCpuMonitor(CgroupDiscoverer cgroupDiscoverer, String feed)
   {
     super(feed);
     this.cgroupDiscoverer = cgroupDiscoverer;
-    this.dimensions = dimensions;
-    
+
     // Check if we're running on cgroups v2
     this.isRunningOnCgroupsV2 = cgroupDiscoverer.getCgroupVersion().equals(CgroupVersion.V2);
     if (isRunningOnCgroupsV2) {
-      this.cgroupV2CpuMonitor = new CgroupV2CpuMonitor(cgroupDiscoverer, dimensions, feed);
+      this.cgroupV2CpuMonitor = new CgroupV2CpuMonitor(cgroupDiscoverer, feed);
       LOG.info("Detected cgroups v2, using CgroupV2CpuMonitor behavior for accurate metrics");
     } else {
       this.cgroupV2CpuMonitor = null;
@@ -65,19 +63,14 @@ public CgroupCpuMonitor(CgroupDiscoverer cgroupDiscoverer, final Map<String, Str
 
   }
 
-  public CgroupCpuMonitor(final Map<String, String[]> dimensions, String feed)
-  {
-    this(ProcSelfCgroupDiscoverer.autoCgroupDiscoverer(), dimensions, feed);
-  }
-
-  public CgroupCpuMonitor(final Map<String, String[]> dimensions)
+  public CgroupCpuMonitor(String feed)
   {
-    this(dimensions, DEFAULT_METRICS_FEED);
+    this(ProcSelfCgroupDiscoverer.autoCgroupDiscoverer(), feed);
   }
 
   public CgroupCpuMonitor()
   {
-    this(ImmutableMap.of());
+    this(DEFAULT_METRICS_FEED);
   }
 
   @Override
@@ -96,7 +89,6 @@ private boolean doMonitorV1(ServiceEmitter emitter)
     long now = Instant.now().getEpochSecond();
 
     final ServiceMetricEvent.Builder builder = builder();
-    MonitorUtils.addDimensionsToBuilder(builder, dimensions);
     builder.setDimension("cgroupversion", cgroupDiscoverer.getCgroupVersion().name());
 
     emitter.emit(builder.setMetric("cgroup/cpu/shares", cpuSnapshot.getShares()));

processing/src/main/java/org/apache/druid/java/util/metrics/CgroupCpuSetMonitor.java

@@ -19,52 +19,41 @@
 
 package org.apache.druid.java.util.metrics;
 
-import com.google.common.collect.ImmutableMap;
 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.cgroups.CgroupDiscoverer;
 import org.apache.druid.java.util.metrics.cgroups.CpuSet;
 import org.apache.druid.java.util.metrics.cgroups.ProcSelfCgroupDiscoverer;
 
-import java.util.Map;
-
 /**
  * Monitor that reports CPU set metrics from cgroups both v1 and v2.
  */
 
 public class CgroupCpuSetMonitor extends FeedDefiningMonitor
 {
   final CgroupDiscoverer cgroupDiscoverer;
-  final Map<String, String[]> dimensions;
 
-  public CgroupCpuSetMonitor(CgroupDiscoverer cgroupDiscoverer, final Map<String, String[]> dimensions, String feed)
+  public CgroupCpuSetMonitor(CgroupDiscoverer cgroupDiscoverer, String feed)
   {
     super(feed);
     this.cgroupDiscoverer = cgroupDiscoverer;
-    this.dimensions = dimensions;
-  }
-
-  public CgroupCpuSetMonitor(final Map<String, String[]> dimensions, String feed)
-  {
-    this(ProcSelfCgroupDiscoverer.autoCgroupDiscoverer(), dimensions, feed);
   }
 
-  public CgroupCpuSetMonitor(final Map<String, String[]> dimensions)
+  public CgroupCpuSetMonitor(String feed)
   {
-    this(dimensions, DEFAULT_METRICS_FEED);
+    this(ProcSelfCgroupDiscoverer.autoCgroupDiscoverer(), feed);
   }
 
   public CgroupCpuSetMonitor()
   {
-    this(ImmutableMap.of());
+    this(DEFAULT_METRICS_FEED);
   }
 
   @Override
   public boolean doMonitor(ServiceEmitter emitter)
   {
     final CpuSet.CpuSetMetric cpusetSnapshot = cgroupDiscoverer.getCpuSetMetrics();
     final ServiceMetricEvent.Builder builder = builder();
-    MonitorUtils.addDimensionsToBuilder(builder, dimensions);
     builder.setDimension("cgroupversion", cgroupDiscoverer.getCgroupVersion().name());
     emitter.emit(builder.setMetric(
         "cgroup/cpuset/cpu_count",

processing/src/main/java/org/apache/druid/java/util/metrics/CgroupDiskMonitor.java

@@ -33,41 +33,34 @@
 public class CgroupDiskMonitor extends FeedDefiningMonitor
 {
   private static final Logger LOG = new Logger(CgroupDiskMonitor.class);
-  final CgroupDiscoverer cgroupDiscoverer;
-  final Map<String, String[]> dimensions;
+  private final CgroupDiscoverer cgroupDiscoverer;
   private final KeyedDiff diff = new KeyedDiff();
   private final boolean isRunningOnCgroupsV2;
   private final CgroupV2DiskMonitor cgroupV2DiskMonitor;
 
-  public CgroupDiskMonitor(CgroupDiscoverer cgroupDiscoverer, final Map<String, String[]> dimensions, String feed)
+  public CgroupDiskMonitor(CgroupDiscoverer cgroupDiscoverer, String feed)
   {
     super(feed);
     this.cgroupDiscoverer = cgroupDiscoverer;
-    this.dimensions = dimensions;
-    
+
     // Check if we're running on cgroups v2
     this.isRunningOnCgroupsV2 = cgroupDiscoverer.getCgroupVersion().equals(CgroupVersion.V2);
     if (isRunningOnCgroupsV2) {
-      this.cgroupV2DiskMonitor = new CgroupV2DiskMonitor(cgroupDiscoverer, dimensions, feed);
+      this.cgroupV2DiskMonitor = new CgroupV2DiskMonitor(cgroupDiscoverer, feed);
       LOG.info("Detected cgroups v2, using CgroupV2DiskMonitor behavior for accurate metrics");
     } else {
       this.cgroupV2DiskMonitor = null;
     }
   }
 
-  public CgroupDiskMonitor(final Map<String, String[]> dimensions, String feed)
-  {
-    this(ProcSelfCgroupDiscoverer.autoCgroupDiscoverer(), dimensions, feed);
-  }
-
-  public CgroupDiskMonitor(final Map<String, String[]> dimensions)
+  public CgroupDiskMonitor(String feed)
   {
-    this(dimensions, DEFAULT_METRICS_FEED);
+    this(ProcSelfCgroupDiscoverer.autoCgroupDiscoverer(), feed);
   }
 
   public CgroupDiskMonitor()
   {
-    this(ImmutableMap.of());
+    this(DEFAULT_METRICS_FEED);
   }
 
   @Override
@@ -97,7 +90,6 @@ private boolean doMonitorV1(ServiceEmitter emitter)
       if (stats != null) {
         final ServiceMetricEvent.Builder builder = builder()
             .setDimension("diskName", entry.getValue().getDiskName());
-        MonitorUtils.addDimensionsToBuilder(builder, dimensions);
         builder.setDimension("cgroupversion", cgroupDiscoverer.getCgroupVersion());
         for (Map.Entry<String, Long> stat : stats.entrySet()) {
           emitter.emit(builder.setMetric(stat.getKey(), stat.getValue()));