You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by yh...@apache.org on 2025/08/26 14:53:38 UTC
(beam) branch master updated: Fix dup client cache key for BigQueryIO streaming STORAGE_WRITE (#35822)
This is an automated email from the ASF dual-hosted git repository.
yhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 3aa27c366ea Fix dup client cache key for BigQueryIO streaming STORAGE_WRITE (#35822)
3aa27c366ea is described below
commit 3aa27c366ea300a0714a9faa37b759e74ac5aaa0
Author: Yi Hu <ya...@google.com>
AuthorDate: Tue Aug 26 10:53:32 2025 -0400
Fix dup client cache key for BigQueryIO streaming STORAGE_WRITE (#35822)
---
.../bigquery/StorageApiWritesShardedRecords.java | 31 +++++++++++++---------
.../bigquery/TwoLevelMessageConverterCache.java | 5 ++++
2 files changed, 24 insertions(+), 12 deletions(-)
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java
index d905c4bf93c..a441803cc4f 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java
@@ -182,11 +182,11 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
}
};
- private static final Cache<ShardedKey<?>, AppendClientInfo> APPEND_CLIENTS =
+ private static final Cache<KV<String, ShardedKey<?>>, AppendClientInfo> APPEND_CLIENTS =
CacheBuilder.newBuilder()
.expireAfterAccess(5, TimeUnit.MINUTES)
.removalListener(
- (RemovalNotification<ShardedKey<?>, AppendClientInfo> removal) -> {
+ (RemovalNotification<KV<String, ShardedKey<?>>, AppendClientInfo> removal) -> {
final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
if (appendClientInfo != null) {
appendClientInfo.close();
@@ -580,14 +580,18 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
};
AtomicReference<AppendClientInfo> appendClientInfo =
- new AtomicReference<>(APPEND_CLIENTS.get(element.getKey(), getAppendClientInfo));
+ new AtomicReference<>(
+ APPEND_CLIENTS.get(
+ messageConverters.getAppendClientKey(element.getKey()), getAppendClientInfo));
String currentStream = getOrCreateStream.get();
if (!currentStream.equals(appendClientInfo.get().getStreamName())) {
// Cached append client is inconsistent with persisted state. Throw away cached item and
// force it to be
// recreated.
- APPEND_CLIENTS.invalidate(element.getKey());
- appendClientInfo.set(APPEND_CLIENTS.get(element.getKey(), getAppendClientInfo));
+ APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
+ appendClientInfo.set(
+ APPEND_CLIENTS.get(
+ messageConverters.getAppendClientKey(element.getKey()), getAppendClientInfo));
}
TableSchema updatedSchemaValue = updatedSchema.read();
@@ -596,8 +600,9 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
appendClientInfo.set(
AppendClientInfo.of(
updatedSchemaValue, appendClientInfo.get().getCloseAppendClient(), false));
- APPEND_CLIENTS.invalidate(element.getKey());
- APPEND_CLIENTS.put(element.getKey(), appendClientInfo.get());
+ APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
+ APPEND_CLIENTS.put(
+ messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
}
}
@@ -664,9 +669,10 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
Consumer<Iterable<AppendRowsContext>> clearClients =
contexts -> {
- APPEND_CLIENTS.invalidate(element.getKey());
+ APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
appendClientInfo.set(appendClientInfo.get().withNoAppendClient());
- APPEND_CLIENTS.put(element.getKey(), appendClientInfo.get());
+ APPEND_CLIENTS.put(
+ messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
for (AppendRowsContext context : contexts) {
if (context.client != null) {
// Unpin in a different thread, as it may execute a blocking close.
@@ -960,8 +966,9 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
appendClientInfo.set(
AppendClientInfo.of(
newSchema.get(), appendClientInfo.get().getCloseAppendClient(), false));
- APPEND_CLIENTS.invalidate(element.getKey());
- APPEND_CLIENTS.put(element.getKey(), appendClientInfo.get());
+ APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
+ APPEND_CLIENTS.put(
+ messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
LOG.debug(
"Fetched updated schema for table {}:\n\t{}", tableId, updatedSchemaReturned);
updatedSchema.write(newSchema.get());
@@ -993,7 +1000,7 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
streamName.clear();
streamOffset.clear();
// Make sure that the stream object is closed.
- APPEND_CLIENTS.invalidate(key);
+ APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(key));
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TwoLevelMessageConverterCache.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TwoLevelMessageConverterCache.java
index 5f90e1dd395..0ce7c7573c9 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TwoLevelMessageConverterCache.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TwoLevelMessageConverterCache.java
@@ -20,6 +20,7 @@ package org.apache.beam.sdk.io.gcp.bigquery;
import java.io.Serializable;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService;
import org.apache.beam.sdk.io.gcp.bigquery.StorageApiDynamicDestinations.MessageConverter;
+import org.apache.beam.sdk.util.ShardedKey;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
@@ -71,4 +72,8 @@ class TwoLevelMessageConverterCache<DestinationT extends @NonNull Object, Elemen
KV.of(operationName, destination),
() -> dynamicDestinations.getMessageConverter(destination, datasetService)));
}
+
+ public KV<String, ShardedKey<?>> getAppendClientKey(ShardedKey<DestinationT> shardedKey) {
+ return KV.of(operationName, shardedKey);
+ }
}