Skip to content

Fixed flaky file-source unit tests #6984 #6995

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
dlvenable merged 1 commit into from
Jul 21, 2026
Merged

Fixed flaky file-source unit tests #6984 #6995

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
30 changes: 21 additions & 9 deletions .../file-source/src/main/java/org/opensearch/dataprepper/plugins/source/file/FileReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.time.Clock;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
Expand Down Expand Up @@ -82,6 +83,7 @@ public final class FileReader implements Runnable {
private final InputCodec codec;
private final boolean tailMode;
private final DecompressionEngine decompressionEngine;
private final Clock clock;

private final AtomicLong readOffset;
private final StringBuilder partialLine;
Expand All @@ -102,11 +104,21 @@ public FileReader(final Path path,
final CheckpointEntry checkpointEntry,
final FileReaderContext context,
final Runnable onComplete) {
this(path, fileIdentity, checkpointEntry, context, onComplete, Clock.systemDefaultZone());
}

FileReader(final Path path,
final FileIdentity fileIdentity,
final CheckpointEntry checkpointEntry,
final FileReaderContext context,
final Runnable onComplete,
final Clock clock) {
this.path = Objects.requireNonNull(path, "path must not be null");
this.fileIdentity = Objects.requireNonNull(fileIdentity, "fileIdentity must not be null");
this.checkpointEntry = Objects.requireNonNull(checkpointEntry, "checkpointEntry must not be null");
Objects.requireNonNull(context, "context must not be null");
this.onComplete = Objects.requireNonNull(onComplete, "onComplete must not be null");
this.clock = Objects.requireNonNull(clock, "clock must not be null");

this.buffer = context.getBuffer();
this.eventFactory = context.getEventFactory();
Expand Down Expand Up @@ -145,9 +157,9 @@ public FileReader(final Path path,
this.cachedAbsolutePath = path.toAbsolutePath().toString();
this.currentBatchCount = 0;
this.batchStartOffset = readOffset.get();
this.batchOpenedAtMillis = System.currentTimeMillis();
this.batchOpenedAtMillis = clock.millis();
this.lastRotationType = RotationType.NO_ROTATION;
this.lastActivityMillis = System.currentTimeMillis();
this.lastActivityMillis = clock.millis();
}

@Override
Expand Down Expand Up @@ -244,7 +256,7 @@ private void readFile() {

private void readLoop(final FileChannel channel, final long timeoutMillis, final boolean isDraining) throws IOException {
final ByteBuffer byteBuffer = ByteBuffer.allocate(readBufferSize);
final long loopStart = System.currentTimeMillis();
final long loopStart = clock.millis();
final ByteArrayOutputStream codecAccumulator = codec != null ? new ByteArrayOutputStream() : null;
long codecBytesAccumulated = 0;
final CharsetDecoder decoder = codec == null ? encoding.newDecoder()
Expand All @@ -254,7 +266,7 @@ private void readLoop(final FileChannel channel, final long timeoutMillis, final
final ByteBuffer decoderCarryover = codec == null ? ByteBuffer.allocate(8) : null;

while (!Thread.currentThread().isInterrupted()) {
final long elapsed = System.currentTimeMillis() - loopStart;
final long elapsed = clock.millis() - loopStart;
if (elapsed >= timeoutMillis) {
if (isDraining) {
long currentFileSize = 0;
Expand Down Expand Up @@ -305,10 +317,10 @@ private void readLoop(final FileChannel channel, final long timeoutMillis, final
}

readOffset.addAndGet(Math.max(0, bytesRead));
lastActivityMillis = System.currentTimeMillis();
lastActivityMillis = clock.millis();

if (acknowledgementsEnabled && currentAckSet != null && currentBatchCount > 0) {
final long batchAge = System.currentTimeMillis() - batchOpenedAtMillis;
final long batchAge = clock.millis() - batchOpenedAtMillis;
if (batchAge >= batchTimeout.toMillis()) {
completePendingAckSet();
}
Expand Down Expand Up @@ -421,9 +433,9 @@ private void writeRecordWithRetry(final Record<Object> record, final Event event
long backpressureStartNanos = 0;
boolean backpressureActive = false;
final long maxRetryMillis = maxReadTimePerFile.toMillis();
final long retryStart = System.currentTimeMillis();
final long retryStart = clock.millis();
while (!written && !Thread.currentThread().isInterrupted()) {
if (System.currentTimeMillis() - retryStart > maxRetryMillis) {
if (clock.millis() - retryStart > maxRetryMillis) {
LOG.warn("Backpressure retry timeout exceeded for file {}. Event may be lost.", path);
metrics.getDataLossEvents().increment();
break;
Expand Down Expand Up @@ -468,7 +480,7 @@ private void ensureAckSet() {
if (currentAckSet == null) {
final long capturedBatchStart = readOffset.get();
batchStartOffset = capturedBatchStart;
batchOpenedAtMillis = System.currentTimeMillis();
batchOpenedAtMillis = clock.millis();
currentAckSet = acknowledgementSetManager.create(
result -> handleAcknowledgement(result, capturedBatchStart, currentBatchEndOffset),
acknowledgmentTimeout);
Expand Down
25 changes: 17 additions & 8 deletions ...urce/src/test/java/org/opensearch/dataprepper/plugins/source/file/FileReaderPoolTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,9 +103,18 @@ private FileReaderPool createPoolWithoutMetrics(final int maxActiveFiles, final
Duration.ofMinutes(30), createReaderContext());
}

private FileReaderPool createPoolWithNonRunningExecutor(final int maxActiveFiles) {
when(metrics.getActiveFileCount()).thenReturn(new AtomicLong(0));
final ExecutorService nonRunningExecutor = mock(ExecutorService.class);
return new FileReaderPool(
checkpointRegistry, metrics, maxActiveFiles,
Duration.ofMinutes(30), createReaderContext(),
() -> nonRunningExecutor);
}

@Test
void addFile_submits_reader_when_under_max_active_files() {
FileReaderPool pool = createPool(10, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(10);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity = mock(FileIdentity.class);
Expand All @@ -120,7 +129,7 @@ void addFile_submits_reader_when_under_max_active_files() {

@Test
void addFile_is_idempotent_for_same_identity() {
FileReaderPool pool = createPool(10, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(10);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity = mock(FileIdentity.class);
Expand All @@ -135,7 +144,7 @@ void addFile_is_idempotent_for_same_identity() {

@Test
void addFile_queues_pending_when_at_max_active_files() {
FileReaderPool pool = createPool(1, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(1);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity1 = mock(FileIdentity.class);
Expand All @@ -151,7 +160,7 @@ void addFile_queues_pending_when_at_max_active_files() {

@Test
void addFile_does_not_add_pending_duplicate_to_queue() {
FileReaderPool pool = createPool(1, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(1);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity1 = mock(FileIdentity.class);
Expand All @@ -167,7 +176,7 @@ void addFile_does_not_add_pending_duplicate_to_queue() {

@Test
void addFile_queues_multiple_pending_files() {
FileReaderPool pool = createPool(1, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(1);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity1 = mock(FileIdentity.class);
Expand Down Expand Up @@ -220,7 +229,7 @@ void closeInactiveReaders_removes_inactive_readers() throws Exception {
void closeReaderForPath_removes_matching_reader() {
Counter filesClosed = mock(Counter.class);
lenient().when(metrics.getFilesClosed()).thenReturn(filesClosed);
FileReaderPool pool = createPool(10, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(10);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity = mock(FileIdentity.class);
Expand All @@ -236,7 +245,7 @@ void closeReaderForPath_removes_matching_reader() {

@Test
void closeReaderForPath_does_nothing_when_no_match() {
FileReaderPool pool = createPool(10, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(10);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity = mock(FileIdentity.class);
Expand All @@ -251,7 +260,7 @@ void closeReaderForPath_does_nothing_when_no_match() {
void closeReaderForPath_promotes_pending_files() {
Counter filesClosed = mock(Counter.class);
lenient().when(metrics.getFilesClosed()).thenReturn(filesClosed);
FileReaderPool pool = createPool(1, 2);
FileReaderPool pool = createPoolWithNonRunningExecutor(1);
when(checkpointRegistry.getOrCreate(anyString())).thenReturn(new CheckpointEntry());

FileIdentity identity1 = mock(FileIdentity.class);
Expand Down
41 changes: 30 additions & 11 deletions ...e-source/src/test/java/org/opensearch/dataprepper/plugins/source/file/FileReaderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.time.Clock;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.TimeUnit;
Expand All @@ -50,6 +51,7 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyDouble;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.eq;
Expand Down Expand Up @@ -630,14 +632,18 @@ void run_drain_timeout_logs_data_loss_when_unread_data() throws Exception {
when(rotationDetector.checkRotation(any(), any(), any(long.class)))
.thenReturn(new RotationResult(RotationType.CREATE_RENAME, newIdentity));

final AtomicLong fakeNowMillis = new AtomicLong(0);
final Clock clock = mock(Clock.class);
when(clock.millis()).thenAnswer(inv -> fakeNowMillis.get());

FileChannel mockChannel = mock(FileChannel.class);
when(fileOps.openReadChannel(testFile)).thenReturn(mockChannel);
when(mockChannel.position(anyLong())).thenReturn(mockChannel);
lenient().when(mockChannel.read(any(ByteBuffer.class))).thenAnswer(inv -> {
Thread.sleep(5);
when(mockChannel.read(any(ByteBuffer.class))).thenAnswer(inv -> {
ByteBuffer buf = inv.getArgument(0);
byte[] data = "A".repeat(buf.remaining()).getBytes();
buf.put(data, 0, Math.min(data.length, buf.remaining()));
fakeNowMillis.set(1000);
return buf.position();
});
when(mockChannel.size()).thenReturn(100000L);
Expand All @@ -649,14 +655,14 @@ void run_drain_timeout_logs_data_loss_when_unread_data() throws Exception {
when(metrics.getFilesRotated()).thenReturn(filesRotated);
when(metrics.getFilesClosed()).thenReturn(filesClosed);
when(metrics.getFilesOpened()).thenReturn(mock(Counter.class));
lenient().when(metrics.getBytesRead()).thenReturn(bytesRead);
when(metrics.getBytesRead()).thenReturn(bytesRead);
when(metrics.getDataLossEvents()).thenReturn(dataLossEvents);

FileReaderContext context = new FileReaderContext(
buffer, eventFactory, fileOps, metrics, rotationDetector,
acknowledgementSetManager, false, StandardCharsets.UTF_8,
4096, 1048576, 5000, Duration.ofSeconds(30),
Duration.ofMillis(1), StartPosition.BEGINNING, false,
Duration.ofMillis(100), StartPosition.BEGINNING, false,
Duration.ofSeconds(30), 1000,
Duration.ofSeconds(5), 3, null, true, null);

Expand All @@ -668,7 +674,7 @@ void run_drain_timeout_logs_data_loss_when_unread_data() throws Exception {

fileIdentity = mock(FileIdentity.class);
final FileReader reader = new FileReader(testFile, fileIdentity, checkpointEntry, context,
() -> onCompleteCalled.set(true));
() -> onCompleteCalled.set(true), clock);
reader.run();

verify(dataLossEvents).increment();
Expand All @@ -682,14 +688,18 @@ void run_max_read_time_reached_breaks_without_data_loss() throws Exception {
when(rotationDetector.checkRotation(any(), any(), any(long.class)))
.thenReturn(RotationResult.NO_ROTATION);

final AtomicLong fakeNowMillis = new AtomicLong(0);
final Clock clock = mock(Clock.class);
when(clock.millis()).thenAnswer(inv -> fakeNowMillis.get());

FileChannel mockChannel = mock(FileChannel.class);
when(fileOps.openReadChannel(testFile)).thenReturn(mockChannel);
when(mockChannel.position(anyLong())).thenReturn(mockChannel);
when(mockChannel.read(any(ByteBuffer.class))).thenAnswer(inv -> {
Thread.sleep(5);
ByteBuffer buf = inv.getArgument(0);
byte[] data = "A".repeat(buf.remaining()).getBytes();
buf.put(data, 0, Math.min(data.length, buf.remaining()));
fakeNowMillis.set(1000);
return buf.position();
});

Expand All @@ -709,16 +719,18 @@ void run_max_read_time_reached_breaks_without_data_loss() throws Exception {
FileReaderContext context = new FileReaderContext(
buffer, eventFactory, fileOps, metrics, rotationDetector,
acknowledgementSetManager, false, StandardCharsets.UTF_8,
4096, 1048576, 5000, Duration.ofMillis(1),
4096, 1048576, 5000, Duration.ofMillis(100),
Duration.ofSeconds(30), StartPosition.BEGINNING, false,
Duration.ofSeconds(30), 1000,
Duration.ofSeconds(5), 3, null, true, null);

fileIdentity = mock(FileIdentity.class);
final FileReader reader = new FileReader(testFile, fileIdentity, checkpointEntry, context,
() -> onCompleteCalled.set(true));
() -> onCompleteCalled.set(true), clock);
reader.run();

verify(mockChannel, times(1)).read(any(ByteBuffer.class));
verify(bytesRead).increment(anyDouble());
assertThat(onCompleteCalled.get(), equalTo(true));
}

Expand Down Expand Up @@ -1251,14 +1263,18 @@ void run_drain_timeout_handles_io_exception_on_channel_size() throws Exception {
when(rotationDetector.checkRotation(any(), any(), any(long.class)))
.thenReturn(new RotationResult(RotationType.CREATE_RENAME, newIdentity));

final AtomicLong fakeNowMillis = new AtomicLong(0);
final Clock clock = mock(Clock.class);
when(clock.millis()).thenAnswer(inv -> fakeNowMillis.get());

FileChannel mockChannel = mock(FileChannel.class);
when(fileOps.openReadChannel(testFile)).thenReturn(mockChannel);
when(mockChannel.position(anyLong())).thenReturn(mockChannel);
when(mockChannel.read(any(ByteBuffer.class))).thenAnswer(inv -> {
Thread.sleep(5);
ByteBuffer buf = inv.getArgument(0);
byte[] data = "A".repeat(buf.remaining()).getBytes();
buf.put(data, 0, Math.min(data.length, buf.remaining()));
fakeNowMillis.set(1000);
return buf.position();
});
when(mockChannel.size()).thenThrow(new IOException("channel closed"));
Expand All @@ -1277,7 +1293,7 @@ void run_drain_timeout_handles_io_exception_on_channel_size() throws Exception {
buffer, eventFactory, fileOps, metrics, rotationDetector,
acknowledgementSetManager, false, StandardCharsets.UTF_8,
4096, 1048576, 5000, Duration.ofSeconds(30),
Duration.ofMillis(1), StartPosition.BEGINNING, false,
Duration.ofMillis(100), StartPosition.BEGINNING, false,
Duration.ofSeconds(30), 1000,
Duration.ofSeconds(5), 3, null, true, null);

Expand All @@ -1289,8 +1305,11 @@ void run_drain_timeout_handles_io_exception_on_channel_size() throws Exception {

fileIdentity = mock(FileIdentity.class);
final FileReader reader = new FileReader(testFile, fileIdentity, checkpointEntry, context,
() -> onCompleteCalled.set(true));
() -> onCompleteCalled.set(true), clock);
reader.run();

verify(mockChannel, times(1)).read(any(ByteBuffer.class));
verify(mockChannel).size();
}

@Test
Expand Down
Loading