Skip to content
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

feat: only close writers on flush #26

Merged
merged 3 commits into from
Apr 4, 2024
Merged

feat: only close writers on flush #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d762d9f
feat: only close writers on flush
gintarasm Apr 2, 2024
59d9662
fix: do not flush on close
gintarasm Apr 3, 2024
837ce1a
refactor: do not flush on close
gintarasm Apr 3, 2024
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
35 changes: 17 additions & 18 deletions src/main/java/com/vinted/flink/bigquery/sink/async/AsyncBigQuerySinkWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ public class AsyncBigQuerySinkWriter<A> extends AsyncSinkWriter<Rows<A>, StreamR

private final Executor appendExecutor;

protected transient Queue<StreamWriter> writersToClose = new ArrayDeque<>();
protected transient Map<String, StreamWriter> streamMap = new ConcurrentHashMap<>();

public AsyncBigQuerySinkWriter(ExecutorProvider executorProvider, AsyncClientProvider clientProvider, ElementConverter<Rows<A>, StreamRequest> elementConverter, Sink.InitContext context, AsyncSinkWriterConfiguration configuration, Collection<BufferedRequestState<StreamRequest>> bufferedRequestStates) {
Expand Down Expand Up @@ -94,11 +95,7 @@ protected final void recreateStreamWriter(String traceId, String streamName, Str
streamMap.replaceAll((key, writer) -> {
var newWriter = writer;
if (writer.getWriterId().equals(writerId)) {
try {
writer.close();
} catch (Exception e) {
logger.trace("Trace-id {} Could not close writer for {}", traceId, streamName);
}
writersToClose.add(writer);
newWriter = this.clientProvider.getWriter(streamName, table);
registerInflightMetric(newWriter);
}
Expand Down Expand Up @@ -230,21 +227,23 @@ protected long getSizeInBytes(StreamRequest StreamRequest) {
}

@Override
public void close() {
logger.info("Closing BigQuery write stream");
try {
flush(true);
streamMap.values().forEach(stream -> {
try {
stream.close();
} catch (Exception e) {
throw new RuntimeException(e);
}
});
} catch (InterruptedException e) {
throw new RuntimeException(e);
public void flush(boolean flush) throws InterruptedException {
super.flush(flush);

while (writersToClose.peek() != null) {
var writer = writersToClose.poll();
try {
writer.close();
} catch (Exception e) {
logger.error("Could not close unused writer", e);
}
}
}

@Override
public void close() {
logger.info("Closing BigQuery write stream");
streamMap.values().forEach(StreamWriter::close);
}


Expand Down
4 changes: 2 additions & 2 deletions src/test/java/com/vinted/flink/bigquery/AsyncBigQuerySinkTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ public void shouldRecreateWriterAndRetryFailingWithMaximumRequestCallbackWaitTim

@Test
public void shouldFailAndNotRetryWhenUnknownErrorReceived(@FlinkTest.FlinkParam FlinkTest.PipelineRunner runner, @FlinkTest.FlinkParam MockAsyncProtoClientProvider mockClientProvider) throws Exception {
mockClientProvider.givenStreamIsFinalized(stream);
mockClientProvider.givenFailingAppendWithStatus(Status.UNKNOWN);

assertThatThrownBy(() -> {
runner
Expand All @@ -76,7 +76,7 @@ public void shouldFailAndNotRetryWhenUnknownErrorReceived(@FlinkTest.FlinkParam
}).isInstanceOf(JobExecutionException.class);


verify(mockClientProvider.getMockProtoWriter(), times(2)).append(any());
verify(mockClientProvider.getMockProtoWriter(), times(1)).append(any());
}

@Test
Expand Down
Loading