diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/CompactionTestBase.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/CompactionTestBase.java index b714cd00859e..aefb9770bab4 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/CompactionTestBase.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/CompactionTestBase.java @@ -97,8 +97,8 @@ protected List getSegmentIntervals() protected void verifySegmentsCount(int numExpectedSegments) { - // Ensure that Broker has synced latest segments from the Coordinator - broker.latchableEmitter().waitForNextEvent(event -> event.hasMetricName("segment/metadataCache/sync/time")); + // The Overlord state is verified immediately; the Broker-side sys.segments view is polled + // until it matches, since a single metadata cache sync may still observe the previous state. cluster.callApi().verifyNumVisibleSegmentsIs(numExpectedSegments, dataSource, overlord); } diff --git a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java index 6ae8750b8d8e..091c86c1da9c 100644 --- a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java +++ b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java @@ -61,6 +61,7 @@ import java.util.TreeSet; import java.util.function.Function; import java.util.function.Predicate; +import java.util.stream.Collectors; /** * Contains various utility methods to interact with an {@link EmbeddedDruidCluster}. @@ -281,21 +282,43 @@ public List getSortedSegmentIntervals(String dataSource, EmbeddedOverl */ public void verifyNumVisibleSegmentsIs(int numExpectedSegments, String dataSource, EmbeddedOverlord overlord) { - int segmentCount = getVisibleUsedSegments(dataSource, overlord).size(); + final Set visibleSegments = getVisibleUsedSegments(dataSource, overlord); + final int segmentCount = visibleSegments.size(); Assertions.assertEquals( numExpectedSegments, segmentCount, "Segment count mismatch" ); - Assertions.assertEquals( - String.valueOf(segmentCount), - runSql( - "SELECT COUNT(*) FROM sys.segments WHERE datasource='%s'" - + " AND is_overshadowed = 0 AND is_available = 1", - dataSource - ), - "Segment count mismatch in sys.segments table" - ); + + // The Broker learns about segment changes asynchronously from the Coordinator, so the + // sys.segments table may briefly lag behind the metadata store. Match the segment IDs as well + // as the count because compaction can replace segments without changing their number. + final String expectedCount = String.valueOf(segmentCount); + final String expectedSegmentIds = visibleSegments + .stream() + .map(segment -> StringUtils.format("'%s'", StringUtils.escapeSql(segment.getId().toString()))) + .collect(Collectors.joining(", ")); + final String segmentIdFilter = expectedSegmentIds.isEmpty() + ? "" + : " AND segment_id IN (" + expectedSegmentIds + ")"; + final String sql = "SELECT COUNT(*) FROM sys.segments WHERE datasource='%s'" + + " AND is_overshadowed = 0 AND is_available = 1%s"; + try { + waitForResult(() -> runSql(sql, StringUtils.escapeSql(dataSource), segmentIdFilter), expectedCount::equals) + .withTimeoutMillis(60_000) + .go(); + } + catch (ISE e) { + throw new AssertionError( + StringUtils.format( + "Segment count mismatch in sys.segments table: expected[%s] segments for datasource[%s]. %s", + expectedCount, + dataSource, + e.getMessage() + ), + e + ); + } } /**