Skip to content

Fix flaky unit tests in SegmentBootstrapperTest and KinesisIndexTaskTest #17841

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kfaraz merged 6 commits into from
Mar 29, 2025
Merged

Fix flaky unit tests in SegmentBootstrapperTest and KinesisIndexTaskTest #17841

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0231b52
Fix flakiness inSegmentBootstrapperTest
capistrant Mar 25, 2025
896cc22
Fix flaky KinesisIndexTask tests.
capistrant Mar 25, 2025
07c24d0
Fix one additional flake from SegmentBootstrapperTest found after add…
capistrant Mar 27, 2025
ab9162f
refactor asserts in SegmentBootstrapperTest after code review suggest…
capistrant Mar 27, 2025
45ee39a
Simplify the KinesisIndexTaskTests who wait on given task state
capistrant Mar 27, 2025
7a1a6b6
Complete javadoc for new method
capistrant Mar 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions ...ndexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -793,7 +793,7 @@ public void testRunWithMinimumMessageTime() throws Exception

final ListenableFuture<TaskStatus> future = runTask(task);

waitUntil(task, this::isTaskReading);
waitUntil(task, this::isTaskPublishing);

// Wait for task to exit
Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());
Expand Down Expand Up @@ -856,7 +856,7 @@ public void testRunWithMaximumMessageTime() throws Exception

final ListenableFuture<TaskStatus> future = runTask(task);

waitUntil(task, this::isTaskReading);
waitUntil(task, this::isTaskPublishing);

// Wait for task to exit
Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());
Expand Down Expand Up @@ -915,7 +915,7 @@ public void testRunWithTransformSpec() throws Exception
);

final ListenableFuture<TaskStatus> future = runTask(task);
waitUntil(task, this::isTaskReading);
waitUntil(task, this::isTaskPublishing);

// Wait for task to exit
Assert.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());
Expand Down Expand Up @@ -2461,6 +2461,16 @@ private boolean isTaskReading(KinesisIndexTask task)
return task.getRunner().getStatus() == SeekableStreamIndexTaskRunner.Status.READING;
}

/**
* Return true if specified task is in PUBLISHING state
* @param task {@link KinesisIndexTask} having its state checked
* @return true if task is in PUBLISHING state, otherwise false
*/
private boolean isTaskPublishing(KinesisIndexTask task)
{
return task.getRunner().getStatus() == SeekableStreamIndexTaskRunner.Status.PUBLISHING;
}

private static KinesisRecordEntity kjb(
String timestamp,
String dim1,
Expand Down
42 changes: 28 additions & 14 deletions server/src/test/java/org/apache/druid/server/coordination/SegmentBootstrapperTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,11 +150,12 @@ public void testStartStop() throws Exception
Assert.assertEquals(2L, segmentManager.getDataSourceCounts().get("test_two" + i).longValue());
}

Assert.assertEquals(ImmutableList.copyOf(segments), segmentAnnouncer.getObservedSegments());

final ImmutableList<DataSegment> expectedBootstrapSegments = ImmutableList.copyOf(segments);
Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());

assertUnsortedListsAreEqual(expectedBootstrapSegments, segmentAnnouncer.getObservedSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());

Assert.assertEquals(ImmutableList.of(), cacheManager.getObservedSegments());
Assert.assertEquals(ImmutableList.of(), cacheManager.getObservedSegmentsLoadedIntoPageCache());

Expand Down Expand Up @@ -211,11 +212,12 @@ public void testLoadCachedSegments() throws Exception
Assert.assertEquals(2L, segmentManager.getDataSourceCounts().get("test_two" + i).longValue());
}

Assert.assertEquals(ImmutableList.copyOf(segments), segmentAnnouncer.getObservedSegments());

final ImmutableList<DataSegment> expectedBootstrapSegments = ImmutableList.copyOf(segments);
Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());

assertUnsortedListsAreEqual(expectedBootstrapSegments, segmentAnnouncer.getObservedSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());

Assert.assertEquals(ImmutableList.of(), cacheManager.getObservedSegments());
Assert.assertEquals(ImmutableList.of(), cacheManager.getObservedSegmentsLoadedIntoPageCache());

Expand Down Expand Up @@ -270,10 +272,10 @@ public void testLoadBootstrapSegments() throws Exception

final ImmutableList<DataSegment> expectedBootstrapSegments = ImmutableList.copyOf(segments);

Assert.assertEquals(expectedBootstrapSegments, segmentAnnouncer.getObservedSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, segmentAnnouncer.getObservedSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());

Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());
serviceEmitter.verifyValue("segment/bootstrap/count", expectedBootstrapSegments.size());
serviceEmitter.verifyEmitted("segment/bootstrap/time", 1);

Expand Down Expand Up @@ -393,10 +395,10 @@ public void testLoadOnlyRequiredBootstrapSegments() throws Exception

final ImmutableList<DataSegment> expectedBootstrapSegments = ImmutableList.of(ds1Segment2, ds1Segment1);

Assert.assertEquals(expectedBootstrapSegments, segmentAnnouncer.getObservedSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, segmentAnnouncer.getObservedSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
assertUnsortedListsAreEqual(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());

Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegments());
Assert.assertEquals(expectedBootstrapSegments, cacheManager.getObservedBootstrapSegmentsLoadedIntoPageCache());
serviceEmitter.verifyValue("segment/bootstrap/count", expectedBootstrapSegments.size());
serviceEmitter.verifyEmitted("segment/bootstrap/time", 1);

Expand Down Expand Up @@ -440,4 +442,16 @@ public void testLoadBootstrapSegmentsWhenExceptionThrown() throws Exception

bootstrapper.stop();
}

/**
* Given two lists, assert they are equivalent and contain the same set of entries irrespecive of entry ordering
* @param expected The expected result list
* @param actual The actual result list
* @param <T> Object type stored in the list parameters
*/
private static <T> void assertUnsortedListsAreEqual(List<T> expected, List<T> actual)
{
Assert.assertEquals(expected.size(), actual.size());
Assert.assertEquals(Set.copyOf(expected), Set.copyOf(actual));
}
}
16 changes: 10 additions & 6 deletions server/src/test/java/org/apache/druid/server/coordination/TestSegmentCacheManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,9 @@
import org.apache.druid.timeline.DataSegment;
import org.joda.time.Interval;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;

/**
Expand All @@ -59,11 +59,15 @@ class TestSegmentCacheManager extends NoopSegmentCacheManager
TestSegmentCacheManager(final Set<DataSegment> segmentsToCache)
{
this.cachedSegments = ImmutableList.copyOf(segmentsToCache);
this.observedBootstrapSegments = new ArrayList<>();
this.observedBootstrapSegmentsLoadedIntoPageCache = new ArrayList<>();
this.observedSegments = new ArrayList<>();
this.observedSegmentsLoadedIntoPageCache = new ArrayList<>();
this.observedSegmentsRemovedFromCache = new ArrayList<>();

// While inneficient, these CopyOnWriteArrayList objects greatly simplify meeting the thread
// safety mandate from SegmentCacheManager. For testing, this should be ok.
this.observedBootstrapSegments = new CopyOnWriteArrayList<>();
this.observedBootstrapSegmentsLoadedIntoPageCache = new CopyOnWriteArrayList<>();
this.observedSegments = new CopyOnWriteArrayList<>();
this.observedSegmentsLoadedIntoPageCache = new CopyOnWriteArrayList<>();
this.observedSegmentsRemovedFromCache = new CopyOnWriteArrayList<>();

this.observedShutdownBootstrapCount = new AtomicInteger(0);
}

Expand Down
Loading