Skip to content

[FLINK-32446][Connectors/MongoDB] MongoWriter should regularly check whether the last write time is over limit #12

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
merged 2 commits into from
Jun 28, 2023
Merged

[FLINK-32446][Connectors/MongoDB] MongoWriter should regularly check whether the last write time is over limit #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
471c180
[FLINK-32446][Connectors/MongoDB] MongoWriter should regularly check …
Jiabao-Sun Jun 27, 2023
d1454b3
Address Leonard's comments
Jiabao-Sun Jun 28, 2023
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
83 changes: 75 additions & 8 deletions ...tor-mongodb/src/main/java/org/apache/flink/connector/mongodb/sink/writer/MongoWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
import org.apache.flink.util.Collector;
import org.apache.flink.util.FlinkRuntimeException;
import org.apache.flink.util.concurrent.ExecutorThreadFactory;

import com.mongodb.MongoException;
import com.mongodb.client.MongoClient;
Expand All @@ -45,6 +46,10 @@
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

import static org.apache.flink.util.Preconditions.checkNotNull;

Expand All @@ -68,11 +73,18 @@ public class MongoWriter<IN> implements SinkWriter<IN> {
private final Collector<WriteModel<BsonDocument>> collector;
private final Counter numRecordsOut;
private final MongoClient mongoClient;
private final long batchIntervalMs;
private final int batchSize;

private boolean checkpointInProgress = false;
private volatile long lastSendTime = 0L;
private volatile long ackTime = Long.MAX_VALUE;

private transient volatile boolean closed = false;
private transient ScheduledExecutorService scheduler;
private transient ScheduledFuture<?> scheduledFuture;
private transient volatile Exception flushException;

public MongoWriter(
MongoConnectionOptions connectionOptions,
MongoWriteOptions writeOptions,
Expand All @@ -83,6 +95,8 @@ public MongoWriter(
this.writeOptions = checkNotNull(writeOptions);
this.serializationSchema = checkNotNull(serializationSchema);
this.flushOnCheckpoint = flushOnCheckpoint;
this.batchIntervalMs = writeOptions.getBatchIntervalMs();
this.batchSize = writeOptions.getBatchSize();

checkNotNull(initContext);
this.mailboxExecutor = checkNotNull(initContext.getMailboxExecutor());
Expand All @@ -105,10 +119,37 @@ public MongoWriter(

// Initialize the mongo client.
this.mongoClient = MongoClients.create(connectionOptions.getUri());

boolean flushOnlyOnCheckpoint = batchIntervalMs == -1 && batchSize == -1;

if (!flushOnlyOnCheckpoint && batchIntervalMs > 0) {
this.scheduler =
Executors.newScheduledThreadPool(1, new ExecutorThreadFactory("mongo-writer"));

this.scheduledFuture =
this.scheduler.scheduleWithFixedDelay(
() -> {
synchronized (MongoWriter.this) {
if (!closed && isOverMaxBatchIntervalLimit()) {
try {
doBulkWrite();
} catch (Exception e) {
flushException = e;
}
}
}
},
batchIntervalMs,
batchIntervalMs,
TimeUnit.MILLISECONDS);
}
}

@Override
public void write(IN element, Context context) throws IOException, InterruptedException {
public synchronized void write(IN element, Context context)
Copy link

@6harat 6harat Jun 28, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Jiabao-Sun just curious, wouldn't adding synchronized here hurt the performance for write operation?

Copy link
Contributor Author

@Jiabao-Sun Jiabao-Sun Jun 28, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @6harat with this comment.
Adding synchronized here because we may add an other thread to check whether the last write time is over limit.
The thread may call doBulkWrite method and clear the bulkRequests list which is write by write method introduces a race condition.

throws IOException, InterruptedException {
checkFlushException();

// do not allow new bulk writes until all actions are flushed
while (checkpointInProgress) {
mailboxExecutor.yield();
Expand All @@ -122,7 +163,9 @@ public void write(IN element, Context context) throws IOException, InterruptedEx
}

@Override
public void flush(boolean endOfInput) throws IOException {
public synchronized void flush(boolean endOfInput) throws IOException {
checkFlushException();

checkpointInProgress = true;
while (!bulkRequests.isEmpty() && (flushOnCheckpoint || endOfInput)) {
doBulkWrite();
Expand All @@ -131,8 +174,28 @@ public void flush(boolean endOfInput) throws IOException {
}

@Override
public void close() {
mongoClient.close();
public synchronized void close() throws Exception {
if (!closed) {
if (scheduledFuture != null) {
scheduledFuture.cancel(false);
scheduler.shutdown();
}

if (!bulkRequests.isEmpty()) {
try {
doBulkWrite();
} catch (Exception e) {
LOG.error("Writing records to MongoDB failed when closing MongoWriter", e);
throw new IOException("Writing records to MongoDB failed.", e);
} finally {
mongoClient.close();
closed = true;
}
} else {
mongoClient.close();
closed = true;
}
}
}

@VisibleForTesting
Expand Down Expand Up @@ -172,13 +235,17 @@ void doBulkWrite() throws IOException {
}

private boolean isOverMaxBatchSizeLimit() {
int bulkActions = writeOptions.getBatchSize();
return bulkActions != -1 && bulkRequests.size() >= bulkActions;
return batchSize != -1 && bulkRequests.size() >= batchSize;
}

private boolean isOverMaxBatchIntervalLimit() {
long bulkFlushInterval = writeOptions.getBatchIntervalMs();
long lastSentInterval = System.currentTimeMillis() - lastSendTime;
return bulkFlushInterval != -1 && lastSentInterval >= bulkFlushInterval;
return batchIntervalMs != -1 && lastSentInterval >= batchIntervalMs;
}

private void checkFlushException() {
if (flushException != null) {
throw new RuntimeException("Writing records to MongoDB failed.", flushException);
}
}
}
7 changes: 4 additions & 3 deletions ...ngodb/src/test/java/org/apache/flink/connector/mongodb/sink/writer/MongoWriterITCase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@

import static org.apache.flink.connector.mongodb.testutils.MongoTestUtil.assertThatIdsAreNotWritten;
import static org.apache.flink.connector.mongodb.testutils.MongoTestUtil.assertThatIdsAreWritten;
import static org.apache.flink.connector.mongodb.testutils.MongoTestUtil.assertThatIdsAreWrittenWithMaxWaitTime;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;

Expand Down Expand Up @@ -138,12 +139,12 @@ void testWriteOnBatchIntervalFlush() throws Exception {
createWriter(collection, batchSize, batchIntervalMs, flushOnCheckpoint)) {
writer.write(buildMessage(1), null);
writer.write(buildMessage(2), null);
writer.doBulkWrite();
writer.write(buildMessage(3), null);
writer.write(buildMessage(4), null);
writer.doBulkWrite();
}

assertThatIdsAreWritten(collectionOf(collection), 1, 2, 3, 4);
assertThatIdsAreWrittenWithMaxWaitTime(collectionOf(collection), 10000L, 1, 2, 3, 4);
}
}

@Test
Expand Down
13 changes: 13 additions & 0 deletions ...tor-mongodb/src/test/java/org/apache/flink/connector/mongodb/testutils/MongoTestUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,4 +77,17 @@ public static void assertThatIdsAreWritten(MongoCollection<Document> coll, Integ

assertThat(actualIds).containsExactlyInAnyOrder(ids);
}

public static void assertThatIdsAreWrittenWithMaxWaitTime(
MongoCollection<Document> coll, long maxWaitTimeMs, Integer... ids)
throws InterruptedException {
long startTimeMillis = System.currentTimeMillis();
while (System.currentTimeMillis() - startTimeMillis < maxWaitTimeMs) {
if (coll.countDocuments(Filters.in("_id", ids)) == ids.length) {
break;
}
Thread.sleep(1000L);
}
assertThatIdsAreWritten(coll, ids);
}
}