You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "reuvenlax (via GitHub)" <gi...@apache.org> on 2026/04/22 03:51:28 UTC

[PR] Refactor storageapi conversion [beam]

reuvenlax opened a new pull request, #38264:
URL: https://github.com/apache/beam/pull/38264

   **Please** add a meaningful description for your change here
   
   ------------------------
   
   Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
   
    - [ ] Mention the appropriate issue in your description (for example: `addresses #123`), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment `fixes #<ISSUE NUMBER>` instead.
    - [ ] Update `CHANGES.md` with noteworthy changes.
    - [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
   
   See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://github.com/apache/beam/blob/master/CONTRIBUTING.md#make-the-reviewers-job-easier).
   
   To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
   
   GitHub Actions Tests Status (on master branch)
   ------------------------------------------------------------------------------------------------
   [![Build python source distribution and wheels](https://github.com/apache/beam/actions/workflows/build_wheels.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
   [![Python tests](https://github.com/apache/beam/actions/workflows/python_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
   [![Java tests](https://github.com/apache/beam/actions/workflows/java_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
   [![Go tests](https://github.com/apache/beam/actions/workflows/go_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
   
   See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI or the [workflows README](https://github.com/apache/beam/blob/master/.github/workflows/README.md) to see a list of phrases to trigger workflows.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3149558552


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo get(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    return appendCache.get(key, wrapWithPin(loader));
+  }

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   This method should be `synchronized` to maintain the atomicity protocol described in the class Javadoc. Without synchronization, a concurrent access via this method could trigger cache maintenance and evict an entry that another thread is currently pinning in `getAndPin`.
   
   ```suggestion
     public synchronized AppendClientInfo get(KeyT key, Callable<AppendClientInfo> loader)
         throws Exception {
       return appendCache.get(key, wrapWithPin(loader));
     }
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo get(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    return appendCache.get(key, wrapWithPin(loader));
+  }
+
+  public AppendClientInfo putAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = wrapWithPin(loader).call();
+      appendCache.put(key, info);
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo put(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    AppendClientInfo info = wrapWithPin(loader).call();
+    appendCache.put(key, info);
+    return info;
+  }

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   This method should be `synchronized` to ensure that the cache access is atomic relative to other operations like `getAndPin` and `invalidate`. This prevents races where an entry might be evicted or invalidated while being updated.
   
   ```suggestion
     public synchronized AppendClientInfo put(KeyT key, Callable<AppendClientInfo> loader)
         throws Exception {
       AppendClientInfo info = wrapWithPin(loader).call();
       appendCache.put(key, info);
       return info;
     }
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo get(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    return appendCache.get(key, wrapWithPin(loader));
+  }
+
+  public AppendClientInfo putAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = wrapWithPin(loader).call();
+      appendCache.put(key, info);
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo put(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    AppendClientInfo info = wrapWithPin(loader).call();
+    appendCache.put(key, info);
+    return info;
+  }
+
+  public void invalidate(KeyT key, AppendClientInfo expectedClient) {
+    // The default stream is cached across multiple different DoFns. If they all try
+    // and
+    // invalidate, then we can get races between threads invalidating and recreating
+    // streams. For this reason,
+    // we check to see that the cache still contains the object we created before
+    // invalidating (in case another
+    // thread has already invalidated and recreated the stream).
+    synchronized (this) {
+      AppendClientInfo cachedAppendClient = appendCache.getIfPresent(key);
+      if (cachedAppendClient != null
+          && System.identityHashCode(cachedAppendClient)
+              == System.identityHashCode(expectedClient)) {
+        appendCache.invalidate(key);
+      }
+    }
+  }
+
+  public void invalidate(KeyT key) {
+    synchronized (this) {
+      appendCache.invalidate(key);
+    }
+  }
+
+  public void tickle(KeyT key) {
+    appendCache.getIfPresent(key);
+  }

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   The `tickle` method should be `synchronized`. Since it performs a cache access (`getIfPresent`), it can trigger maintenance and eviction. To protect the atomicity of the pinning protocol used in `getAndPin`, all cache access methods must be synchronized on the same monitor.
   
   ```suggestion
     public synchronized void tickle(KeyT key) {
       appendCache.getIfPresent(key);
     }
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java:
##########
@@ -199,4 +219,38 @@ public TableRow toTableRow(ByteString protoBytes, Predicate<String> includeField
       throw new RuntimeException(e);
     }
   }
+
+  public void pinAppendClient() {
+    BigQueryServices.StreamAppendClient client =
+        Preconditions.checkStateNotNull(getStreamAppendClient());
+    client.pin();
+  }
+
+  public void unpinAppendClient(@Nullable ExecutorService executor) {
+    BigQueryServices.StreamAppendClient client =
+        Preconditions.checkStateNotNull(getStreamAppendClient());
+    if (executor != null) {
+      runAsyncIgnoreFailure(executor, client::unpin);
+    } else {
+      client.unpin();
+    }
+  }
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  private static void runAsyncIgnoreFailure(ExecutorService executor, ThrowingRunnable task) {
+    executor.submit(
+        () -> {
+          try {
+            task.run();
+          } catch (Throwable e) {
+            String msg =
+                e.toString()
+                    + "\n"
+                    + Arrays.stream(e.getStackTrace())
+                        .map(StackTraceElement::toString)
+                        .collect(Collectors.joining("\n"));
+            System.err.println("Exception happened while executing async task. Ignoring: " + msg);

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   Avoid using `System.err.println` for error reporting. Use the class logger instead. Additionally, let the logger handle the exception object directly rather than manually formatting the stack trace.
   
   ```java
               LOG.error("Exception happened while executing async task. Ignoring", e);
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -554,6 +553,214 @@ private CreateRetryManagerResult<DestinationT> createRetryManager(
           retryManager, failedRows, recordsAppended, histogramUpdates);
     }
 
+    private void handleAppendFailure(
+        Iterable<AppendRowsContext<DestinationT>> failedContexts,
+        TableReference tableReference,
+        String shortTableId,
+        AppendClientInfo appendClientInfo,
+        Callable<Boolean> tryCreateTable,
+        BiConsumer<Iterable<AppendRowsContext<DestinationT>>, Boolean> initializeContexts,
+        Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients,
+        ValueState<Long> streamOffset,
+        MultiOutputReceiver o) {
+      // The first context is always the one that fails.
+      AppendRowsContext<DestinationT> failedContext =
+          Preconditions.checkStateNotNull(Iterables.getFirst(failedContexts, null));
+      BigQuerySinkMetrics.reportFailedRPCMetrics(
+          failedContext, BigQuerySinkMetrics.RpcMethod.APPEND_ROWS, shortTableId);
+      String errorCode = BigQuerySinkMetrics.throwableToGRPCCodeString(failedContext.getError());
+
+      // AppendSerializationError means that BigQuery detected errors on individual rows, e.g.
+      // a row not conforming to bigQuery invariants. These errors are persistent, so we redirect
+      // those rows to the
+      // failedInserts PCollection, and retry with the remaining rows.
+      if (failedContext.getError() != null
+          && failedContext.getError() instanceof Exceptions.AppendSerializationError) {
+        Exceptions.AppendSerializationError error =
+            Preconditions.checkArgumentNotNull(
+                (Exceptions.AppendSerializationError) failedContext.getError());
+
+        Set<Integer> failedRowIndices = error.getRowIndexToErrorMessage().keySet();
+        for (int failedIndex : failedRowIndices) {
+          // Convert the message to a TableRow and send it to the failedRows collection.
+          TableRow failedRow = failedContext.failsafeTableRows.get(failedIndex);
+          if (failedRow == null) {
+            ByteString protoBytes = failedContext.protoRows.getSerializedRows(failedIndex);
+            failedRow = appendClientInfo.toTableRow(protoBytes, Predicates.alwaysTrue());
+          }
+          org.joda.time.Instant timestamp = failedContext.timestamps.get(failedIndex);
+          o.get(failedRowsTag)
+              .outputWithTimestamp(
+                  new BigQueryStorageApiInsertError(
+                      failedRow,
+                      error.getRowIndexToErrorMessage().get(failedIndex),
+                      tableReference),
+                  timestamp);
+        }
+        int failedRows = failedRowIndices.size();
+        rowsSentToFailedRowsCollection.inc(failedRows);
+        BigQuerySinkMetrics.appendRowsRowStatusCounter(
+                BigQuerySinkMetrics.RowStatus.FAILED, errorCode, shortTableId)
+            .inc(failedRows);
+
+        // Remove the failed row from the payload, so we retry the batch without the failed
+        // rows.
+        ProtoRows.Builder retryRows = ProtoRows.newBuilder();
+        @Nullable List<org.joda.time.Instant> timestamps = Lists.newArrayList();
+        for (int i = 0; i < failedContext.protoRows.getSerializedRowsCount(); ++i) {
+          if (!failedRowIndices.contains(i)) {
+            ByteString rowBytes = failedContext.protoRows.getSerializedRows(i);
+            retryRows.addSerializedRows(rowBytes);
+            timestamps.add(failedContext.timestamps.get(i));
+          }
+        }
+        failedContext.protoRows = retryRows.build();
+        failedContext.timestamps = timestamps;
+        int retriedRows = failedContext.protoRows.getSerializedRowsCount();
+        BigQuerySinkMetrics.appendRowsRowStatusCounter(
+                BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableId)
+            .inc(retriedRows);
+
+        // Since we removed rows, we need to update the insert offsets for all remaining rows.
+        long offset = failedContext.offset;
+        for (AppendRowsContext<DestinationT> context : failedContexts) {
+          context.offset = offset;
+          offset += context.protoRows.getSerializedRowsCount();
+        }
+        streamOffset.write(offset);
+        return;
+      }
+
+      Throwable error = Preconditions.checkStateNotNull(failedContext.getError());
+      Status.Code statusCode = Status.fromThrowable(error).getCode();
+
+      // This means that the offset we have stored does not match the current end of
+      // the stream in the Storage API. Usually this happens because a crash or a bundle
+      // failure
+      // happened after an append but before the worker could checkpoint it's
+      // state. The records that were appended in a failed bundle will be retried,
+      // meaning that the unflushed tail of the stream must be discarded to prevent
+      // duplicates.
+      boolean offsetMismatch =
+          statusCode.equals(Code.OUT_OF_RANGE) || statusCode.equals(Code.ALREADY_EXISTS);
+
+      boolean quotaError = statusCode.equals(Code.RESOURCE_EXHAUSTED);
+      if (!offsetMismatch) {
+        // Don't log errors for expected offset mismatch. These will be logged as warnings
+        // below.
+        LOG.error("Got error {} closing {}", failedContext.getError(), failedContext.streamName);
+      }
+
+      try {
+        // TODO: Only do this on explicit NOT_FOUND errors once BigQuery reliably produces
+        // them.
+        tryCreateTable.call();
+      } catch (Exception e) {
+        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(
+              BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableId)
+          .inc(retriedRows);
+
+      // Schema mismatched exceptions can happen if the table was recently updated. Since
+      // vortex caches schemas
+      // we might see the new schema before vortex does. In this case, we simply need to
+      // retry.
+      Exceptions.@Nullable StorageException storageException = Exceptions.toStorageException(error);
+      boolean schemaMismatchError =
+          (storageException instanceof Exceptions.SchemaMismatchedException);
+      if (!schemaMismatchError) {
+        // There's no special error code for missing required fields, and that can also happen
+        // due to vortex
+        // being delayed at seeing a new schema. We're forced to parse the description to
+        // determine that this has happened.
+        // TODO: Vortex team to introduce a special storage error code for this, so we don't
+        // have to parse
+        // descriptions.
+        Status status = Status.fromThrowable(error);
+        if (status.getCode() == Code.INVALID_ARGUMENT) {
+          String description = status.getDescription();
+          schemaMismatchError = description != null && description.contains("incompatible fields");
+        }
+      }
+      if (schemaMismatchError) {
+        LOG.info(
+            "Vortex failed stream open due to incompatible fields. This is likely because the BigTable "
+                + "schema was recently updated and Vortex hasn't noticed yet, so retrying. error {}",

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   Typo: 'BigTable' should be 'BigQuery' in this context, as this class handles BigQuery Storage API writes.
   
   ```java
               "Vortex failed stream open due to incompatible fields. This is likely because the BigQuery "
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java:
##########
@@ -134,7 +143,18 @@ public AppendClientInfo withAppendClient(
   public void close() {
     BigQueryServices.StreamAppendClient client = getStreamAppendClient();
     if (client != null) {
-      getCloseAppendClient().accept(client);
+      try {
+        getCloseAppendClient().accept(client);
+      } catch (Exception e) {
+        // We ignore errors when closing clients.
+        String msg =
+            e
+                + "\n"
+                + Arrays.stream(e.getStackTrace())
+                    .map(StackTraceElement::toString)
+                    .collect(Collectors.joining("\n"));
+        LOG.warn("Caught exception whilw trying to close append client. Ignoring {}", msg);

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   The manual formatting of the stack trace is unnecessary and inefficient. SLF4J loggers can handle exceptions directly, which includes printing the stack trace. This also fixes a typo ('whilw' -> 'while').
   
   ```java
           LOG.warn("Caught exception while trying to close append client. Ignoring", e);
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -149,44 +141,12 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
   private final Coder<TableRow> succussfulRowsCoder;
 
   private final TupleTag<KV<String, Operation>> flushTag = new TupleTag<>("flushTag");
-  private static final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
-
-  private static final Cache<KV<String, ShardedKey<?>>, AppendClientInfo> APPEND_CLIENTS =
-      CacheBuilder.newBuilder()
-          .expireAfterAccess(5, TimeUnit.MINUTES)
-          .removalListener(
-              (RemovalNotification<KV<String, ShardedKey<?>>, AppendClientInfo> removal) -> {
-                final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
-                if (appendClientInfo != null) {
-                  appendClientInfo.close();
-                }
-              })
-          .build();
 
-  static void clearCache() {
-    APPEND_CLIENTS.invalidateAll();
-  }
+  private static AppendClientCache<KV<String, ShardedKey<?>>> APPEND_CLIENTS =
+      new AppendClientCache<>(Duration.standardMinutes(5));

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   The `APPEND_CLIENTS` cache should be declared as `final` to ensure it is not reassigned and to follow best practices for static constants.
   
   ```suggestion
     private static final AppendClientCache<KV<String, ShardedKey<?>>> APPEND_CLIENTS =
         new AppendClientCache<>(Duration.standardMinutes(5));
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4329426311

   /gemini review


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3155920985


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo putAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = wrapWithPin(loader).call();

Review Comment:
   I'm not sure I understand - do you mean unpin other clients when inserting this one?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "ahmedabu98 (via GitHub)" <gi...@apache.org>.
ahmedabu98 commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3156441559


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo putAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = wrapWithPin(loader).call();

Review Comment:
   Sorry I think I just got mixed up on the unpinning logic. 
   I thought we were leaking resources by not unpinning the consuming `processElement`s, but I just did another pass and it looks fine to me.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "ahmedabu98 (via GitHub)" <gi...@apache.org>.
ahmedabu98 commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3156445042


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  /** "Refresh" an object by invalidating the old cache entry. */
+  public AppendClientInfo refreshObjectAndAndPin(KeyT key, Callable<AppendClientInfo> loader)

Review Comment:
   nit: typo
   ```suggestion
     public AppendClientInfo refreshObjectAndPin(KeyT key, Callable<AppendClientInfo> loader)
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax merged PR #38264:
URL: https://github.com/apache/beam/pull/38264


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4329796959

   /gemini review


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3149662316


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   It is recommended to use a named `ThreadFactory` with daemon threads for the `ExecutorService`. This improves observability in thread dumps and ensures that the threads do not prevent the JVM from exiting if the process is shutting down.
   
   ```java
     private final ExecutorService closeWriterExecutor =
         Executors.newCachedThreadPool(
             new org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder()
                 .setDaemon(true)
                 .setNameFormat("AppendClientCache-closer-%d")
                 .build());
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java:
##########
@@ -134,7 +143,18 @@ public AppendClientInfo withAppendClient(
   public void close() {
     BigQueryServices.StreamAppendClient client = getStreamAppendClient();
     if (client != null) {
-      getCloseAppendClient().accept(client);
+      try {
+        getCloseAppendClient().accept(client);
+      } catch (Exception e) {
+        // We ignore errors when closing clients.
+        String msg =
+            e
+                + "\n"
+                + Arrays.stream(e.getStackTrace())
+                    .map(StackTraceElement::toString)
+                    .collect(Collectors.joining("\n"));
+        LOG.warn("Caught exception whilw trying to close append client. Ignoring {}", msg);

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   The manual formatting of the stack trace is redundant and inefficient. SLF4J's logger can handle exceptions directly, which includes printing the stack trace. Additionally, there is a typo in the log message ('whilw' instead of 'while').
   
   ```java
           LOG.warn("Caught exception while trying to close append client. Ignoring", e);
   ```



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   Using `synchronized(this)` on the entire `getAndPin` method creates a global bottleneck for all cache accesses. Since the `loader` (which creates the `StreamAppendClient`) can be slow due to network calls, this will block all other threads trying to access the cache for different keys. Consider using a more granular locking mechanism, such as Guava's `Striped` locks, to allow concurrent access for different keys.



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java:
##########
@@ -199,4 +219,32 @@ public TableRow toTableRow(ByteString protoBytes, Predicate<String> includeField
       throw new RuntimeException(e);
     }
   }
+
+  public void pinAppendClient() {
+    BigQueryServices.StreamAppendClient client =
+        Preconditions.checkStateNotNull(getStreamAppendClient());
+    client.pin();
+  }
+
+  public void unpinAppendClient(@Nullable ExecutorService executor) {
+    BigQueryServices.StreamAppendClient client =
+        Preconditions.checkStateNotNull(getStreamAppendClient());
+    if (executor != null) {
+      runAsyncIgnoreFailure(executor, client::unpin);
+    } else {
+      client.unpin();
+    }
+  }
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  private static void runAsyncIgnoreFailure(ExecutorService executor, ThrowingRunnable task) {
+    executor.submit(
+        () -> {
+          try {
+            task.run();
+          } catch (Throwable e) {
+            LOG.info("Exception happened while executing async task. Ignoring: ", e);

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   Logging an exception at `INFO` level is unusual. Even if the failure is ignored, it is generally better to log it at `WARN` or `ERROR` level to make it more visible in logs when troubleshooting background tasks.
   
   ```suggestion
               LOG.warn("Exception happened while executing async task. Ignoring: ", e);
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4329425947

   /gemini summary


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4329425208

   /gemini


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4329428889

   R: @ahmedabu98 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3155960906


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -671,475 +880,223 @@ public void process(
                 AppendClientInfo.of(
                         Preconditions.checkStateNotNull(tableSchema),
                         descriptor,
-                        // Make sure that the client is always closed in a different thread
-                        // to
-                        // avoid blocking.
-                        client ->
-                            runAsyncIgnoreFailure(
-                                closeWriterExecutor,
-                                () -> {
-                                  // Remove the pin that is "owned" by the cache.
-                                  client.unpin();
-                                  client.close();
-                                }))
+                        AutoCloseable::close)
                     .withAppendClient(
                         writeStreamService,
                         getOrCreateStream,
                         false,
                         defaultMissingValueInterpretation);
-            // This pin is "owned" by the cache.
-            Preconditions.checkStateNotNull(info.getStreamAppendClient()).pin();
             return info;
           };
 
-      AtomicReference<AppendClientInfo> appendClientInfo =
-          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(messageConverters.getAppendClientKey(element.getKey()));
-        appendClientInfo.set(
-            APPEND_CLIENTS.get(
-                messageConverters.getAppendClientKey(element.getKey()), getAppendClientInfo));
-      }
-
-      TableSchema updatedSchemaValue = updatedSchema.read();
-      if (autoUpdateSchema && updatedSchemaValue != null) {
-        if (appendClientInfo.get().hasSchemaChanged(updatedSchemaValue)) {
-          appendClientInfo.set(
-              AppendClientInfo.of(
-                  updatedSchemaValue, appendClientInfo.get().getCloseAppendClient(), false));
-          APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-          APPEND_CLIENTS.put(
-              messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
+      // The StreamWriter has two pins on it. The static cache holds a pin, as it continues to cache
+      // values after this
+      // method exits, so must hold the pin. The local AppendClientHolder also holds a pin, as the
+      // cache could in
+      // theory evict the object during execution and we want a pin held throughout the execution of
+      // this function.
+      try (AppendClientHolder appendClientHolder =
+          new AppendClientHolder(element.getKey(), getAppendClientInfo)) {
+        String currentStream = getOrCreateStream.get();
+        if (!currentStream.equals(appendClientHolder.get().getStreamName())) {
+          // Cached append client is inconsistent with persisted state. Throw away cached item and
+          // force it to be recreated.
+          appendClientHolder.invalidateAndReset();
         }
-      }
 
-      // 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) -> {
-            try {
-              if (isFailure) {
-                // Clear the stream name, forcing a new one to be created.
-                streamName.write("");
-              }
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              StreamAppendClient streamAppendClient =
-                  Preconditions.checkArgumentNotNull(
-                      appendClientInfo.get().getStreamAppendClient());
-              String streamNameRead = Preconditions.checkArgumentNotNull(streamName.read());
-              long currentOffset = Preconditions.checkArgumentNotNull(streamOffset.read());
-              for (AppendRowsContext<DestinationT> context : contexts) {
-                context.streamName = streamNameRead;
-                streamAppendClient.pin();
-                context.client = appendClientInfo.get().getStreamAppendClient();
-                context.offset = currentOffset;
-                ++context.tryIteration;
-                currentOffset = context.offset + context.protoRows.getSerializedRowsCount();
-              }
-              streamOffset.write(currentOffset);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
-
-      Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients =
-          (contexts) -> {
-            APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-            appendClientInfo.set(appendClientInfo.get().withNoAppendClient());
-            APPEND_CLIENTS.put(
-                messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
-            for (AppendRowsContext<DestinationT> context : contexts) {
-              if (context.client != null) {
-                // Unpin in a different thread, as it may execute a blocking close.
-                runAsyncIgnoreFailure(closeWriterExecutor, context.client::unpin);
-                context.client = null;
-              }
-            }
-          };
-
-      Function<AppendRowsContext<DestinationT>, ApiFuture<AppendRowsResponse>> runOperation =
-          context -> {
-            if (context.protoRows.getSerializedRowsCount() == 0) {
-              // This might happen if all rows in a batch failed and were sent to the failed-rows
-              // PCollection.
-              return ApiFutures.immediateFuture(AppendRowsResponse.newBuilder().build());
-            }
-            try {
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              return Preconditions.checkStateNotNull(appendClientInfo.get().getStreamAppendClient())
-                  .appendRows(context.offset, context.protoRows);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
+        TableSchema updatedSchemaValue = autoUpdateSchema ? updatedSchema.read() : null;
+        if (autoUpdateSchema && updatedSchemaValue != null) {
+          if (appendClientHolder.get().hasSchemaChanged(updatedSchemaValue)) {
+            appendClientHolder.invalidateAndReset();
+          }
+        }
 
-      Function<Iterable<AppendRowsContext<DestinationT>>, RetryType> onError =
-          failedContexts -> {
-            // The first context is always the one that fails.
-            AppendRowsContext<DestinationT> failedContext =
-                Preconditions.checkStateNotNull(Iterables.getFirst(failedContexts, null));
-            BigQuerySinkMetrics.reportFailedRPCMetrics(
-                failedContext, BigQuerySinkMetrics.RpcMethod.APPEND_ROWS, shortTableId);
-            String errorCode =
-                BigQuerySinkMetrics.throwableToGRPCCodeString(failedContext.getError());
-
-            // AppendSerializationError means that BigQuery detected errors on individual rows, e.g.
-            // a row not conforming
-            // to bigQuery invariants. These errors are persistent, so we redirect those rows to the
-            // failedInserts
-            // PCollection, and retry with the remaining rows.
-            if (failedContext.getError() != null
-                && failedContext.getError() instanceof Exceptions.AppendSerializationError) {
-              Exceptions.AppendSerializationError error =
-                  Preconditions.checkArgumentNotNull(
-                      (Exceptions.AppendSerializationError) failedContext.getError());
-
-              Set<Integer> failedRowIndices = error.getRowIndexToErrorMessage().keySet();
-              for (int failedIndex : failedRowIndices) {
-                // Convert the message to a TableRow and send it to the failedRows collection.
-                TableRow failedRow = failedContext.failsafeTableRows.get(failedIndex);
-                if (failedRow == null) {
-                  // TODO: MAKE SURE WE USE UPDATED DESCRIPTOR
-                  ByteString protoBytes = failedContext.protoRows.getSerializedRows(failedIndex);
-                  failedRow =
-                      appendClientInfo.get().toTableRow(protoBytes, Predicates.alwaysTrue());
+        // 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) -> {
+              try {
+                if (isFailure) {
+                  // Clear the stream name, forcing a new one to be created.
+                  streamName.write("");
                 }
-                org.joda.time.Instant timestamp = failedContext.timestamps.get(failedIndex);
-                o.get(failedRowsTag)
-                    .outputWithTimestamp(
-                        new BigQueryStorageApiInsertError(
-                            failedRow,
-                            error.getRowIndexToErrorMessage().get(failedIndex),
-                            tableReference),
-                        timestamp);
-              }
-              int failedRows = failedRowIndices.size();
-              rowsSentToFailedRowsCollection.inc(failedRows);
-              BigQuerySinkMetrics.appendRowsRowStatusCounter(
-                      BigQuerySinkMetrics.RowStatus.FAILED, errorCode, shortTableId)
-                  .inc(failedRows);
-
-              // Remove the failed row from the payload, so we retry the batch without the failed
-              // rows.
-              ProtoRows.Builder retryRows = ProtoRows.newBuilder();
-              @Nullable List<org.joda.time.Instant> timestamps = Lists.newArrayList();
-              for (int i = 0; i < failedContext.protoRows.getSerializedRowsCount(); ++i) {
-                if (!failedRowIndices.contains(i)) {
-                  ByteString rowBytes = failedContext.protoRows.getSerializedRows(i);
-                  retryRows.addSerializedRows(rowBytes);
-                  timestamps.add(failedContext.timestamps.get(i));
+                String streamNameRead = Preconditions.checkArgumentNotNull(streamName.read());
+                long currentOffset = Preconditions.checkArgumentNotNull(streamOffset.read());
+                for (AppendRowsContext<DestinationT> context : contexts) {
+                  context.streamName = streamNameRead;
+                  context.offset = currentOffset;
+                  ++context.tryIteration;
+                  currentOffset = context.offset + context.protoRows.getSerializedRowsCount();
                 }
+                streamOffset.write(currentOffset);
+              } catch (Exception e) {
+                throw new RuntimeException(e);
               }
-              failedContext.protoRows = retryRows.build();
-              failedContext.timestamps = timestamps;
-              int retriedRows = failedContext.protoRows.getSerializedRowsCount();
-              BigQuerySinkMetrics.appendRowsRowStatusCounter(
-                      BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableId)
-                  .inc(retriedRows);
-
-              // Since we removed rows, we need to update the insert offsets for all remaining rows.
-              long offset = failedContext.offset;
-              for (AppendRowsContext<DestinationT> context : failedContexts) {
-                context.offset = offset;
-                offset += context.protoRows.getSerializedRowsCount();
+            };
+
+        Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients =
+            (contexts) -> {
+              try {
+                appendClientHolder.invalidateAndReset();
+              } catch (Exception e) {
+                throw new RuntimeException(e);
               }
-              streamOffset.write(offset);
-              return RetryType.RETRY_ALL_OPERATIONS;
-            }
-
-            Throwable error = Preconditions.checkStateNotNull(failedContext.getError());
-
-            Status.Code statusCode = Status.fromThrowable(error).getCode();
-
-            // This means that the offset we have stored does not match the current end of
-            // the stream in the Storage API. Usually this happens because a crash or a bundle
-            // failure
-            // happened after an append but before the worker could checkpoint it's
-            // state. The records that were appended in a failed bundle will be retried,
-            // meaning that the unflushed tail of the stream must be discarded to prevent
-            // duplicates.
-            boolean offsetMismatch =
-                statusCode.equals(Code.OUT_OF_RANGE) || statusCode.equals(Code.ALREADY_EXISTS);
-
-            boolean quotaError = statusCode.equals(Code.RESOURCE_EXHAUSTED);
-            if (!offsetMismatch) {
-              // Don't log errors for expected offset mismatch. These will be logged as warnings
-              // below.
-              LOG.error(
-                  "Got error {} closing {}", failedContext.getError(), failedContext.streamName);
-            }
-
-            try {
-              // TODO: Only do this on explicit NOT_FOUND errors once BigQuery reliably produces
-              // them.
-              tryCreateTable.call();
-            } catch (Exception e) {
-              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(
-                    BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableId)
-                .inc(retriedRows);
-
-            // Schema mismatched exceptions can happen if the table was recently updated. Since
-            // vortex caches schemas
-            // we might see the new schema before vortex does. In this case, we simply need to
-            // retry.
-            Exceptions.@Nullable StorageException storageException =
-                Exceptions.toStorageException(error);
-            boolean schemaMismatchError =
-                (storageException instanceof Exceptions.SchemaMismatchedException);
-            if (!schemaMismatchError) {
-              // There's no special error code for missing required fields, and that can also happen
-              // due to vortex
-              // being delayed at seeing a new schema. We're forced to parse the description to
-              // determine that this has happened.
-              // TODO: Vortex team to introduce a special storage error code for this, so we don't
-              // have to parse
-              // descriptions.
-              Status status = Status.fromThrowable(error);
-              if (status.getCode() == Code.INVALID_ARGUMENT) {
-                String description = status.getDescription();
-                schemaMismatchError =
-                    description != null && description.contains("incompatible fields");
+            };
+
+        Function<AppendRowsContext<DestinationT>, ApiFuture<AppendRowsResponse>> runOperation =
+            context -> {
+              if (context.protoRows.getSerializedRowsCount() == 0) {
+                // This might happen if all rows in a batch failed and were sent to the failed-rows
+                // PCollection.
+                return ApiFutures.immediateFuture(AppendRowsResponse.newBuilder().build());
               }
-            }
-            if (schemaMismatchError) {
-              LOG.info(
-                  "Vortex failed stream open due to incompatible fields. This is likely because the BigTable "
-                      + "schema was recently updated and Vortex hasn't noticed yet, so retrying. error {}",
-                  Preconditions.checkStateNotNull(error).toString());
-            }
-
-            boolean explicitStreamFinalized =
-                failedContext.getError() instanceof StreamFinalizedException;
-            // This implies that the stream doesn't exist or has already been finalized. In this
-            // case we have no choice but to create a new stream.
-            boolean streamDoesNotExist =
-                explicitStreamFinalized
-                    || statusCode.equals(Code.INVALID_ARGUMENT)
-                    || statusCode.equals(Code.NOT_FOUND)
-                    || statusCode.equals(Code.FAILED_PRECONDITION);
-            streamDoesNotExist = streamDoesNotExist && !schemaMismatchError;
-
-            if (offsetMismatch || streamDoesNotExist) {
-              appendOffsetFailures.inc();
-              LOG.warn(
-                  "Append to {} failed. Will retry with a new stream",
-                  failedContext,
-                  failedContext.getError());
-              // 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)));
-              // Reinitialize all contexts with the new stream and new offsets.
-              initializeContexts.accept(failedContexts, true);
-
-              // Offset failures imply that all subsequent parallel appends will also fail.
-              // Retry them all.
+              try {
+                return appendClientHolder
+                    .getStreamAppendClient()
+                    .appendRows(context.offset, context.protoRows);
+              } catch (Exception e) {
+                throw new RuntimeException(e);
+              }
+            };
+
+        Function<Iterable<AppendRowsContext<DestinationT>>, RetryType> onError =
+            failedContexts -> {
+              handleAppendFailure(
+                  failedContexts,
+                  tableReference,
+                  shortTableId,
+                  appendClientHolder.get(),
+                  tryCreateTable,
+                  initializeContexts,
+                  clearClients,
+                  streamOffset,
+                  o);
               return RetryType.RETRY_ALL_OPERATIONS;
-            }
-
-            return RetryType.RETRY_ALL_OPERATIONS;
-          };
+            };
+        Consumer<AppendRowsContext<DestinationT>> onSuccess =
+            context -> handleAppendSuccess(context, shortTableId, appendClientHolder.get(), o);
+
+        BackOff backoff =
+            FluentBackoff.DEFAULT
+                .withInitialBackoff(Duration.standardSeconds(1))
+                .withMaxBackoff(Duration.standardMinutes(1))
+                .withMaxRetries(500)
+                .withThrottledTimeCounter(
+                    BigQuerySinkMetrics.throttledTimeCounter(
+                        BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM))
+                .backoff();
+        CreateRetryManagerResult<DestinationT> createRetryManagerResult;
+        do {
+          // Each ProtoRows object contains at most 1MB of rows.
+          // TODO: Push messageFromTableRow up to top level. That we we cans skip TableRow entirely
+          // if
+          // already proto or already schema.
+          Iterable<SplittingIterable.Value> messages =
+              new SplittingIterable(
+                  element.getValue(),
+                  splitSize,
+                  // Unknown field merger
+                  (bytes, tableRow) ->
+                      appendClientHolder.get().mergeNewFields(bytes, tableRow, ignoreUnknownValues),
+                  // Convert back to TableRow
+                  bytes -> appendClientHolder.get().toTableRow(bytes, Predicates.alwaysTrue()),
+                  // Failed rows consumer
+                  (failedRow, errorMessage) -> {
+                    o.get(failedRowsTag)
+                        .outputWithTimestamp(
+                            new BigQueryStorageApiInsertError(
+                                failedRow.getValue(), errorMessage, tableReference),
+                            failedRow.getTimestamp());
+                    rowsSentToFailedRowsCollection.inc();
+                    BigQuerySinkMetrics.appendRowsRowStatusCounter(
+                            BigQuerySinkMetrics.RowStatus.FAILED,
+                            BigQuerySinkMetrics.PAYLOAD_TOO_LARGE,
+                            shortTableId)
+                        .inc(1);
+                  },
+                  // Get the currently-known TableSchema hash
+                  () -> appendClientHolder.get().getTableSchemaHash(),
+                  () ->
+                      TableRowToStorageApiProto.wrapDescriptorProto(
+                          messageConverter.getDescriptor(false)),
+                  autoUpdateSchema,
+                  elementTs);
+
+          createRetryManagerResult =
+              createRetryManager(
+                  element.getKey(),
+                  messages,
+                  runOperation,
+                  onError,
+                  onSuccess,
+                  appendClientHolder.get(),
+                  tableReference);
+          if (createRetryManagerResult.getSchemaMismatchSeen()) {
+            // TODO: The call to updateSchemaFromTable will throttle the DoFn (both because of the
+            // RPC
+            // call and because
+            // the cache has a delay on refresh). We should update throttling counters here as well.
+            LOG.info("Schema out of date: refreshing table schema for {}", tableId);
+            // Force the message converter to get the schema again from the table.
+            messageConverter.updateSchemaFromTable();
+            // Close all RPC clients that were opened with the old descriptor. Clear the cache,
+            // forcing us to create a new append client with the updated descriptor.
+            appendClientHolder.invalidateAndReset();
+          }
+        } while (createRetryManagerResult.getSchemaMismatchSeen()
+            && BackOffUtils.next(Sleeper.DEFAULT, backoff));
+
+        // Output any rows that failed along they way.
+        createRetryManagerResult
+            .getFailedRows()
+            .forEach(
+                tv -> o.get(failedRowsTag).outputWithTimestamp(tv.getValue(), tv.getTimestamp()));
+        rowsSentToFailedRowsCollection.inc(createRetryManagerResult.getFailedRows().size());
+        BigQuerySinkMetrics.appendRowsRowStatusCounter(
+                BigQuerySinkMetrics.RowStatus.FAILED,
+                BigQuerySinkMetrics.PAYLOAD_TOO_LARGE,
+                shortTableId)
+            .inc(createRetryManagerResult.getFailedRows().size());
+
+        recordsAppended.inc(createRetryManagerResult.getRecordsAppended());
+        createRetryManagerResult.getHistogramValues().forEach(appendSizeDistribution::update);
+
+        Instant now = Instant.now();
+
+        RetryManager<AppendRowsResponse, AppendRowsContext<DestinationT>> retryManager =
+            Preconditions.checkStateNotNull(createRetryManagerResult.getRetryManager());
+        int numAppends = retryManager.getRemainingOperationCount();
+        Iterable<AppendRowsContext<DestinationT>> contexts = retryManager.getRemainingContexts();
+
+        if (numAppends > 0) {
+          initializeContexts.accept(contexts, false);
+          retryManager.run(true);
 
-      Consumer<AppendRowsContext<DestinationT>> onSuccess =
-          context -> {
-            AppendRowsResponse response = Preconditions.checkStateNotNull(context.getResult());
-            o.get(flushTag)
-                .output(
-                    KV.of(
-                        context.streamName,
-                        new Operation(
-                            context.offset + context.protoRows.getSerializedRowsCount() - 1,
-                            false)));
-            int flushedRows = context.protoRows.getSerializedRowsCount();
-            flushesScheduled.inc(flushedRows);
-            BigQuerySinkMetrics.reportSuccessfulRpcMetrics(
-                context, BigQuerySinkMetrics.RpcMethod.APPEND_ROWS, shortTableId);
-            BigQuerySinkMetrics.appendRowsRowStatusCounter(
-                    BigQuerySinkMetrics.RowStatus.SUCCESSFUL, BigQuerySinkMetrics.OK, shortTableId)
-                .inc(flushedRows);
-
-            if (successfulRowsTag != null) {
-              for (int i = 0; i < context.protoRows.getSerializedRowsCount(); ++i) {
-                ByteString protoBytes = context.protoRows.getSerializedRows(i);
-                org.joda.time.Instant timestamp = context.timestamps.get(i);
-                o.get(successfulRowsTag)
-                    .outputWithTimestamp(
-                        appendClientInfo.get().toTableRow(protoBytes, successfulRowsPredicate),
-                        timestamp);
+          appendSplitDistribution.update(numAppends);
+          if (autoUpdateSchema) {
+            @Nullable
+            StreamAppendClient streamAppendClient = appendClientHolder.getStreamAppendClient();
+            TableSchema originalSchema = appendClientHolder.get().getTableSchema();
+
+            @Nullable
+            TableSchema updatedSchemaReturned =
+                (streamAppendClient != null) ? streamAppendClient.getUpdatedSchema() : null;
+            // Update the table schema and clear the append client.
+            if (updatedSchemaReturned != null) {
+              Optional<TableSchema> newSchema =
+                  TableSchemaUpdateUtils.getUpdatedSchema(originalSchema, updatedSchemaReturned);
+              if (newSchema.isPresent()) {
+                APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
+                LOG.debug(
+                    "Fetched updated schema for table {}:\n\t{}", tableId, updatedSchemaReturned);

Review Comment:
   Yes, it makes the pinning logic simpler. the local AppendClientInfo will be unpinned at the end of the try block, and the next time processElement is called it will create a new append client.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3156297033


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo putAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = wrapWithPin(loader).call();

Review Comment:
   hmm, that's actually a good point. The only place we call it is immediately after a call to invalidate, but there's a race condition there where somebody else could sneak in.
   
   I don't think this could actually cause problems, since the pin owned by the DoFn is always released based on the instance variable (this.appendClientInfo), so even if we "lose" the cache entry we won't lose the pin. However to make things a bit clearer, I got rid of this method.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4436898886

   I don't quite understand the fix.
   
   On Tue, May 12, 2026 at 8:11 PM Yi Hu ***@***.***> wrote:
   
   > *Abacn* left a comment (apache/beam#38264)
   > <https://github.com/apache/beam/pull/38264#issuecomment-4436811897>
   >
   > It appears this change has caused periodic watermark stuck for long
   > running streaming jobs (commented internally). Tried a fix #38410
   > <https://github.com/apache/beam/pull/38410> but not sure if it fixes the
   > exact issue, For the upcoming release also prepared a revert #38479
   > <https://github.com/apache/beam/pull/38479> if hasn't figured out a way
   > to test it
   >
   > —
   > Reply to this email directly, view it on GitHub
   > <https://github.com/apache/beam/pull/38264#issuecomment-4436811897>, or
   > unsubscribe
   > <https://github.com/notifications/unsubscribe-auth/AFAYJVODDN7FPVAFUWDIY4T42PRUVAVCNFSM6AAAAACYBYX7MCVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHM2DIMZWHAYTCOBZG4>
   > .
   > Triage notifications on the go with GitHub Mobile for iOS
   > <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
   > or Android
   > <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.
   >
   > You are receiving this because you modified the open/close state.Message
   > ID: ***@***.***>
   >
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4340278800

   Assigning reviewers:
   
   R: @kennknowles for label java.
   
   
   Note: If you would like to opt out of this review, comment `assign to next reviewer`.
   
   Available commands:
   - `stop reviewer notifications` - opt out of the automated review tooling
   - `remind me after tests pass` - tag the comment author after tests pass
   - `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
   
   The PR bot will only process comments in the main thread (not review comments).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3155980812


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -282,7 +242,6 @@ static class AppendRowsContext<DestinationT>
       extends RetryManager.Operation.Context<AppendRowsResponse> {
     final ShardedKey<DestinationT> key;
     String streamName = "";
-    @Nullable StreamAppendClient client = null;

Review Comment:
   There's a comment in front of the try-with-resources block explaining this.
   
   Basically pins are ref counts. Since all calls happen in the context of processElements, all we really need is one pin for the entire function - there was never a need for per-RPC pins which was hard to follow. Instead we add one pin for the duration of the function, managed via a try-with-resources block. Since the cache holds onto objects longer, the cache also manages one additional pin per object (though that can be invalidated at any time)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4436811897

   It appears this change has caused periodic watermark stuck for long running streaming jobs (commented internally). Tried a fix #38410 but not sure if it fixes the exact issue, For the upcoming release also prepared a revert #38479 if hasn't figured out a way to test it


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor storageapi conversion [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4295787712

   Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "reuvenlax (via GitHub)" <gi...@apache.org>.
reuvenlax commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3155988295


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java:
##########
@@ -134,7 +143,18 @@ public AppendClientInfo withAppendClient(
   public void close() {
     BigQueryServices.StreamAppendClient client = getStreamAppendClient();
     if (client != null) {
-      getCloseAppendClient().accept(client);
+      try {
+        getCloseAppendClient().accept(client);
+      } catch (Exception e) {
+        // We ignore errors when closing clients.
+        String msg =
+            e
+                + "\n"
+                + Arrays.stream(e.getStackTrace())
+                    .map(StackTraceElement::toString)
+                    .collect(Collectors.joining("\n"));
+        LOG.warn("Caught exception whilw trying to close append client. Ignoring {}", msg);

Review Comment:
   done



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -671,475 +880,223 @@ public void process(
                 AppendClientInfo.of(
                         Preconditions.checkStateNotNull(tableSchema),
                         descriptor,
-                        // Make sure that the client is always closed in a different thread
-                        // to
-                        // avoid blocking.
-                        client ->
-                            runAsyncIgnoreFailure(
-                                closeWriterExecutor,
-                                () -> {
-                                  // Remove the pin that is "owned" by the cache.
-                                  client.unpin();
-                                  client.close();
-                                }))
+                        AutoCloseable::close)
                     .withAppendClient(
                         writeStreamService,
                         getOrCreateStream,
                         false,
                         defaultMissingValueInterpretation);
-            // This pin is "owned" by the cache.
-            Preconditions.checkStateNotNull(info.getStreamAppendClient()).pin();
             return info;
           };
 
-      AtomicReference<AppendClientInfo> appendClientInfo =
-          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(messageConverters.getAppendClientKey(element.getKey()));
-        appendClientInfo.set(
-            APPEND_CLIENTS.get(
-                messageConverters.getAppendClientKey(element.getKey()), getAppendClientInfo));
-      }
-
-      TableSchema updatedSchemaValue = updatedSchema.read();
-      if (autoUpdateSchema && updatedSchemaValue != null) {
-        if (appendClientInfo.get().hasSchemaChanged(updatedSchemaValue)) {
-          appendClientInfo.set(
-              AppendClientInfo.of(
-                  updatedSchemaValue, appendClientInfo.get().getCloseAppendClient(), false));
-          APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-          APPEND_CLIENTS.put(
-              messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
+      // The StreamWriter has two pins on it. The static cache holds a pin, as it continues to cache
+      // values after this
+      // method exits, so must hold the pin. The local AppendClientHolder also holds a pin, as the
+      // cache could in
+      // theory evict the object during execution and we want a pin held throughout the execution of
+      // this function.
+      try (AppendClientHolder appendClientHolder =
+          new AppendClientHolder(element.getKey(), getAppendClientInfo)) {
+        String currentStream = getOrCreateStream.get();
+        if (!currentStream.equals(appendClientHolder.get().getStreamName())) {
+          // Cached append client is inconsistent with persisted state. Throw away cached item and
+          // force it to be recreated.
+          appendClientHolder.invalidateAndReset();
         }
-      }
 
-      // 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) -> {
-            try {
-              if (isFailure) {
-                // Clear the stream name, forcing a new one to be created.
-                streamName.write("");
-              }
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              StreamAppendClient streamAppendClient =
-                  Preconditions.checkArgumentNotNull(
-                      appendClientInfo.get().getStreamAppendClient());
-              String streamNameRead = Preconditions.checkArgumentNotNull(streamName.read());
-              long currentOffset = Preconditions.checkArgumentNotNull(streamOffset.read());
-              for (AppendRowsContext<DestinationT> context : contexts) {
-                context.streamName = streamNameRead;
-                streamAppendClient.pin();
-                context.client = appendClientInfo.get().getStreamAppendClient();
-                context.offset = currentOffset;
-                ++context.tryIteration;
-                currentOffset = context.offset + context.protoRows.getSerializedRowsCount();
-              }
-              streamOffset.write(currentOffset);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
-
-      Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients =
-          (contexts) -> {
-            APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-            appendClientInfo.set(appendClientInfo.get().withNoAppendClient());
-            APPEND_CLIENTS.put(
-                messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
-            for (AppendRowsContext<DestinationT> context : contexts) {
-              if (context.client != null) {
-                // Unpin in a different thread, as it may execute a blocking close.
-                runAsyncIgnoreFailure(closeWriterExecutor, context.client::unpin);
-                context.client = null;
-              }
-            }
-          };
-
-      Function<AppendRowsContext<DestinationT>, ApiFuture<AppendRowsResponse>> runOperation =
-          context -> {
-            if (context.protoRows.getSerializedRowsCount() == 0) {
-              // This might happen if all rows in a batch failed and were sent to the failed-rows
-              // PCollection.
-              return ApiFutures.immediateFuture(AppendRowsResponse.newBuilder().build());
-            }
-            try {
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              return Preconditions.checkStateNotNull(appendClientInfo.get().getStreamAppendClient())
-                  .appendRows(context.offset, context.protoRows);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
+        TableSchema updatedSchemaValue = autoUpdateSchema ? updatedSchema.read() : null;
+        if (autoUpdateSchema && updatedSchemaValue != null) {
+          if (appendClientHolder.get().hasSchemaChanged(updatedSchemaValue)) {
+            appendClientHolder.invalidateAndReset();
+          }
+        }

Review Comment:
   done



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -487,6 +446,46 @@ public void onTeardown() {
       }
     }
 
+    // Holder for an AppendClientHolder. Maintains a pin on the client as long as it's active.
+    private class AppendClientHolder implements AutoCloseable {

Review Comment:
   Correct. Keeping one pin for the duration of processElements.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "ahmedabu98 (via GitHub)" <gi...@apache.org>.
ahmedabu98 commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3155947515


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo putAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = wrapWithPin(loader).call();

Review Comment:
   Yes, unpinning a client if it already exists for this key



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "ahmedabu98 (via GitHub)" <gi...@apache.org>.
ahmedabu98 commented on code in PR #38264:
URL: https://github.com/apache/beam/pull/38264#discussion_r3154653273


##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -487,6 +446,46 @@ public void onTeardown() {
       }
     }
 
+    // Holder for an AppendClientHolder. Maintains a pin on the client as long as it's active.
+    private class AppendClientHolder implements AutoCloseable {

Review Comment:
   Instead of pinning for each append context, we're using one AppendClientHolder that holds a pin on their behalf?



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -282,7 +242,6 @@ static class AppendRowsContext<DestinationT>
       extends RetryManager.Operation.Context<AppendRowsResponse> {
     final ShardedKey<DestinationT> key;
     String streamName = "";
-    @Nullable StreamAppendClient client = null;

Review Comment:
   Hmmm so we're no longer pinning the client for each AppenRowsContext? Can you briefly explain this change somewhere? maybe in the PR description?



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -671,475 +880,223 @@ public void process(
                 AppendClientInfo.of(
                         Preconditions.checkStateNotNull(tableSchema),
                         descriptor,
-                        // Make sure that the client is always closed in a different thread
-                        // to
-                        // avoid blocking.
-                        client ->
-                            runAsyncIgnoreFailure(
-                                closeWriterExecutor,
-                                () -> {
-                                  // Remove the pin that is "owned" by the cache.
-                                  client.unpin();
-                                  client.close();
-                                }))
+                        AutoCloseable::close)
                     .withAppendClient(
                         writeStreamService,
                         getOrCreateStream,
                         false,
                         defaultMissingValueInterpretation);
-            // This pin is "owned" by the cache.
-            Preconditions.checkStateNotNull(info.getStreamAppendClient()).pin();
             return info;
           };
 
-      AtomicReference<AppendClientInfo> appendClientInfo =
-          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(messageConverters.getAppendClientKey(element.getKey()));
-        appendClientInfo.set(
-            APPEND_CLIENTS.get(
-                messageConverters.getAppendClientKey(element.getKey()), getAppendClientInfo));
-      }
-
-      TableSchema updatedSchemaValue = updatedSchema.read();
-      if (autoUpdateSchema && updatedSchemaValue != null) {
-        if (appendClientInfo.get().hasSchemaChanged(updatedSchemaValue)) {
-          appendClientInfo.set(
-              AppendClientInfo.of(
-                  updatedSchemaValue, appendClientInfo.get().getCloseAppendClient(), false));
-          APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-          APPEND_CLIENTS.put(
-              messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
+      // The StreamWriter has two pins on it. The static cache holds a pin, as it continues to cache
+      // values after this
+      // method exits, so must hold the pin. The local AppendClientHolder also holds a pin, as the
+      // cache could in
+      // theory evict the object during execution and we want a pin held throughout the execution of
+      // this function.
+      try (AppendClientHolder appendClientHolder =
+          new AppendClientHolder(element.getKey(), getAppendClientInfo)) {
+        String currentStream = getOrCreateStream.get();
+        if (!currentStream.equals(appendClientHolder.get().getStreamName())) {
+          // Cached append client is inconsistent with persisted state. Throw away cached item and
+          // force it to be recreated.
+          appendClientHolder.invalidateAndReset();
         }
-      }
 
-      // 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) -> {
-            try {
-              if (isFailure) {
-                // Clear the stream name, forcing a new one to be created.
-                streamName.write("");
-              }
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              StreamAppendClient streamAppendClient =
-                  Preconditions.checkArgumentNotNull(
-                      appendClientInfo.get().getStreamAppendClient());
-              String streamNameRead = Preconditions.checkArgumentNotNull(streamName.read());
-              long currentOffset = Preconditions.checkArgumentNotNull(streamOffset.read());
-              for (AppendRowsContext<DestinationT> context : contexts) {
-                context.streamName = streamNameRead;
-                streamAppendClient.pin();
-                context.client = appendClientInfo.get().getStreamAppendClient();
-                context.offset = currentOffset;
-                ++context.tryIteration;
-                currentOffset = context.offset + context.protoRows.getSerializedRowsCount();
-              }
-              streamOffset.write(currentOffset);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
-
-      Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients =
-          (contexts) -> {
-            APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-            appendClientInfo.set(appendClientInfo.get().withNoAppendClient());
-            APPEND_CLIENTS.put(
-                messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
-            for (AppendRowsContext<DestinationT> context : contexts) {
-              if (context.client != null) {
-                // Unpin in a different thread, as it may execute a blocking close.
-                runAsyncIgnoreFailure(closeWriterExecutor, context.client::unpin);
-                context.client = null;
-              }
-            }
-          };
-
-      Function<AppendRowsContext<DestinationT>, ApiFuture<AppendRowsResponse>> runOperation =
-          context -> {
-            if (context.protoRows.getSerializedRowsCount() == 0) {
-              // This might happen if all rows in a batch failed and were sent to the failed-rows
-              // PCollection.
-              return ApiFutures.immediateFuture(AppendRowsResponse.newBuilder().build());
-            }
-            try {
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              return Preconditions.checkStateNotNull(appendClientInfo.get().getStreamAppendClient())
-                  .appendRows(context.offset, context.protoRows);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
+        TableSchema updatedSchemaValue = autoUpdateSchema ? updatedSchema.read() : null;
+        if (autoUpdateSchema && updatedSchemaValue != null) {
+          if (appendClientHolder.get().hasSchemaChanged(updatedSchemaValue)) {
+            appendClientHolder.invalidateAndReset();
+          }
+        }

Review Comment:
   nit: merge if conditions ?



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java:
##########
@@ -134,7 +143,18 @@ public AppendClientInfo withAppendClient(
   public void close() {
     BigQueryServices.StreamAppendClient client = getStreamAppendClient();
     if (client != null) {
-      getCloseAppendClient().accept(client);
+      try {
+        getCloseAppendClient().accept(client);
+      } catch (Exception e) {
+        // We ignore errors when closing clients.
+        String msg =
+            e
+                + "\n"
+                + Arrays.stream(e.getStackTrace())
+                    .map(StackTraceElement::toString)
+                    .collect(Collectors.joining("\n"));
+        LOG.warn("Caught exception whilw trying to close append client. Ignoring {}", msg);

Review Comment:
   Gemini has a point here. Why not use `LOG.warn("Caught exception while trying to close append client. Ignoring", e);`?



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientCache.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+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;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates the cache of {@link AppendClientInfo} objects and the synchronization protocol
+ * required to use them safely. The Guava cache object is thread-safe. However our protocol requires
+ * that client pin the StreamAppendClient after looking up the cache, and we must ensure that the
+ * cache is not accessed in between the lookup and the pin (any access of the cache could trigger
+ * element expiration).
+ */
+class AppendClientCache<KeyT extends @NonNull Object> {
+  private static final Logger LOG = LoggerFactory.getLogger(AppendClientCache.class);
+  private final ExecutorService closeWriterExecutor = Executors.newCachedThreadPool();
+
+  private final Cache<KeyT, AppendClientInfo> appendCache;
+
+  @SuppressWarnings({"FutureReturnValueIgnored"})
+  AppendClientCache(Duration expireAfterAccess) {
+    this.appendCache =
+        CacheBuilder.newBuilder()
+            .expireAfterAccess(expireAfterAccess.getMillis(), TimeUnit.MILLISECONDS)
+            .removalListener(
+                (RemovalNotification<KeyT, AppendClientInfo> removal) -> {
+                  LOG.info("Expiring append client for {}", removal.getKey());
+                  final @Nullable AppendClientInfo appendClientInfo = removal.getValue();
+                  if (appendClientInfo != null) {
+                    // Remove the pin owned by the cache itself. Since the client has not been
+                    // marked as closed, we
+                    // can call unpin in this thread without worrying about blocking the thread.
+                    appendClientInfo.unpinAppendClient(null);
+                    // Close the client in another thread to avoid blocking the main thread.
+                    closeWriterExecutor.submit(appendClientInfo::close);
+                  }
+                })
+            .build();
+  }
+
+  // The cache itself always own one pin on the object. This Callable is always used to ensure that
+  // the cache
+  // adds a pin before loading a value.
+  private static Callable<AppendClientInfo> wrapWithPin(Callable<AppendClientInfo> loader) {
+    return () -> {
+      AppendClientInfo client = loader.call();
+      client.pinAppendClient();
+      return client;
+    };
+  }
+
+  /**
+   * Atomically get an append client from the cache and add a pin. This pin is owned by the client,
+   * which has the responsibility of removing it. If the client is not in the cache, loader will be
+   * used to load the client; in this case an additional pin will be added owned by the cache,
+   * removed when the item is evicted.
+   */
+  public AppendClientInfo getAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = appendCache.get(key, wrapWithPin(loader));
+      info.pinAppendClient();
+      return info;
+    }
+  }
+
+  public AppendClientInfo putAndPin(KeyT key, Callable<AppendClientInfo> loader) throws Exception {
+    synchronized (this) {
+      AppendClientInfo info = wrapWithPin(loader).call();

Review Comment:
   Do we need to unpin any existing AppendClients?



##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -671,475 +880,223 @@ public void process(
                 AppendClientInfo.of(
                         Preconditions.checkStateNotNull(tableSchema),
                         descriptor,
-                        // Make sure that the client is always closed in a different thread
-                        // to
-                        // avoid blocking.
-                        client ->
-                            runAsyncIgnoreFailure(
-                                closeWriterExecutor,
-                                () -> {
-                                  // Remove the pin that is "owned" by the cache.
-                                  client.unpin();
-                                  client.close();
-                                }))
+                        AutoCloseable::close)
                     .withAppendClient(
                         writeStreamService,
                         getOrCreateStream,
                         false,
                         defaultMissingValueInterpretation);
-            // This pin is "owned" by the cache.
-            Preconditions.checkStateNotNull(info.getStreamAppendClient()).pin();
             return info;
           };
 
-      AtomicReference<AppendClientInfo> appendClientInfo =
-          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(messageConverters.getAppendClientKey(element.getKey()));
-        appendClientInfo.set(
-            APPEND_CLIENTS.get(
-                messageConverters.getAppendClientKey(element.getKey()), getAppendClientInfo));
-      }
-
-      TableSchema updatedSchemaValue = updatedSchema.read();
-      if (autoUpdateSchema && updatedSchemaValue != null) {
-        if (appendClientInfo.get().hasSchemaChanged(updatedSchemaValue)) {
-          appendClientInfo.set(
-              AppendClientInfo.of(
-                  updatedSchemaValue, appendClientInfo.get().getCloseAppendClient(), false));
-          APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-          APPEND_CLIENTS.put(
-              messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
+      // The StreamWriter has two pins on it. The static cache holds a pin, as it continues to cache
+      // values after this
+      // method exits, so must hold the pin. The local AppendClientHolder also holds a pin, as the
+      // cache could in
+      // theory evict the object during execution and we want a pin held throughout the execution of
+      // this function.
+      try (AppendClientHolder appendClientHolder =
+          new AppendClientHolder(element.getKey(), getAppendClientInfo)) {
+        String currentStream = getOrCreateStream.get();
+        if (!currentStream.equals(appendClientHolder.get().getStreamName())) {
+          // Cached append client is inconsistent with persisted state. Throw away cached item and
+          // force it to be recreated.
+          appendClientHolder.invalidateAndReset();
         }
-      }
 
-      // 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) -> {
-            try {
-              if (isFailure) {
-                // Clear the stream name, forcing a new one to be created.
-                streamName.write("");
-              }
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              StreamAppendClient streamAppendClient =
-                  Preconditions.checkArgumentNotNull(
-                      appendClientInfo.get().getStreamAppendClient());
-              String streamNameRead = Preconditions.checkArgumentNotNull(streamName.read());
-              long currentOffset = Preconditions.checkArgumentNotNull(streamOffset.read());
-              for (AppendRowsContext<DestinationT> context : contexts) {
-                context.streamName = streamNameRead;
-                streamAppendClient.pin();
-                context.client = appendClientInfo.get().getStreamAppendClient();
-                context.offset = currentOffset;
-                ++context.tryIteration;
-                currentOffset = context.offset + context.protoRows.getSerializedRowsCount();
-              }
-              streamOffset.write(currentOffset);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
-
-      Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients =
-          (contexts) -> {
-            APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
-            appendClientInfo.set(appendClientInfo.get().withNoAppendClient());
-            APPEND_CLIENTS.put(
-                messageConverters.getAppendClientKey(element.getKey()), appendClientInfo.get());
-            for (AppendRowsContext<DestinationT> context : contexts) {
-              if (context.client != null) {
-                // Unpin in a different thread, as it may execute a blocking close.
-                runAsyncIgnoreFailure(closeWriterExecutor, context.client::unpin);
-                context.client = null;
-              }
-            }
-          };
-
-      Function<AppendRowsContext<DestinationT>, ApiFuture<AppendRowsResponse>> runOperation =
-          context -> {
-            if (context.protoRows.getSerializedRowsCount() == 0) {
-              // This might happen if all rows in a batch failed and were sent to the failed-rows
-              // PCollection.
-              return ApiFutures.immediateFuture(AppendRowsResponse.newBuilder().build());
-            }
-            try {
-              appendClientInfo.set(
-                  appendClientInfo
-                      .get()
-                      .withAppendClient(
-                          writeStreamService,
-                          getOrCreateStream,
-                          false,
-                          defaultMissingValueInterpretation));
-              return Preconditions.checkStateNotNull(appendClientInfo.get().getStreamAppendClient())
-                  .appendRows(context.offset, context.protoRows);
-            } catch (Exception e) {
-              throw new RuntimeException(e);
-            }
-          };
+        TableSchema updatedSchemaValue = autoUpdateSchema ? updatedSchema.read() : null;
+        if (autoUpdateSchema && updatedSchemaValue != null) {
+          if (appendClientHolder.get().hasSchemaChanged(updatedSchemaValue)) {
+            appendClientHolder.invalidateAndReset();
+          }
+        }
 
-      Function<Iterable<AppendRowsContext<DestinationT>>, RetryType> onError =
-          failedContexts -> {
-            // The first context is always the one that fails.
-            AppendRowsContext<DestinationT> failedContext =
-                Preconditions.checkStateNotNull(Iterables.getFirst(failedContexts, null));
-            BigQuerySinkMetrics.reportFailedRPCMetrics(
-                failedContext, BigQuerySinkMetrics.RpcMethod.APPEND_ROWS, shortTableId);
-            String errorCode =
-                BigQuerySinkMetrics.throwableToGRPCCodeString(failedContext.getError());
-
-            // AppendSerializationError means that BigQuery detected errors on individual rows, e.g.
-            // a row not conforming
-            // to bigQuery invariants. These errors are persistent, so we redirect those rows to the
-            // failedInserts
-            // PCollection, and retry with the remaining rows.
-            if (failedContext.getError() != null
-                && failedContext.getError() instanceof Exceptions.AppendSerializationError) {
-              Exceptions.AppendSerializationError error =
-                  Preconditions.checkArgumentNotNull(
-                      (Exceptions.AppendSerializationError) failedContext.getError());
-
-              Set<Integer> failedRowIndices = error.getRowIndexToErrorMessage().keySet();
-              for (int failedIndex : failedRowIndices) {
-                // Convert the message to a TableRow and send it to the failedRows collection.
-                TableRow failedRow = failedContext.failsafeTableRows.get(failedIndex);
-                if (failedRow == null) {
-                  // TODO: MAKE SURE WE USE UPDATED DESCRIPTOR
-                  ByteString protoBytes = failedContext.protoRows.getSerializedRows(failedIndex);
-                  failedRow =
-                      appendClientInfo.get().toTableRow(protoBytes, Predicates.alwaysTrue());
+        // 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) -> {
+              try {
+                if (isFailure) {
+                  // Clear the stream name, forcing a new one to be created.
+                  streamName.write("");
                 }
-                org.joda.time.Instant timestamp = failedContext.timestamps.get(failedIndex);
-                o.get(failedRowsTag)
-                    .outputWithTimestamp(
-                        new BigQueryStorageApiInsertError(
-                            failedRow,
-                            error.getRowIndexToErrorMessage().get(failedIndex),
-                            tableReference),
-                        timestamp);
-              }
-              int failedRows = failedRowIndices.size();
-              rowsSentToFailedRowsCollection.inc(failedRows);
-              BigQuerySinkMetrics.appendRowsRowStatusCounter(
-                      BigQuerySinkMetrics.RowStatus.FAILED, errorCode, shortTableId)
-                  .inc(failedRows);
-
-              // Remove the failed row from the payload, so we retry the batch without the failed
-              // rows.
-              ProtoRows.Builder retryRows = ProtoRows.newBuilder();
-              @Nullable List<org.joda.time.Instant> timestamps = Lists.newArrayList();
-              for (int i = 0; i < failedContext.protoRows.getSerializedRowsCount(); ++i) {
-                if (!failedRowIndices.contains(i)) {
-                  ByteString rowBytes = failedContext.protoRows.getSerializedRows(i);
-                  retryRows.addSerializedRows(rowBytes);
-                  timestamps.add(failedContext.timestamps.get(i));
+                String streamNameRead = Preconditions.checkArgumentNotNull(streamName.read());
+                long currentOffset = Preconditions.checkArgumentNotNull(streamOffset.read());
+                for (AppendRowsContext<DestinationT> context : contexts) {
+                  context.streamName = streamNameRead;
+                  context.offset = currentOffset;
+                  ++context.tryIteration;
+                  currentOffset = context.offset + context.protoRows.getSerializedRowsCount();
                 }
+                streamOffset.write(currentOffset);
+              } catch (Exception e) {
+                throw new RuntimeException(e);
               }
-              failedContext.protoRows = retryRows.build();
-              failedContext.timestamps = timestamps;
-              int retriedRows = failedContext.protoRows.getSerializedRowsCount();
-              BigQuerySinkMetrics.appendRowsRowStatusCounter(
-                      BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableId)
-                  .inc(retriedRows);
-
-              // Since we removed rows, we need to update the insert offsets for all remaining rows.
-              long offset = failedContext.offset;
-              for (AppendRowsContext<DestinationT> context : failedContexts) {
-                context.offset = offset;
-                offset += context.protoRows.getSerializedRowsCount();
+            };
+
+        Consumer<Iterable<AppendRowsContext<DestinationT>>> clearClients =
+            (contexts) -> {
+              try {
+                appendClientHolder.invalidateAndReset();
+              } catch (Exception e) {
+                throw new RuntimeException(e);
               }
-              streamOffset.write(offset);
-              return RetryType.RETRY_ALL_OPERATIONS;
-            }
-
-            Throwable error = Preconditions.checkStateNotNull(failedContext.getError());
-
-            Status.Code statusCode = Status.fromThrowable(error).getCode();
-
-            // This means that the offset we have stored does not match the current end of
-            // the stream in the Storage API. Usually this happens because a crash or a bundle
-            // failure
-            // happened after an append but before the worker could checkpoint it's
-            // state. The records that were appended in a failed bundle will be retried,
-            // meaning that the unflushed tail of the stream must be discarded to prevent
-            // duplicates.
-            boolean offsetMismatch =
-                statusCode.equals(Code.OUT_OF_RANGE) || statusCode.equals(Code.ALREADY_EXISTS);
-
-            boolean quotaError = statusCode.equals(Code.RESOURCE_EXHAUSTED);
-            if (!offsetMismatch) {
-              // Don't log errors for expected offset mismatch. These will be logged as warnings
-              // below.
-              LOG.error(
-                  "Got error {} closing {}", failedContext.getError(), failedContext.streamName);
-            }
-
-            try {
-              // TODO: Only do this on explicit NOT_FOUND errors once BigQuery reliably produces
-              // them.
-              tryCreateTable.call();
-            } catch (Exception e) {
-              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(
-                    BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableId)
-                .inc(retriedRows);
-
-            // Schema mismatched exceptions can happen if the table was recently updated. Since
-            // vortex caches schemas
-            // we might see the new schema before vortex does. In this case, we simply need to
-            // retry.
-            Exceptions.@Nullable StorageException storageException =
-                Exceptions.toStorageException(error);
-            boolean schemaMismatchError =
-                (storageException instanceof Exceptions.SchemaMismatchedException);
-            if (!schemaMismatchError) {
-              // There's no special error code for missing required fields, and that can also happen
-              // due to vortex
-              // being delayed at seeing a new schema. We're forced to parse the description to
-              // determine that this has happened.
-              // TODO: Vortex team to introduce a special storage error code for this, so we don't
-              // have to parse
-              // descriptions.
-              Status status = Status.fromThrowable(error);
-              if (status.getCode() == Code.INVALID_ARGUMENT) {
-                String description = status.getDescription();
-                schemaMismatchError =
-                    description != null && description.contains("incompatible fields");
+            };
+
+        Function<AppendRowsContext<DestinationT>, ApiFuture<AppendRowsResponse>> runOperation =
+            context -> {
+              if (context.protoRows.getSerializedRowsCount() == 0) {
+                // This might happen if all rows in a batch failed and were sent to the failed-rows
+                // PCollection.
+                return ApiFutures.immediateFuture(AppendRowsResponse.newBuilder().build());
               }
-            }
-            if (schemaMismatchError) {
-              LOG.info(
-                  "Vortex failed stream open due to incompatible fields. This is likely because the BigTable "
-                      + "schema was recently updated and Vortex hasn't noticed yet, so retrying. error {}",
-                  Preconditions.checkStateNotNull(error).toString());
-            }
-
-            boolean explicitStreamFinalized =
-                failedContext.getError() instanceof StreamFinalizedException;
-            // This implies that the stream doesn't exist or has already been finalized. In this
-            // case we have no choice but to create a new stream.
-            boolean streamDoesNotExist =
-                explicitStreamFinalized
-                    || statusCode.equals(Code.INVALID_ARGUMENT)
-                    || statusCode.equals(Code.NOT_FOUND)
-                    || statusCode.equals(Code.FAILED_PRECONDITION);
-            streamDoesNotExist = streamDoesNotExist && !schemaMismatchError;
-
-            if (offsetMismatch || streamDoesNotExist) {
-              appendOffsetFailures.inc();
-              LOG.warn(
-                  "Append to {} failed. Will retry with a new stream",
-                  failedContext,
-                  failedContext.getError());
-              // 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)));
-              // Reinitialize all contexts with the new stream and new offsets.
-              initializeContexts.accept(failedContexts, true);
-
-              // Offset failures imply that all subsequent parallel appends will also fail.
-              // Retry them all.
+              try {
+                return appendClientHolder
+                    .getStreamAppendClient()
+                    .appendRows(context.offset, context.protoRows);
+              } catch (Exception e) {
+                throw new RuntimeException(e);
+              }
+            };
+
+        Function<Iterable<AppendRowsContext<DestinationT>>, RetryType> onError =
+            failedContexts -> {
+              handleAppendFailure(
+                  failedContexts,
+                  tableReference,
+                  shortTableId,
+                  appendClientHolder.get(),
+                  tryCreateTable,
+                  initializeContexts,
+                  clearClients,
+                  streamOffset,
+                  o);
               return RetryType.RETRY_ALL_OPERATIONS;
-            }
-
-            return RetryType.RETRY_ALL_OPERATIONS;
-          };
+            };
+        Consumer<AppendRowsContext<DestinationT>> onSuccess =
+            context -> handleAppendSuccess(context, shortTableId, appendClientHolder.get(), o);
+
+        BackOff backoff =
+            FluentBackoff.DEFAULT
+                .withInitialBackoff(Duration.standardSeconds(1))
+                .withMaxBackoff(Duration.standardMinutes(1))
+                .withMaxRetries(500)
+                .withThrottledTimeCounter(
+                    BigQuerySinkMetrics.throttledTimeCounter(
+                        BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM))
+                .backoff();
+        CreateRetryManagerResult<DestinationT> createRetryManagerResult;
+        do {
+          // Each ProtoRows object contains at most 1MB of rows.
+          // TODO: Push messageFromTableRow up to top level. That we we cans skip TableRow entirely
+          // if
+          // already proto or already schema.
+          Iterable<SplittingIterable.Value> messages =
+              new SplittingIterable(
+                  element.getValue(),
+                  splitSize,
+                  // Unknown field merger
+                  (bytes, tableRow) ->
+                      appendClientHolder.get().mergeNewFields(bytes, tableRow, ignoreUnknownValues),
+                  // Convert back to TableRow
+                  bytes -> appendClientHolder.get().toTableRow(bytes, Predicates.alwaysTrue()),
+                  // Failed rows consumer
+                  (failedRow, errorMessage) -> {
+                    o.get(failedRowsTag)
+                        .outputWithTimestamp(
+                            new BigQueryStorageApiInsertError(
+                                failedRow.getValue(), errorMessage, tableReference),
+                            failedRow.getTimestamp());
+                    rowsSentToFailedRowsCollection.inc();
+                    BigQuerySinkMetrics.appendRowsRowStatusCounter(
+                            BigQuerySinkMetrics.RowStatus.FAILED,
+                            BigQuerySinkMetrics.PAYLOAD_TOO_LARGE,
+                            shortTableId)
+                        .inc(1);
+                  },
+                  // Get the currently-known TableSchema hash
+                  () -> appendClientHolder.get().getTableSchemaHash(),
+                  () ->
+                      TableRowToStorageApiProto.wrapDescriptorProto(
+                          messageConverter.getDescriptor(false)),
+                  autoUpdateSchema,
+                  elementTs);
+
+          createRetryManagerResult =
+              createRetryManager(
+                  element.getKey(),
+                  messages,
+                  runOperation,
+                  onError,
+                  onSuccess,
+                  appendClientHolder.get(),
+                  tableReference);
+          if (createRetryManagerResult.getSchemaMismatchSeen()) {
+            // TODO: The call to updateSchemaFromTable will throttle the DoFn (both because of the
+            // RPC
+            // call and because
+            // the cache has a delay on refresh). We should update throttling counters here as well.
+            LOG.info("Schema out of date: refreshing table schema for {}", tableId);
+            // Force the message converter to get the schema again from the table.
+            messageConverter.updateSchemaFromTable();
+            // Close all RPC clients that were opened with the old descriptor. Clear the cache,
+            // forcing us to create a new append client with the updated descriptor.
+            appendClientHolder.invalidateAndReset();
+          }
+        } while (createRetryManagerResult.getSchemaMismatchSeen()
+            && BackOffUtils.next(Sleeper.DEFAULT, backoff));
+
+        // Output any rows that failed along they way.
+        createRetryManagerResult
+            .getFailedRows()
+            .forEach(
+                tv -> o.get(failedRowsTag).outputWithTimestamp(tv.getValue(), tv.getTimestamp()));
+        rowsSentToFailedRowsCollection.inc(createRetryManagerResult.getFailedRows().size());
+        BigQuerySinkMetrics.appendRowsRowStatusCounter(
+                BigQuerySinkMetrics.RowStatus.FAILED,
+                BigQuerySinkMetrics.PAYLOAD_TOO_LARGE,
+                shortTableId)
+            .inc(createRetryManagerResult.getFailedRows().size());
+
+        recordsAppended.inc(createRetryManagerResult.getRecordsAppended());
+        createRetryManagerResult.getHistogramValues().forEach(appendSizeDistribution::update);
+
+        Instant now = Instant.now();
+
+        RetryManager<AppendRowsResponse, AppendRowsContext<DestinationT>> retryManager =
+            Preconditions.checkStateNotNull(createRetryManagerResult.getRetryManager());
+        int numAppends = retryManager.getRemainingOperationCount();
+        Iterable<AppendRowsContext<DestinationT>> contexts = retryManager.getRemainingContexts();
+
+        if (numAppends > 0) {
+          initializeContexts.accept(contexts, false);
+          retryManager.run(true);
 
-      Consumer<AppendRowsContext<DestinationT>> onSuccess =
-          context -> {
-            AppendRowsResponse response = Preconditions.checkStateNotNull(context.getResult());
-            o.get(flushTag)
-                .output(
-                    KV.of(
-                        context.streamName,
-                        new Operation(
-                            context.offset + context.protoRows.getSerializedRowsCount() - 1,
-                            false)));
-            int flushedRows = context.protoRows.getSerializedRowsCount();
-            flushesScheduled.inc(flushedRows);
-            BigQuerySinkMetrics.reportSuccessfulRpcMetrics(
-                context, BigQuerySinkMetrics.RpcMethod.APPEND_ROWS, shortTableId);
-            BigQuerySinkMetrics.appendRowsRowStatusCounter(
-                    BigQuerySinkMetrics.RowStatus.SUCCESSFUL, BigQuerySinkMetrics.OK, shortTableId)
-                .inc(flushedRows);
-
-            if (successfulRowsTag != null) {
-              for (int i = 0; i < context.protoRows.getSerializedRowsCount(); ++i) {
-                ByteString protoBytes = context.protoRows.getSerializedRows(i);
-                org.joda.time.Instant timestamp = context.timestamps.get(i);
-                o.get(successfulRowsTag)
-                    .outputWithTimestamp(
-                        appendClientInfo.get().toTableRow(protoBytes, successfulRowsPredicate),
-                        timestamp);
+          appendSplitDistribution.update(numAppends);
+          if (autoUpdateSchema) {
+            @Nullable
+            StreamAppendClient streamAppendClient = appendClientHolder.getStreamAppendClient();
+            TableSchema originalSchema = appendClientHolder.get().getTableSchema();
+
+            @Nullable
+            TableSchema updatedSchemaReturned =
+                (streamAppendClient != null) ? streamAppendClient.getUpdatedSchema() : null;
+            // Update the table schema and clear the append client.
+            if (updatedSchemaReturned != null) {
+              Optional<TableSchema> newSchema =
+                  TableSchemaUpdateUtils.getUpdatedSchema(originalSchema, updatedSchemaReturned);
+              if (newSchema.isPresent()) {
+                APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey()));
+                LOG.debug(
+                    "Fetched updated schema for table {}:\n\t{}", tableId, updatedSchemaReturned);

Review Comment:
   I noticed we removed `APPEND_CLIENTS.put(...)` from here. Are we relying on the next process call to create the append client for `element.getKey()` ?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor StorageAPI BigQuery sink [beam]

Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4329426995

   ## Summary of Changes
   
   This pull request refactors the StorageAPI BigQuery sink to improve the maintainability and reliability of its caching and client-pinning infrastructure. By consolidating complex, distributed logic into a dedicated cache manager, the changes reduce code duplication and simplify the synchronization protocol required for safe append-client management. These improvements lead to a more robust system for handling stream lifecycle and resource cleanup.
   
   ### Highlights
   
   * **Consolidated Caching Logic**: Introduced a new `AppendClientCache` class to centralize and simplify the caching and pinning logic previously duplicated across `StorageApiWriteShardedRecords` and `StorageApiWriteUnshardedRecords`.
   * **Improved Pinning Protocol**: Refactored the append-client pinning mechanism to be more robust and easier to maintain, ensuring safe synchronization between cache lookups and client usage.
   * **Error Handling and Cleanup**: Enhanced error handling during client closure and unpinning, including better support for asynchronous cleanup to avoid blocking main processing threads.
   
   🧠 **New Feature in Public Preview:** You can now enable **Memory** to help **Gemini Code Assist** learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. **Click [here](https://codeassist.google/code-review/login) to enable Memory in your admin console.**
   
   <details>
   <summary><b>Activity</b></summary>
   
   * The pull request was created by reuvenlax.
   * Automated checks initially failed, preventing review requests.
   * The author requested a summary and review via automated tools.
   </details>
   
   
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Refactor storageapi conversion [beam]

Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #38264:
URL: https://github.com/apache/beam/pull/38264#issuecomment-4293385899

   ## Summary of Changes
   
   Hello, I'm Gemini Code Assist[^1]! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
   
   This pull request introduces robust support for automatic schema evolution in the BigQuery Storage Write API. It enables pipelines to handle schema changes dynamically by buffering records that encounter schema mismatches, patching the BigQuery table schema, and then retrying the buffered records. This change significantly improves the resilience of streaming pipelines writing to BigQuery.
   
   ### Highlights
   
   * **Storage API Schema Evolution**: Implemented support for automatic BigQuery table schema updates when using the Storage Write API, allowing for field additions and relaxations.
   * **Schema Buffering and Patching**: Introduced new DoFns (ConvertMessagesDoFn, PatchTableSchemaDoFn, SchemaUpdateHoldingFn) to buffer records that fail schema validation until the table schema is successfully updated.
   * **Append Client Management**: Refactored AppendClientCache to improve thread safety and resource management, including better handling of pinned clients and asynchronous closing.
   * **Schema Hash Verification**: Added schema hashing to StorageApiWritePayload to detect when a client's view of the schema is out of date compared to the BigQuery table.
   
   🧠 **New Feature in Public Preview:** You can now enable **Memory** to help **Gemini Code Assist** learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. **Click [here](https://codeassist.google/code-review/login) to enable Memory in your admin console.**
   
   <details>
   <summary><b>Using Gemini Code Assist</b></summary>
   <br>
   
   The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
   
   <b>Invoking Gemini</b>
   
   You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page.
   
   Feature | Command | Description
   --- | --- | ---
   Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
   Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
   Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments.
   Help | `/gemini help` | Displays a list of available commands.
   
   <b>Customization</b>
   
   To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
   
   <b>Limitations & Feedback</b>
   
   Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up [here](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
   </details>
   
   
   [^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org