minor: Add segment/allocated/count metric

docs/operations/metrics.md

@@ -351,6 +351,7 @@ If the JVM does not support CPU time measurement for the current thread, `ingest
 |`segment/added/bytes`|Size in bytes of new segments created.| `dataSource`, `taskId`, `taskType`, `groupId`, `interval`, `tags`|Varies|
 |`segment/moved/bytes`|Size in bytes of segments moved/archived via the Move Task.| `dataSource`, `taskId`, `taskType`, `groupId`, `interval`, `tags`|Varies|
 |`segment/nuked/bytes`|Size in bytes of segments deleted via the Kill Task.| `dataSource`, `taskId`, `taskType`, `groupId`, `interval`, `tags`|Varies|
+|`segment/allocated/count`|Number of segments successfully allocated by the Overlord to an append (realtime or batch) task. May be emitted multiple times for a single allocated segment ID if multiple tasks request an allocation with the same parameters.|`id`, `dataSource`, `taskId`, `taskType`, `groupId`, `tags`, `supervisorId`|Always 1|
 |`task/success/count`|Number of successful tasks per emission period. This metric is available only if the `TaskCountStatsMonitor` module is included.| `dataSource`,`taskType`, `supervisorId`|Varies|
 |`task/failed/count`|Number of failed tasks per emission period. This metric is available only if the `TaskCountStatsMonitor` module is included.|`dataSource`,`taskType`, `supervisorId`|Varies|
 |`task/running/count`|Number of current running tasks. This metric is available only if the `TaskCountStatsMonitor` module is included.|`dataSource`,`taskType`, `supervisorId`|Varies|

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

@@ -3479,8 +3479,8 @@ public void testKafkaTaskContainsAllTaskDimensions()
     for (Event event : emitter.getEvents()) {
       if (event instanceof ServiceMetricEvent) {
         EventMap observedEvent = event.toMap();
+        // Do not verify emission of "id" dimension as that is deprecated in favor of "taskId"
         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"));

indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java

@@ -26,6 +26,7 @@
 import org.apache.druid.error.DruidException;
 import org.apache.druid.indexing.common.LockGranularity;
 import org.apache.druid.indexing.common.TaskLockType;
+import org.apache.druid.indexing.common.task.IndexTaskUtils;
 import org.apache.druid.indexing.common.task.PendingSegmentAllocatingTask;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator;
@@ -51,7 +52,6 @@
 import java.util.Set;
 import java.util.concurrent.Future;
 import java.util.concurrent.ThreadLocalRandom;
-import java.util.stream.Collectors;
 
 /**
  * Allocates a pending segment for a given timestamp.
@@ -241,6 +241,7 @@ public SegmentIdWithShardSpec perform(
         identifier = tryAllocateSubsequentSegment(toolbox, task, rowInterval, usedSegmentsForRow.iterator().next());
       }
       if (identifier != null) {
+        IndexTaskUtils.emitSegmentAllocateMetric(identifier, task, toolbox.getEmitter());
         return identifier;
       }
 
@@ -287,7 +288,7 @@ private SegmentIdWithShardSpec tryAllocateFirstSegment(TaskActionToolbox toolbox
     final List<Interval> tryIntervals = Granularity.granularitiesFinerThan(preferredSegmentGranularity)
                                                    .stream()
                                                    .map(granularity -> granularity.bucket(timestamp))
-                                                   .collect(Collectors.toList());
+                                                   .toList();
     for (Interval tryInterval : tryIntervals) {
       if (tryInterval.contains(rowInterval)) {
         final SegmentIdWithShardSpec identifier = tryAllocate(toolbox, task, tryInterval, rowInterval, false);

indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocationQueue.java

@@ -745,6 +745,7 @@ void handleResult(SegmentAllocateResult result, SegmentAllocateRequest request)
 
       if (result.isSuccess()) {
         emitTaskMetric("task/action/success/count", 1L, request);
+        IndexTaskUtils.emitSegmentAllocateMetric(result.getSegmentId(), request.getTask(), emitter);
         requestToFuture.remove(request).complete(result.getSegmentId());
       } else if (request.canRetry()) {
         log.debug(

indexing-service/src/main/java/org/apache/druid/indexing/common/task/IndexTaskUtils.java

@@ -22,8 +22,10 @@
 import org.apache.druid.indexer.TaskStatus;
 import org.apache.druid.indexing.common.actions.TaskActionToolbox;
 import org.apache.druid.indexing.overlord.SegmentPublishResult;
+import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.emitter.service.SegmentMetadataEvent;
+import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
 import org.apache.druid.metadata.PendingSegmentRecord;
 import org.apache.druid.query.DruidMetrics;
@@ -107,6 +109,23 @@ public static void emitSegmentPublishMetrics(
     }
   }
 
+  /**
+   * Emits the metric {@code segment/allocated/count}.
+   */
+  public static void emitSegmentAllocateMetric(SegmentIdWithShardSpec allocatedId, Task task, ServiceEmitter emitter)
+  {
+    final ServiceMetricEvent.Builder metricBuilder = new ServiceMetricEvent.Builder();
+    IndexTaskUtils.setTaskDimensions(metricBuilder, task);
+    if (task instanceof SeekableStreamIndexTask<?, ?, ?>) {
+      metricBuilder.setDimension(
+          DruidMetrics.SUPERVISOR_ID,
+          ((SeekableStreamIndexTask<?, ?, ?>) task).getSupervisorId()
+      );
+    }
+    metricBuilder.setDimension(DruidMetrics.ID, allocatedId.toString());
+    emitter.emit(metricBuilder.setMetric("segment/allocated/count", 1));
+  }
+
   /**
    * Gets total row count of the given segments. Legacy segments do not have the
    * row count populated in the metadata and thus do not contribute to the row

indexing-service/src/main/java/org/apache/druid/indexing/scheduledbatch/ScheduledBatchTaskManager.java

@@ -308,11 +308,19 @@ private synchronized void stopScheduling()
       statusTracker.cleanupStaleTaskStatuses(supervisorId);
     }
 
+    /**
+     * Emits a metric with all the dimensions applicable to this supervisor.
+     * The {@link #supervisorId} is added to both {@link DruidMetrics#SUPERVISOR_ID}
+     * and {@link DruidMetrics#ID} dimensions for backward compatibility.
+     * {@link DruidMetrics#ID} is deprecated because it is ambiguous and will be
+     * removed in a future release.
+     */
     private void emitMetric(final String metricName, final int value)
     {
       emitter.emit(
           ServiceMetricEvent.builder()
                             .setDimension(DruidMetrics.ID, supervisorId)
+                            .setDimension(DruidMetrics.SUPERVISOR_ID, supervisorId)
                             .setDimension(DruidMetrics.DATASOURCE, dataSource)
                             .setMetric(metricName, value)
       );

indexing-service/src/test/java/org/apache/druid/indexing/common/task/concurrent/ConcurrentReplaceAndAppendTest.java

@@ -40,6 +40,7 @@
 import org.apache.druid.indexing.common.task.NoopTaskContextEnricher;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.common.task.TestAppenderatorsManager;
+import org.apache.druid.indexing.overlord.SegmentPublishResult;
 import org.apache.druid.indexing.overlord.Segments;
 import org.apache.druid.indexing.overlord.TaskQueue;
 import org.apache.druid.indexing.overlord.TaskRunner;
@@ -54,6 +55,7 @@
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.metadata.PendingSegmentRecord;
 import org.apache.druid.segment.IndexIO;
 import org.apache.druid.segment.TestDataSource;
 import org.apache.druid.segment.column.ColumnConfig;
@@ -83,6 +85,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 /**
  * Contains tests to verify behaviour of concurrently running REPLACE and APPEND
@@ -977,16 +980,29 @@ public void testSegmentIsAllocatedAtLatestVersion()
 
     final String v1 = replaceTask.acquireReplaceLockOn(JAN_23).getVersion();
     final DataSegment segmentV10 = createSegment(JAN_23, v1);
-    replaceTask.commitReplaceSegments(segmentV10);
+    final SegmentPublishResult replacePublishResult = replaceTask.commitReplaceSegments(segmentV10);
+
+    final List<PendingSegmentRecord> upgradedPendingSegments = replacePublishResult.getUpgradedPendingSegments();
+    Assert.assertNotNull(upgradedPendingSegments);
+    final SegmentIdWithShardSpec pendingSegmentV11 = upgradedPendingSegments.getFirst().getId();
+
     verifyIntervalHasUsedSegments(JAN_23, segmentV10);
     verifyIntervalHasVisibleSegments(JAN_23, segmentV10);
 
+    // New pending segment is allocated at v1
     final SegmentIdWithShardSpec pendingSegmentV12
         = appendTask.allocateSegmentForTimestamp(JAN_23.getStart(), Granularities.MONTH);
     Assert.assertNotEquals(pendingSegmentV01.asSegmentId(), pendingSegmentV12.asSegmentId());
+    Assert.assertNotEquals(pendingSegmentV11.asSegmentId(), pendingSegmentV12.asSegmentId());
     Assert.assertEquals(v1, pendingSegmentV12.getVersion());
     Assert.assertEquals(JAN_23, pendingSegmentV12.getInterval());
 
+    // Verify that no new segment has been allocated at v0
+    verifyIntervalHasPendingSegments(
+        JAN_23,
+        pendingSegmentV01, pendingSegmentV11, pendingSegmentV12
+    );
+
     replaceTask.releaseLock(JAN_23);
     final ActionsTestTask replaceTask2 = createAndStartTask();
     final String v2 = replaceTask2.acquireReplaceLockOn(JAN_23).getVersion();
@@ -1205,6 +1221,37 @@ public void test_allocateCommitDelete_createsFreshVersion_uptoMaxAllowedRetries(
     );
   }
 
+  @Test
+  public void test_concurrentReplace_onIntervalWithPendingSegment_upgradesIt()
+  {
+    // Allocate a segment on an empty interval
+    final SegmentIdWithShardSpec pendingSegmentV01
+        = appendTask.allocateSegmentForTimestamp(JAN_23.getStart(), Granularities.MONTH);
+    Assert.assertEquals(SEGMENT_V0, pendingSegmentV01.getVersion());
+    Assert.assertEquals(JAN_23, pendingSegmentV01.getInterval());
+
+    // Replace the segments in the interval
+    final String v1 = replaceTask.acquireReplaceLockOn(JAN_23).getVersion();
+    final DataSegment segmentV10 = createSegment(JAN_23, v1);
+    final SegmentPublishResult replacePublishResult = replaceTask.commitReplaceSegments(segmentV10);
+
+    // Verify that pendingSegmentV01 has been upgraded to version v1
+    final List<PendingSegmentRecord> upgradedPendingSegments = replacePublishResult.getUpgradedPendingSegments();
+    Assert.assertNotNull(upgradedPendingSegments);
+    Assert.assertEquals(1, upgradedPendingSegments.size());
+    PendingSegmentRecord upgradedPendingSegment = upgradedPendingSegments.getFirst();
+    Assert.assertEquals(
+        pendingSegmentV01.asSegmentId().toString(),
+        upgradedPendingSegment.getUpgradedFromSegmentId()
+    );
+    Assert.assertEquals(v1, upgradedPendingSegment.getId().getVersion());
+
+    verifyIntervalHasPendingSegments(
+        JAN_23,
+        pendingSegmentV01, upgradedPendingSegment.getId()
+    );
+  }
+
   @Nullable
   private DataSegment findSegmentWith(String version, Map<String, Object> loadSpec, Set<DataSegment> segments)
   {
@@ -1227,6 +1274,17 @@ private static DataSegment asSegment(SegmentIdWithShardSpec pendingSegment)
                       .build();
   }
 
+  private void verifyIntervalHasPendingSegments(Interval interval, SegmentIdWithShardSpec... expectedPendingSegments)
+  {
+    final Set<SegmentIdWithShardSpec> expected = Set.of(expectedPendingSegments);
+    final Set<SegmentIdWithShardSpec> observed = getStorageCoordinator()
+        .getPendingSegments(TestDataSource.WIKI, interval)
+        .stream()
+        .map(PendingSegmentRecord::getId)
+        .collect(Collectors.toSet());
+    Assert.assertEquals(expected, observed);
+  }
+
   private void verifyIntervalHasUsedSegments(Interval interval, DataSegment... expectedSegments)
   {
     verifySegments(interval, Segments.INCLUDING_OVERSHADOWED, expectedSegments);
@@ -1240,7 +1298,6 @@ private void verifyIntervalHasVisibleSegments(Interval interval, DataSegment...
   private void verifySegments(Interval interval, Segments visibility, DataSegment... expectedSegments)
   {
     try {
-
       Collection<DataSegment> allUsedSegments = dummyTaskActionClient.submit(
           new RetrieveUsedSegmentsAction(
               TestDataSource.WIKI,