Add stream topic label to realtime ingestion delay metrics

pinot-common/src/main/java/org/apache/pinot/common/metrics/AbstractMetrics.java

@@ -32,6 +32,7 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.pinot.common.Utils;
 import org.apache.pinot.spi.metrics.PinotGauge;
 import org.apache.pinot.spi.metrics.PinotMeter;
@@ -565,6 +566,32 @@ public void setOrUpdatePartitionGauge(final String tableName, final int partitio
     setOrUpdateGauge(fullGaugeName, valueSupplier);
   }
 
+  /**
+   * Registers a per-partition table gauge keyed by table name with type, Kafka (or other stream) topic, and partition
+   * group id. The resulting MBean name matches {@code server.yml} rules so Prometheus exports {@code topic} and
+   * {@code partition} labels (see apache/pinot#18099).
+   */
+  public void setOrUpdatePartitionGaugeForStreamTopic(final String tableNameWithType, final String topicName,
+      final int partitionGroupId, final G gauge, final Supplier<Long> valueSupplier) {
+    Preconditions.checkArgument(StringUtils.isNotBlank(topicName), "topicName must not be blank");
+    setOrUpdateTableGauge(composeStreamTopicPartitionKey(tableNameWithType, topicName, partitionGroupId), gauge,
+        valueSupplier);
+  }
+
+  /**
+   * Removes a gauge registered via {@link #setOrUpdatePartitionGaugeForStreamTopic}.
+   */
+  public void removePartitionGaugeForStreamTopic(final String tableNameWithType, final String topicName,
+      final int partitionGroupId, final G gauge) {
+    Preconditions.checkArgument(StringUtils.isNotBlank(topicName), "topicName must not be blank");
+    removeTableGauge(composeStreamTopicPartitionKey(tableNameWithType, topicName, partitionGroupId), gauge);
+  }
+
+  private static String composeStreamTopicPartitionKey(String tableNameWithType, String topicName,
+      int partitionGroupId) {
+    return tableNameWithType + "-" + topicName + "-" + partitionGroupId;
+  }
+
   /**
    * @deprecated please use setOrUpdateGauge(final String metricName, final Supplier<Long> valueSupplier) instead.
    *

pinot-common/src/test/java/org/apache/pinot/common/metrics/AbstractMetricsTest.java

@@ -456,6 +456,24 @@ public void testPartitionGauges() {
     Assert.assertTrue(controllerMetrics.getMetricsRegistry().allMetrics().isEmpty());
   }
 
+  @Test
+  public void testPartitionGaugesWithStreamTopic() {
+    ControllerMetrics controllerMetrics = buildTestMetrics();
+    String tableWithType = "myTable_REALTIME";
+    String topic = "events";
+    int partitionGroupId = 3;
+    String compositeKey = tableWithType + "-" + topic + "-" + partitionGroupId;
+
+    controllerMetrics.setOrUpdatePartitionGaugeForStreamTopic(tableWithType, topic, partitionGroupId,
+        ControllerGauge.VERSION, () -> 7L);
+    Assert.assertEquals(MetricValueUtils.getGaugeValue(controllerMetrics,
+        ControllerGauge.VERSION.getGaugeName() + "." + compositeKey), 7);
+
+    controllerMetrics.removePartitionGaugeForStreamTopic(tableWithType, topic, partitionGroupId,
+        ControllerGauge.VERSION);
+    Assert.assertTrue(controllerMetrics.getMetricsRegistry().allMetrics().isEmpty());
+  }
+
   @Test
   public void testAddCallbackGauges() {
     ControllerMetrics controllerMetrics = buildTestMetrics();

pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/PinotPrometheusMetricsTest.java

@@ -315,6 +315,10 @@ public static class ExportedLabels {
     public static final List<String> PARTITION_TABLENAME_TABLETYPE =
         List.of(PARTITION, "3", TABLE, ExportedLabelValues.TABLENAME, TABLETYPE, TableType.REALTIME.toString());
 
+    public static final List<String> PARTITIONNUM_TABLENAME_TABLETYPE_KAFKATOPIC =
+        List.of(PARTITION, "3", TABLE, ExportedLabelValues.TABLENAME, TABLETYPE, TableType.REALTIME.toString(), TOPIC,
+            KAFKA_TOPIC);
+
     public static final List<String> TABLENAME_TABLETYPE_CONTROLLER_TASKTYPE =
         List.of(TABLE, ExportedLabelValues.TABLENAME, TABLETYPE, TABLETYPE_REALTIME, TASKTYPE,
             CONTROLLER_PERIODIC_TASK_CHC);

pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/ServerPrometheusMetricsTest.java

@@ -54,9 +54,12 @@ public abstract class ServerPrometheusMetricsTest extends PinotPrometheusMetrics
 
   private static final List<ServerGauge> GAUGES_ACCEPTING_PARTITION =
       List.of(ServerGauge.UPSERT_VALID_DOC_ID_SNAPSHOT_COUNT, ServerGauge.UPSERT_PRIMARY_KEYS_IN_SNAPSHOT_COUNT,
-          ServerGauge.REALTIME_INGESTION_OFFSET_LAG, ServerGauge.REALTIME_INGESTION_DELAY_MS,
-          ServerGauge.UPSERT_PRIMARY_KEYS_COUNT, ServerGauge.DEDUP_PRIMARY_KEYS_COUNT,
-          ServerGauge.REALTIME_INGESTION_UPSTREAM_OFFSET, ServerGauge.REALTIME_INGESTION_CONSUMING_OFFSET);
+          ServerGauge.REALTIME_INGESTION_OFFSET_LAG, ServerGauge.UPSERT_PRIMARY_KEYS_COUNT,
+          ServerGauge.DEDUP_PRIMARY_KEYS_COUNT, ServerGauge.REALTIME_INGESTION_UPSTREAM_OFFSET,
+          ServerGauge.REALTIME_INGESTION_CONSUMING_OFFSET);
+
+  private static final List<ServerGauge> GAUGES_ACCEPTING_PARTITION_WITH_STREAM_TOPIC =
+      List.of(ServerGauge.REALTIME_INGESTION_DELAY_MS, ServerGauge.END_TO_END_REALTIME_INGESTION_DELAY_MS);
 
   private static final List<ServerGauge> GAUGES_ACCEPTING_RAW_TABLE_NAME =
       List.of(ServerGauge.REALTIME_OFFHEAP_MEMORY_USED, ServerGauge.REALTIME_SEGMENT_NUM_PARTITIONS,
@@ -122,6 +125,10 @@ public void gaugeTest(ServerGauge serverGauge) {
         addGaugeWithLabels(serverGauge, CLIENT_ID);
         assertGaugeExportedCorrectly(serverGauge.getGaugeName(),
             ExportedLabels.PARTITION_TABLENAME_TABLETYPE_KAFKATOPIC, EXPORTED_METRIC_PREFIX);
+      } else if (GAUGES_ACCEPTING_PARTITION_WITH_STREAM_TOPIC.contains(serverGauge)) {
+        addPartitionGaugeWithStreamTopicLabels(serverGauge, TABLE_NAME_WITH_TYPE);
+        assertGaugeExportedCorrectly(serverGauge.getGaugeName(),
+            ExportedLabels.PARTITIONNUM_TABLENAME_TABLETYPE_KAFKATOPIC, EXPORTED_METRIC_PREFIX);
       } else if (GAUGES_ACCEPTING_PARTITION.contains(serverGauge)) {
         addPartitionGaugeWithLabels(serverGauge, TABLE_NAME_WITH_TYPE);
         assertGaugeExportedCorrectly(serverGauge.getGaugeName(), ExportedLabels.PARTITION_TABLENAME_TABLETYPE,
@@ -145,6 +152,10 @@ private void addPartitionGaugeWithLabels(ServerGauge serverGauge, String labels)
     _serverMetrics.setValueOfPartitionGauge(labels, 3, serverGauge, 100L);
   }
 
+  private void addPartitionGaugeWithStreamTopicLabels(ServerGauge serverGauge, String tableNameWithType) {
+    _serverMetrics.setValueOfTableGauge(tableNameWithType + "-" + KAFKA_TOPIC + "-3", serverGauge, 100L);
+  }
+
   public void addMeterWithLabels(ServerMeter serverMeter, String labels) {
     _serverMetrics.addMeteredTableValue(labels, serverMeter, 4L);
   }

pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java

@@ -35,6 +35,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.BooleanSupplier;
 import javax.annotation.Nullable;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.pinot.common.metrics.ServerGauge;
 import org.apache.pinot.common.metrics.ServerMeter;
 import org.apache.pinot.common.metrics.ServerMetrics;
@@ -98,12 +99,15 @@ private static class IngestionInfo {
     @Nullable
     volatile StreamPartitionMsgOffset _currentOffset;
     volatile long _firstStreamIngestionTimeMs;
+    @Nullable
+    final String _streamTopicName;
 
     IngestionInfo(long ingestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffset,
-        long firstStreamIngestionTimeMs) {
+        long firstStreamIngestionTimeMs, @Nullable String streamTopicName) {
       _ingestionTimeMs = ingestionTimeMs;
       _currentOffset = currentOffset;
       _firstStreamIngestionTimeMs = firstStreamIngestionTimeMs;
+      _streamTopicName = streamTopicName;
     }
 
     void update(long ingestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffset,
@@ -333,9 +337,9 @@ public Thread newThread(Runnable r) {
    */
   private void removePartitionId(int partitionId) {
     _partitionsHostedByThisServer.remove(partitionId);
-    _ingestionInfoMap.remove(partitionId);
+    IngestionInfo removed = _ingestionInfoMap.remove(partitionId);
     _partitionsTracked.computeIfPresent(partitionId, (k, v) -> {
-      removeMetrics(partitionId);
+      removeMetrics(partitionId, removed != null ? removed._streamTopicName : null);
       return null;
     });
 
@@ -394,7 +398,7 @@ void createMetrics(int partitionId) {
     LOGGER.info("Successfully created ingestion metrics for partition id: {}", partitionId);
   }
 
-  private void removeMetrics(int partitionId) {
+  private void removeMetrics(int partitionId, @Nullable String streamTopicName) {
     int streamConfigIndex = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(partitionId);
     StreamMetadataProvider streamMetadataProvider =
         _streamConfigIndexToStreamMetadataProvider.get(streamConfigIndex);
@@ -410,6 +414,12 @@ private void removeMetrics(int partitionId) {
         ServerGauge.END_TO_END_REALTIME_INGESTION_DELAY_MS);
     _serverMetrics.removePartitionGauge(_metricName, partitionId,
         ServerGauge.REALTIME_INGESTION_DELAY_REPORTING_STATUS);
+    if (StringUtils.isNotBlank(streamTopicName)) {
+      _serverMetrics.removePartitionGaugeForStreamTopic(_metricName, streamTopicName, partitionId,
+          ServerGauge.REALTIME_INGESTION_DELAY_MS);
+      _serverMetrics.removePartitionGaugeForStreamTopic(_metricName, streamTopicName, partitionId,
+          ServerGauge.END_TO_END_REALTIME_INGESTION_DELAY_MS);
+    }
 
     LOGGER.info("Successfully removed ingestion metrics for partition id: {}", partitionId);
   }
@@ -420,12 +430,14 @@ private void removeMetrics(int partitionId) {
    * @param segmentName                name of the consuming segment
    * @param partitionId                partition id of the consuming segment (directly passed in to avoid parsing the
    *                                   segment name)
+   * @param streamTopicName            stream topic name for this consumer (e.g. Kafka topic); when blank, metrics use
+   *                                   the legacy name without a topic label
    * @param ingestionTimeMs            ingestion time of the last consumed message (from {@link StreamMessageMetadata})
    * @param firstStreamIngestionTimeMs ingestion time of the last consumed message in the first stream (from
    * {@link StreamMessageMetadata})
    * @param currentOffset              offset of the last consumed message (from {@link StreamMessageMetadata})
    */
-  public void updateMetrics(String segmentName, int partitionId, long ingestionTimeMs,
+  public void updateMetrics(String segmentName, int partitionId, @Nullable String streamTopicName, long ingestionTimeMs,
       long firstStreamIngestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffset) {
     if (!_isServerReadyToServeQueries.getAsBoolean() || _realTimeTableDataManager.isShutDown()) {
       // Do not update the ingestion delay metrics during server startup period
@@ -444,7 +456,7 @@ public void updateMetrics(String segmentName, int partitionId, long ingestionTim
         return v;
       }
       if (v == null) {
-        return new IngestionInfo(ingestionTimeMs, currentOffset, firstStreamIngestionTimeMs);
+        return new IngestionInfo(ingestionTimeMs, currentOffset, firstStreamIngestionTimeMs, streamTopicName);
       }
       v.update(ingestionTimeMs, currentOffset, firstStreamIngestionTimeMs);
       return v;

pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java

@@ -2110,7 +2110,8 @@ private void updateIngestionMetrics(StreamMessageMetadata metadata) {
     if (metadata != null) {
       try {
         _realtimeTableDataManager.updateIngestionMetrics(_segmentNameStr, _partitionGroupId,
-            metadata.getRecordIngestionTimeMs(), metadata.getFirstStreamRecordIngestionTimeMs(), metadata.getOffset());
+            _streamConfig.getTopicName(), metadata.getRecordIngestionTimeMs(),
+            metadata.getFirstStreamRecordIngestionTimeMs(), metadata.getOffset());
       } catch (Exception e) {
         _segmentLogger.warn("Failed to update the ingestion metrics", e);
       }
@@ -2123,8 +2124,8 @@ private void updateIngestionMetrics(StreamMessageMetadata metadata) {
    */
   private void setIngestionDelayToZero() {
     long currentTimeMs = System.currentTimeMillis();
-    _realtimeTableDataManager.updateIngestionMetrics(_segmentNameStr, _partitionGroupId, currentTimeMs, currentTimeMs,
-        null);
+    _realtimeTableDataManager.updateIngestionMetrics(_segmentNameStr, _partitionGroupId, _streamConfig.getTopicName(),
+        currentTimeMs, currentTimeMs, null);
   }
 
   // This should be done during commit? We may not always commit when we build a segment....

pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java

@@ -310,16 +310,18 @@ protected void doShutdown() {
    * @param segmentName                      name of the consuming segment
    * @param partitionId                      partition id of the consuming segment (directly passed in to avoid parsing
    *                                         the segment name)
+   * @param streamTopicName                  stream topic for this consumer (e.g. Kafka topic); used for ingestion
+   *                                         delay metric labels
    * @param ingestionTimeMs                  ingestion time of the last consumed message (from
    *                                         {@link StreamMessageMetadata})
    * @param firstStreamIngestionTimeMs ingestion time of the last consumed message in the first stream (from
    * {@link StreamMessageMetadata})
    * @param currentOffset                    offset of the last consumed message (from {@link StreamMessageMetadata})
    */
-  public void updateIngestionMetrics(String segmentName, int partitionId, long ingestionTimeMs,
-      long firstStreamIngestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffset) {
-    _ingestionDelayTracker.updateMetrics(segmentName, partitionId, ingestionTimeMs, firstStreamIngestionTimeMs,
-        currentOffset);
+  public void updateIngestionMetrics(String segmentName, int partitionId, @Nullable String streamTopicName,
+      long ingestionTimeMs, long firstStreamIngestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffset) {
+    _ingestionDelayTracker.updateMetrics(segmentName, partitionId, streamTopicName, ingestionTimeMs,
+        firstStreamIngestionTimeMs, currentOffset);
   }
 
   /**