fix: retry idle timeout (#4209)

SeqFileWriter#write() which write to sequence file buffer could be slow. And the latency depends on the size of the data. Watchdog could throw idle timeout exceptions and this shouldn't bubble up to the pipeline. Retrying idle timeout exception.

Author: mutianf

bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableOptionsFactory.java

@@ -17,6 +17,7 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalExtensionOnly;
+import com.google.common.annotations.VisibleForTesting;
 
 /**
  * Define {@link org.apache.hadoop.conf.Configuration} names for setting {@link
@@ -332,4 +333,8 @@ public class BigtableOptionsFactory {
    */
   public static final String BIGTABLE_ENABLE_BULK_MUTATION_FLOW_CONTROL =
       "google.bigtable.enable.bulk.mutation.flow.control";
+
+  /** Override idle timeout, for testing only. */
+  @VisibleForTesting
+  public static final String BIGTABLE_TEST_IDLE_TIMEOUT_MS = "google.bigtable.idle.timeout.ms";
 }

bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/wrappers/veneer/BigtableHBaseVeneerSettings.java

@@ -43,6 +43,7 @@
 import static com.google.cloud.bigtable.hbase.BigtableOptionsFactory.BIGTABLE_SERVICE_ACCOUNT_JSON_KEYFILE_LOCATION_KEY;
 import static com.google.cloud.bigtable.hbase.BigtableOptionsFactory.BIGTABLE_SERVICE_ACCOUNT_JSON_VALUE_KEY;
 import static com.google.cloud.bigtable.hbase.BigtableOptionsFactory.BIGTABLE_SERVICE_ACCOUNT_P12_KEYFILE_LOCATION_KEY;
+import static com.google.cloud.bigtable.hbase.BigtableOptionsFactory.BIGTABLE_TEST_IDLE_TIMEOUT_MS;
 import static com.google.cloud.bigtable.hbase.BigtableOptionsFactory.BIGTABLE_USE_BATCH;
 import static com.google.cloud.bigtable.hbase.BigtableOptionsFactory.BIGTABLE_USE_CACHED_DATA_CHANNEL_POOL;
 import static com.google.cloud.bigtable.hbase.BigtableOptionsFactory.BIGTABLE_USE_PLAINTEXT_NEGOTIATION;
@@ -728,6 +729,11 @@ private void configureReadRowsSettings(
           .retrySettings()
           .setTotalTimeout(operationTimeouts.getOperationTimeout().get());
     }
+
+    String idleTimeout = configuration.get(BIGTABLE_TEST_IDLE_TIMEOUT_MS);
+    if (idleTimeout != null) {
+      readRowsSettings.setIdleTimeout(Duration.ofMillis(Long.parseLong(idleTimeout)));
+    }
   }
 
   private void configureRetryableCallSettings(

bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/TemplateUtils.java

@@ -87,13 +87,15 @@ public static CloudBigtableScanConfiguration buildExportConfig(ExportOptions opt
             .withAppProfileId(options.getBigtableAppProfileId())
             .withConfiguration(
                 BigtableOptionsFactory.CUSTOM_USER_AGENT_KEY, "SequenceFileExportJob")
+            .withConfiguration(
+                CloudBigtableIO.Reader.RETRY_IDLE_TIMEOUT,
+                String.valueOf(options.getRetryIdleTimeout()))
             .withScan(
                 new ScanValueProvider(
                     options.getBigtableStartRow(),
                     options.getBigtableStopRow(),
                     options.getBigtableMaxVersions(),
                     options.getBigtableFilter()));
-
     return configBuilder.build();
   }
 }

bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ExportJob.java

