Skip to content

Fix empty stream name encountered in StorageApiFinalizeWrtiesDoFn #38410

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
Abacn merged 4 commits into from
May 15, 2026
+20 −23
Merged

Fix empty stream name encountered in StorageApiFinalizeWrtiesDoFn #38410

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ce0c3d9
Potential fix of empty stream name encountered in StorageApiFinalizeW…
Abacn May 7, 2026
29a4808
Address comments
Abacn May 14, 2026
c6fc974
Address AI comments
Abacn May 14, 2026
ceb7de6
comments
Abacn May 14, 2026
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
39 changes: 20 additions & 19 deletions ...orm/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,6 @@
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
Expand Down Expand Up @@ -559,8 +558,9 @@ private void handleAppendFailure(
String shortTableId,
AppendClientInfo appendClientInfo,
Callable<Boolean> tryCreateTable,
BiConsumer<Iterable<AppendRowsContext<DestinationT>>, Boolean> initializeContexts,
Consumer<Iterable<AppendRowsContext<DestinationT>>> initializeContexts,
Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients,
ValueState<String> streamName,
ValueState<Long> streamOffset,
MultiOutputReceiver o) {
// The first context is always the one that fails.
Expand Down Expand Up @@ -659,15 +659,6 @@ private void handleAppendFailure(
throw new RuntimeException(e);
}

if (!quotaError) {
// For known errors (offset mismatch, not found) we must reestablish
// the streams.
// However we've seen that doing this fixes random stuckness issues by reestablishing
// gRPC connections,
// so we close the clients for all non-quota errors.

clearClients.accept(failedContexts);
}
appendFailures.inc();
int retriedRows = failedContext.protoRows.getSerializedRowsCount();
BigQuerySinkMetrics.appendRowsRowStatusCounter(
Expand Down Expand Up @@ -722,11 +713,24 @@ private void handleAppendFailure(
// Finalize the stream and clear streamName so a new stream will be created.
o.get(flushTag)
.output(KV.of(failedContext.streamName, new Operation(failedContext.offset - 1, true)));

// Clear streamName so a new stream will be created.
try {
streamName.write("");
} catch (Exception e) {
throw new RuntimeException(e);
}
Comment thread
Abacn marked this conversation as resolved.
Outdated

// Re-establish the client with the new stream.
clearClients.accept(failedContexts);

@reuvenlax reuvenlax May 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

clearClients no longer even looks at failedContexts. We should simplify it and remove the parameter

@reuvenlax reuvenlax May 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

We should probably rename clearClients to resetClient

@Abacn Abacn May 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thanks, fixed


// Reinitialize all contexts with the new stream and new offsets.
initializeContexts.accept(failedContexts, true);
initializeContexts.accept(failedContexts);

// Offset failures imply that all subsequent parallel appends will also fail.
// Retry them all.
} else if (!quotaError) {
clearClients.accept(failedContexts);
}
}

Expand Down Expand Up @@ -912,13 +916,9 @@ public void process(

// Initialize stream names and offsets for all contexts. This will be called initially, but
// will also be called if we roll over to a new stream on a retry.
BiConsumer<Iterable<AppendRowsContext<DestinationT>>, Boolean> initializeContexts =
(contexts, isFailure) -> {
Consumer<Iterable<AppendRowsContext<DestinationT>>> initializeContexts =
(contexts) -> {
try {
if (isFailure) {
// Clear the stream name, forcing a new one to be created.
streamName.write("");
}
String streamNameRead = Preconditions.checkArgumentNotNull(streamName.read());

@gemini-code-assist gemini-code-assist Bot May 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

high

Preconditions.checkArgumentNotNull(streamName.read()) only verifies that the value is not null. Since handleAppendFailure now explicitly writes an empty string "" to the state (line 719), this check will pass, and the empty string will be used as the stream name for all contexts. This likely leads to the issue where StorageApiFinalizeWritesDoFn encounters an empty stream name when it receives the output from the flush tag. Consider checking for Strings.isNullOrEmpty() and ensuring a valid stream name is re-established if the state is empty.

long currentOffset = Preconditions.checkArgumentNotNull(streamOffset.read());
for (AppendRowsContext<DestinationT> context : contexts) {
Expand Down Expand Up @@ -968,6 +968,7 @@ public void process(
tryCreateTable,
initializeContexts,
clearClients,
streamName,
streamOffset,
o);
return RetryType.RETRY_ALL_OPERATIONS;
Expand Down Expand Up @@ -1068,7 +1069,7 @@ public void process(
Iterable<AppendRowsContext<DestinationT>> contexts = retryManager.getRemainingContexts();

if (numAppends > 0) {
initializeContexts.accept(contexts, false);
initializeContexts.accept(contexts);
retryManager.run(true);

appendSplitDistribution.update(numAppends);
Expand Down
Loading