minor: Add segment/allocated/count metric

docs/operations/metrics.md

@@ -345,6 +345,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`|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|

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

@@ -26,17 +26,21 @@
 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;
 import org.apache.druid.indexing.overlord.LockRequestForNewSegment;
 import org.apache.druid.indexing.overlord.LockResult;
 import org.apache.druid.indexing.overlord.Segments;
+import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTask;
 import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.partition.NumberedPartialShardSpec;
@@ -241,6 +245,7 @@ public SegmentIdWithShardSpec perform(
         identifier = tryAllocateSubsequentSegment(toolbox, task, rowInterval, usedSegmentsForRow.iterator().next());
       }
       if (identifier != null) {
+        emitSuccessMetric(identifier, task, toolbox);
         return identifier;
       }
 
@@ -280,14 +285,28 @@ public SegmentIdWithShardSpec perform(
     }
   }
 
+  private void emitSuccessMetric(SegmentIdWithShardSpec allocatedId, Task task, TaskActionToolbox toolbox)
+  {
+    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());
+    toolbox.getEmitter().emit(metricBuilder.setMetric("segment/allocated/count", 1));
+  }
+
   private SegmentIdWithShardSpec tryAllocateFirstSegment(TaskActionToolbox toolbox, Task task, Interval rowInterval)
   {
     // No existing segments for this row, but there might still be nearby ones that conflict with our preferred
     // segment granularity. Try that first, and then progressively smaller ones if it fails.
     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/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,