@@ -168,6 +168,13 @@ public interface ExportOptions extends GcpOptions {
 
     @SuppressWarnings("unused")
     void setWait(boolean wait);
+
+    @Description("Get if idle timeout is retried.")
+    @Default.Boolean(true)
+    boolean getRetryIdleTimeout();
+
+    @SuppressWarnings("unused")
+    void setRetryIdleTimeout(boolean retryIdleTimeout);
   }
 
   public static void main(String[] args) {

bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java

@@ -19,6 +19,7 @@
 
 import com.google.bigtable.repackaged.com.google.api.core.InternalApi;
 import com.google.bigtable.repackaged.com.google.api.core.InternalExtensionOnly;
+import com.google.bigtable.repackaged.com.google.api.gax.rpc.WatchdogTimeoutException;
 import com.google.bigtable.repackaged.com.google.bigtable.v2.ReadRowsRequest;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.bigtable.repackaged.com.google.common.annotations.VisibleForTesting;
@@ -32,11 +33,13 @@
 import java.io.ObjectOutputStream;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
@@ -647,6 +650,8 @@ Object readResolve() {
   /** Reads rows for a specific {@link Table}, usually filtered by a {@link Scan}. */
   @VisibleForTesting
   static class Reader extends BoundedReader<Result> {
+    static final String RETRY_IDLE_TIMEOUT = "google.cloud.bigtable.retry.idle.timeout";
+
     private static final Logger READER_LOG = LoggerFactory.getLogger(Reader.class);
 
     private CloudBigtableIO.AbstractSource source;
@@ -657,6 +662,9 @@ static class Reader extends BoundedReader<Result> {
     protected long workStart;
     private final AtomicLong rowsRead = new AtomicLong();
     private final ByteKeyRangeTracker rangeTracker;
+    private transient Result lastScannedRow;
+
+    private final AtomicInteger attempt = new AtomicInteger(3);
 
     @VisibleForTesting
     Reader(CloudBigtableIO.AbstractSource source) {
@@ -690,7 +698,37 @@ void initializeScanner() throws IOException {
     /** Calls {@link ResultScanner#next()}. */
     @Override
     public boolean advance() throws IOException {
+      try {
+        boolean hasMore = tryAdvance();
+        // reset attempt after a success read
+        attempt.set(3);
+        return hasMore;
+      } catch (Throwable e) {
+        // if retry idle timeout is disabled, throw the exception
+        if (!source.getConfiguration().toHBaseConfig().getBoolean(RETRY_IDLE_TIMEOUT, true)) {
+          throw e;
+        }
+        // Exception is not idle timeout, throw it
+        Throwable exception = findCause(e, WatchdogTimeoutException.class);
+        if (exception == null) {
+          throw e;
+        }
+        if (exception.getMessage() == null || !exception.getMessage().contains("idle")) {
+          throw e;
+        }
+        // Run out ot retry attempt, throw the exception
+        if (attempt.decrementAndGet() <= 0) {
+          throw e;
+        }
+        READER_LOG.warn("got idle timeout exception, will try to reset the scanner and retry", e);
+        resetScanner();
+        return tryAdvance();
+      }
+    }
+
+    private boolean tryAdvance() throws IOException {
       Result row = scanner.next();
+      lastScannedRow = row;
       if (row != null && rangeTracker.tryReturnRecordAt(true, ByteKey.copyFrom(row.getRow()))) {
         current = row;
         rowsRead.addAndGet(1l);
@@ -702,6 +740,40 @@ public boolean advance() throws IOException {
       }
     }
 
+    private void resetScanner() throws IOException {
+      CloudBigtableScanConfiguration scanConfiguration = source.getConfiguration();
+      Scan scan;
+      if (lastScannedRow != null) {
+        byte[] rowKey = lastScannedRow.getRow();
+        // ScanConfiguration always gets start key and end key from the RowRange, and it expects
+        // start key to be inclusive and end key to be exclusive.
+        byte[] newStartKey = Arrays.copyOf(rowKey, rowKey.length + 1);
+        scan =
+            scanConfiguration
+                .toBuilder()
+                .withKeys(newStartKey, scanConfiguration.getStopRow())
+                .build()
+                .getScanValueProvider()
+                .get();
+      } else {
+        scan = scanConfiguration.getScanValueProvider().get();
+        READER_LOG.info("last scanned row key is null, haven't read any row yet");
+      }
+
+      scanner =
+          connection
+              .getTable(TableName.valueOf(source.getConfiguration().getTableId()))
+              .getScanner(scan);
+    }
+
+    static Throwable findCause(Throwable e, Class<? extends Throwable> cls) {
+      Throwable throwable = e;
+      while (throwable != null && !cls.isAssignableFrom(e.getClass())) {
+        throwable = throwable.getCause();
+      }
+      return throwable;
+    }
+
     @Override
     public final Double getFractionConsumed() {
       if (rangeTracker.isDone()) {