test: poll sys.segments in verifyNumVisibleSegmentsIs to fix flaky compaction tests

embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/CompactionTestBase.java

@@ -97,8 +97,8 @@ protected List<Interval> 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);
   }
 

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<Interval> getSortedSegmentIntervals(String dataSource, EmbeddedOverl
    */
   public void verifyNumVisibleSegmentsIs(int numExpectedSegments, String dataSource, EmbeddedOverlord overlord)
   {
-    int segmentCount = getVisibleUsedSegments(dataSource, overlord).size();
+    final Set<DataSegment> 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
+      );
+    }
   }
 
   /**