org.apache.beam.sdk.util.BackOff Java Examples

@Override
public boolean advance() throws IOException {
  if (elemsRead >= maxUnboundedBundleSize
      || Instant.now().isAfter(endTime)
      || context.isSinkFullHintSet()) {
    return false;
  }

  BackOff backoff = BACKOFF_FACTORY.backoff();
  while (true) {
    try {
      if (reader.advance()) {
        elemsRead++;
        return true;
      }
    } catch (Exception e) {
      throw new IOException("Failed to advance source: " + reader.getCurrentSource(), e);
    }
    long nextBackoff = backoff.nextBackOffMillis();
    if (nextBackoff == BackOff.STOP) {
      return false;
    }
    Uninterruptibles.sleepUninterruptibly(nextBackoff, TimeUnit.MILLISECONDS);
  }
}
 

private boolean advanceWithBackoff() throws IOException {
  // Try reading from the source with exponential backoff
  final BackOff backoff = backoffFactory.backoff();
  long nextSleep = backoff.nextBackOffMillis();
  while (nextSleep != BackOff.STOP) {
    if (readEndTime != null && Instant.now().isAfter(readEndTime)) {
      finalizeCheckpoint();
      return false;
    }
    if (unboundedReader.advance()) {
      recordsRead++;
      return true;
    }
    Uninterruptibles.sleepUninterruptibly(nextSleep, TimeUnit.MILLISECONDS);
    nextSleep = backoff.nextBackOffMillis();
  }
  finalizeCheckpoint();
  return false;
}
 

private void execute(Retriable retriable, FluentBackoff backoff) throws Exception {
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backOff = backoff.backoff();
  while (true) {
    try {
      retriable.execute();
      break;
    } catch (Exception e) {
      if (retriable.isExceptionRetriable(e) && BackOffUtils.next(sleeper, backOff)) {
        retriable.cleanUpAfterFailure();
      } else {
        throw e;
      }
    }
  }
}
 

private RunQueryResponse runQueryWithRetries(RunQueryRequest request) throws Exception {
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff = RUNQUERY_BACKOFF.backoff();
  while (true) {
    try {
      RunQueryResponse response = datastore.runQuery(request);
      rpcSuccesses.inc();
      return response;
    } catch (DatastoreException exception) {
      rpcErrors.inc();

      if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) {
        throw exception;
      }
      if (!BackOffUtils.next(sleeper, backoff)) {
        LOG.error("Aborting after {} retries.", MAX_RETRIES);
        throw exception;
      }
    }
  }
}
 

private boolean advanceWithBackoff(UnboundedReader<T> reader, Instant endTime)
    throws IOException {
  // Try reading from the source with exponential backoff
  BackOff backoff = BACKOFF_FACTORY.backoff();
  long nextSleep = backoff.nextBackOffMillis();
  while (true) {
    if (nextSleep == BackOff.STOP || (endTime != null && Instant.now().isAfter(endTime))) {
      return false;
    }
    if (reader.advance()) {
      return true;
    }
    Uninterruptibles.sleepUninterruptibly(nextSleep, TimeUnit.MILLISECONDS);
    nextSleep = backoff.nextBackOffMillis();
  }
}
 

@Test
public void testWriteJson() throws Exception {
  File destinationFile = new File(tempFolder.getRoot(), "json-outputs");
  BeamSqlEnv env = BeamSqlEnv.inMemory(new TextTableProvider());
  env.executeDdl(
      String.format(
          "CREATE EXTERNAL TABLE test %s TYPE text LOCATION '%s' TBLPROPERTIES '{\"format\":\"json\"}'",
          SQL_JSON_SCHEMA, destinationFile.getAbsolutePath()));

  BeamSqlRelUtils.toPCollection(
      pipeline, env.parseQuery("INSERT INTO test(name, age) VALUES ('Jack', 13)"));
  pipeline.run();

  assertThat(
      new NumberedShardedFile(destinationFile.getAbsolutePath() + "*")
          .readFilesWithRetries(Sleeper.DEFAULT, BackOff.STOP_BACKOFF),
      containsInAnyOrder(JSON_TEXT));
}
 

@Test
public void testWriteCsv() throws Exception {
  File destinationFile = new File(tempFolder.getRoot(), "csv-outputs");
  BeamSqlEnv env = BeamSqlEnv.inMemory(new TextTableProvider());

  // NumberedShardedFile
  env.executeDdl(
      String.format(
          "CREATE EXTERNAL TABLE test %s TYPE text LOCATION '%s' TBLPROPERTIES '{\"format\":\"csv\"}'",
          SQL_CSV_SCHEMA, destinationFile.getAbsolutePath()));

  BeamSqlRelUtils.toPCollection(
      pipeline, env.parseQuery("INSERT INTO test VALUES ('hello', 42), ('goodbye', 13)"));
  pipeline.run();

  assertThat(
      new NumberedShardedFile(destinationFile.getAbsolutePath() + "*")
          .readFilesWithRetries(Sleeper.DEFAULT, BackOff.STOP_BACKOFF),
      containsInAnyOrder("hello,42", "goodbye,13"));
}
 

@Test
public void testWriteLines() throws Exception {
  File destinationFile = new File(tempFolder.getRoot(), "lines-outputs");
  BeamSqlEnv env = BeamSqlEnv.inMemory(new TextTableProvider());
  env.executeDdl(
      String.format(
          "CREATE EXTERNAL TABLE test %s TYPE text LOCATION '%s' TBLPROPERTIES '{\"format\":\"lines\"}'",
          SQL_LINES_SCHEMA, destinationFile.getAbsolutePath()));

  BeamSqlRelUtils.toPCollection(
      pipeline, env.parseQuery("INSERT INTO test VALUES ('hello'), ('goodbye')"));
  pipeline.run();

  assertThat(
      new NumberedShardedFile(destinationFile.getAbsolutePath() + "*")
          .readFilesWithRetries(Sleeper.DEFAULT, BackOff.STOP_BACKOFF),
      containsInAnyOrder("hello", "goodbye"));
}
 

@Test
public void testInvalidJson() throws Exception {
  File deadLetterFile = new File(tempFolder.getRoot(), "dead-letter-file");
  Files.write(
      tempFolder.newFile("test.json").toPath(), INVALID_JSON_TEXT.getBytes(Charsets.UTF_8));

  BeamSqlEnv env = BeamSqlEnv.inMemory(new TextTableProvider());
  env.executeDdl(
      String.format(
          "CREATE EXTERNAL TABLE test %s TYPE text LOCATION '%s/*' "
              + "TBLPROPERTIES '{\"format\":\"json\", \"deadLetterFile\": \"%s\"}'",
          SQL_JSON_SCHEMA, tempFolder.getRoot(), deadLetterFile.getAbsoluteFile()));

  PCollection<Row> rows =
      BeamSqlRelUtils.toPCollection(pipeline, env.parseQuery("SELECT * FROM test"));

  PAssert.that(rows).empty();

  pipeline.run();
  assertThat(
      new NumberedShardedFile(deadLetterFile.getAbsoluteFile() + "*")
          .readFilesWithRetries(Sleeper.DEFAULT, BackOff.STOP_BACKOFF),
      containsInAnyOrder(INVALID_JSON_TEXT));
}
 

private void flushBatch() throws DatastoreException, IOException, InterruptedException {
  LOG.info("Writing batch of {} entities", entities.size());
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff =
      FluentBackoff.DEFAULT
          .withMaxRetries(MAX_RETRIES)
          .withInitialBackoff(INITIAL_BACKOFF)
          .backoff();

  while (true) {
    // Batch mutate entities.
    try {
      CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
      for (Entity entity : entities) {
        commitRequest.addMutations(mutationBuilder.apply(entity));
      }
      commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
      datastore.commit(commitRequest.build());
      // Break if the commit threw no exception.
      break;
    } catch (DatastoreException exception) {
      LOG.error(
          "Error writing to the Datastore ({}): {}",
          exception.getCode(),
          exception.getMessage());
      if (!BackOffUtils.next(sleeper, backoff)) {
        LOG.error("Aborting after {} retries.", MAX_RETRIES);
        throw exception;
      }
    }
  }
  LOG.info("Successfully wrote {} entities", entities.size());
  entities.clear();
}
 

private void flush() throws Exception {
  BackOff backOff = retryBackoff.backoff();
  int attempt = 0;

  if (buffer.isEmpty()) {
    return;
  }

  batchSize.update(buffer.size());

  while (true) {
    try (ClickHouseStatement statement = connection.createStatement()) {
      statement.sendRowBinaryStream(
          insertSql(schema(), table()),
          stream -> {
            for (Row row : buffer) {
              ClickHouseWriter.writeRow(stream, schema(), row);
            }
          });
      buffer.clear();
      break;
    } catch (SQLException e) {
      if (!BackOffUtils.next(Sleeper.DEFAULT, backOff)) {
        throw e;
      } else {
        retries.inc();
        LOG.warn(String.format(RETRY_ATTEMPT_LOG, attempt), e);
        attempt++;
      }
    }
  }
}
 

/**
 * Returns an adapter to convert from {@link BackOff} to {@link
 * com.google.api.client.util.BackOff}.
 */
public static com.google.api.client.util.BackOff toGcpBackOff(final BackOff backOff) {
  return new com.google.api.client.util.BackOff() {
    @Override
    public void reset() throws IOException {
      backOff.reset();
    }

    @Override
    public long nextBackOffMillis() throws IOException {
      return backOff.nextBackOffMillis();
    }
  };
}
 

public List<Shard> listShards(final String streamName) throws TransientKinesisException {
  return wrapExceptions(
      () -> {
        List<Shard> shards = Lists.newArrayList();
        String lastShardId = null;

        // DescribeStream has limits that can be hit fairly easily if we are attempting
        // to configure multiple KinesisIO inputs in the same account. Retry up to
        // LIST_SHARDS_DESCRIBE_STREAM_MAX_ATTEMPTS times if we end up hitting that limit.
        //
        // Only pass the wrapped exception up once that limit is reached. Use FluentBackoff
        // to implement the retry policy.
        FluentBackoff retryBackoff =
            FluentBackoff.DEFAULT
                .withMaxRetries(LIST_SHARDS_DESCRIBE_STREAM_MAX_ATTEMPTS)
                .withInitialBackoff(LIST_SHARDS_DESCRIBE_STREAM_INITIAL_BACKOFF);
        StreamDescription description = null;
        do {
          BackOff backoff = retryBackoff.backoff();
          Sleeper sleeper = Sleeper.DEFAULT;
          while (true) {
            try {
              description =
                  kinesis.describeStream(streamName, lastShardId).getStreamDescription();
              break;
            } catch (LimitExceededException exc) {
              if (!BackOffUtils.next(sleeper, backoff)) {
                throw exc;
              }
            }
          }

          shards.addAll(description.getShards());
          lastShardId = shards.get(shards.size() - 1).getShardId();
        } while (description.getHasMoreShards());

        return shards;
      });
}
 

@ProcessElement
public void processElement(ProcessContext context) throws Exception {
  PublishRequest request = context.element();
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff = retryBackoff.backoff();
  int attempt = 0;
  while (true) {
    attempt++;
    try {
      PublishResult pr = producer.publish(request);
      context.output(pr);
      break;
    } catch (Exception ex) {
      // Fail right away if there is no retry configuration
      if (spec.getRetryConfiguration() == null
          || !spec.getRetryConfiguration().getRetryPredicate().test(ex)) {
        SNS_WRITE_FAILURES.inc();
        LOG.info("Unable to publish message {} due to {} ", request.getMessage(), ex);
        throw new IOException("Error writing to SNS (no attempt made to retry)", ex);
      }

      if (!BackOffUtils.next(sleeper, backoff)) {
        throw new IOException(
            String.format(
                "Error writing to SNS after %d attempt(s). No more attempts allowed",
                attempt),
            ex);
      } else {
        // Note: this used in test cases to verify behavior
        LOG.warn(String.format(RETRY_ATTEMPT_LOG, attempt), ex);
      }
    }
  }
}
 

private void getConfig(String computation) {
  BackOff backoff =
      FluentBackoff.DEFAULT
          .withInitialBackoff(Duration.millis(100))
          .withMaxBackoff(Duration.standardMinutes(1))
          .withMaxCumulativeBackoff(Duration.standardMinutes(5))
          .backoff();
  while (running.get()) {
    try {
      if (windmillServiceEnabled) {
        getConfigFromDataflowService(computation);
      } else {
        getConfigFromWindmill(computation);
      }
      return;
    } catch (IllegalArgumentException | IOException e) {
      LOG.warn("Error fetching config: ", e);
      try {
        if (!BackOffUtils.next(Sleeper.DEFAULT, backoff)) {
          return;
        }
      } catch (IOException ioe) {
        LOG.warn("Error backing off, will not retry: ", ioe);
        return;
      } catch (InterruptedException ie) {
        Thread.currentThread().interrupt();
        return;
      }
    }
  }
}
 

private void executeBatch() throws SQLException, IOException, InterruptedException {
  if (records.isEmpty()) {
    return;
  }
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff = retryBackOff.backoff();
  while (true) {
    try (PreparedStatement preparedStatement =
        connection.prepareStatement(spec.getStatement().get())) {
      try {
        // add each record in the statement batch
        for (T record : records) {
          processRecord(record, preparedStatement);
        }
        // execute the batch
        preparedStatement.executeBatch();
        // commit the changes
        connection.commit();
        break;
      } catch (SQLException exception) {
        if (!spec.getRetryStrategy().apply(exception)) {
          throw exception;
        }
        LOG.warn("Deadlock detected, retrying", exception);
        // clean up the statement batch and the connection state
        preparedStatement.clearBatch();
        connection.rollback();
        if (!BackOffUtils.next(sleeper, backoff)) {
          // we tried the max number of times
          throw exception;
        }
      }
    }
  }
  records.clear();
}
 

@ProcessElement
public void processElement(ProcessContext context) throws Exception {
  PublishRequest request =
      (PublishRequest) spec.getPublishRequestFn().apply(context.element());
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff = retryBackoff.backoff();
  int attempt = 0;
  while (true) {
    attempt++;
    try {
      PublishResponse pr = producer.publish(request);
      context.output(pr);
      break;
    } catch (Exception ex) {
      // Fail right away if there is no retry configuration
      if (spec.getRetryConfiguration() == null
          || !spec.getRetryConfiguration().getRetryPredicate().test(ex)) {
        SNS_WRITE_FAILURES.inc();
        LOG.info("Unable to publish message {} due to {} ", request.message(), ex);
        throw new IOException("Error writing to SNS (no attempt made to retry)", ex);
      }

      if (!BackOffUtils.next(sleeper, backoff)) {
        throw new IOException(
            String.format(
                "Error writing to SNS after %d attempt(s). No more attempts allowed",
                attempt),
            ex);
      } else {
        // Note: this used in test cases to verify behavior
        LOG.warn(String.format(RETRY_ATTEMPT_LOG, attempt), ex);
      }
    }
  }
}
 

/** Helper for initializing the BackOff used for retries. */
private static BackOff createBackOff() {
  return FluentBackoff.DEFAULT
      .withInitialBackoff(Duration.millis(BACKOFF_INITIAL_INTERVAL_MILLIS))
      .withMaxBackoff(Duration.millis(BACKOFF_MAX_INTERVAL_MILLIS))
      .backoff();
}
 

private <ResponseT> ResponseT callWithBackoff(Supplier<ResponseT> function) {
  BackOff backoff = grpcBackoff();
  int rpcErrors = 0;
  while (true) {
    try {
      return function.get();
    } catch (StatusRuntimeException e) {
      try {
        if (++rpcErrors % 20 == 0) {
          LOG.warn(
              "Many exceptions calling gRPC. Last exception: {} with status {}",
              e,
              e.getStatus());
        }
        if (!BackOffUtils.next(Sleeper.DEFAULT, backoff)) {
          throw new WindmillServerStub.RpcException(e);
        }
      } catch (IOException | InterruptedException i) {
        if (i instanceof InterruptedException) {
          Thread.currentThread().interrupt();
        }
        WindmillServerStub.RpcException rpcException = new WindmillServerStub.RpcException(e);
        rpcException.addSuppressed(i);
        throw rpcException;
      }
    }
  }
}
 

private BackOff grpcBackoff() {
  return FluentBackoff.DEFAULT
      .withInitialBackoff(MIN_BACKOFF)
      .withMaxBackoff(maxBackoff)
      .backoff();
}
 

@BeforeClass
public static void setup() throws IOException, InterruptedException {
  // network sharing doesn't work with ClassRule
  network = Network.newNetwork();

  zookeeper =
      new GenericContainer<>("zookeeper:3.4.13")
          .withStartupAttempts(10)
          .withExposedPorts(2181)
          .withNetwork(network)
          .withNetworkAliases("zookeeper");

  // so far zookeeper container always starts successfully, so no extra retries
  zookeeper.start();

  clickHouse =
      (ClickHouseContainer)
          new ClickHouseContainer(CLICKHOUSE_IMAGE)
              .withStartupAttempts(10)
              .withCreateContainerCmdModifier(
                  // type inference for `(CreateContainerCmd) -> cmd.` doesn't work
                  cmd ->
                      ((CreateContainerCmd) cmd)
                          .withMemory(256 * 1024 * 1024L)
                          .withMemorySwap(4L * 1024 * 1024 * 1024L))
              .withNetwork(network)
              .withClasspathResourceMapping(
                  "config.d/zookeeper_default.xml",
                  "/etc/clickhouse-server/config.d/zookeeper_default.xml",
                  BindMode.READ_ONLY);

  BackOff backOff =
      FluentBackoff.DEFAULT
          .withMaxRetries(3)
          .withInitialBackoff(Duration.standardSeconds(15))
          .backoff();

  // try to start clickhouse-server a couple of times, see BEAM-6639
  while (true) {
    try {
      Unreliables.retryUntilSuccess(
          10,
          () -> {
            DockerClientFactory.instance()
                .checkAndPullImage(DockerClientFactory.instance().client(), CLICKHOUSE_IMAGE);

            return null;
          });

      clickHouse.start();
      break;
    } catch (Exception e) {
      if (!BackOffUtils.next(Sleeper.DEFAULT, backOff)) {
        throw e;
      } else {
        List<Image> images =
            DockerClientFactory.instance().client().listImagesCmd().withShowAll(true).exec();
        String listImagesOutput = "listImagesCmd:\n" + Joiner.on('\n').join(images) + "\n";

        LOG.warn("failed to start clickhouse-server\n\n" + listImagesOutput, e);
      }
    }
  }
}
 

private void flushBatch() throws IOException, InterruptedException {
  if (batch.isEmpty()) {
    return;
  }
  try {
    UpdateRequest updateRequest = new UpdateRequest();
    updateRequest.add(batch);

    Sleeper sleeper = Sleeper.DEFAULT;
    BackOff backoff = retryBackoff.backoff();
    int attempt = 0;
    while (true) {
      attempt++;
      try {
        solrClient.process(spec.getCollection(), updateRequest);
        break;
      } catch (Exception exception) {

        // fail immediately if no retry configuration doesn't handle this
        if (spec.getRetryConfiguration() == null
            || !spec.getRetryConfiguration().getRetryPredicate().test(exception)) {
          throw new IOException(
              "Error writing to Solr (no attempt made to retry)", exception);
        }

        // see if we can pause and try again
        if (!BackOffUtils.next(sleeper, backoff)) {
          throw new IOException(
              String.format(
                  "Error writing to Solr after %d attempt(s). No more attempts allowed",
                  attempt),
              exception);

        } else {
          // Note: this used in test cases to verify behavior
          LOG.warn(String.format(RETRY_ATTEMPT_LOG, attempt), exception);
        }
      }
    }
  } finally {
    batch.clear();
  }
}
 

private void flushBatch() throws IOException, InterruptedException {
  if (batch.isEmpty()) {
    return;
  }

  try {
    // Since each element is a KV<tableName, writeRequest> in the batch, we need to group them
    // by tableName
    Map<String, List<WriteRequest>> mapTableRequest =
        batch.stream()
            .collect(
                Collectors.groupingBy(
                    KV::getKey, Collectors.mapping(KV::getValue, Collectors.toList())));

    BatchWriteItemRequest batchRequest =
        BatchWriteItemRequest.builder().requestItems(mapTableRequest).build();

    Sleeper sleeper = Sleeper.DEFAULT;
    BackOff backoff = retryBackoff.backoff();
    int attempt = 0;
    while (true) {
      attempt++;
      try {
        client.batchWriteItem(batchRequest);
        break;
      } catch (Exception ex) {
        // Fail right away if there is no retry configuration
        if (spec.getRetryConfiguration() == null
            || !spec.getRetryConfiguration().getRetryPredicate().test(ex)) {
          DYNAMO_DB_WRITE_FAILURES.inc();
          LOG.info(
              "Unable to write batch items {} due to {} ",
              batchRequest.requestItems().entrySet(),
              ex);
          throw new IOException("Error writing to DynamoDB (no attempt made to retry)", ex);
        }

        if (!BackOffUtils.next(sleeper, backoff)) {
          throw new IOException(
              String.format(
                  "Error writing to DynamoDB after %d attempt(s). No more attempts allowed",
                  attempt),
              ex);
        } else {
          // Note: this used in test cases to verify behavior
          LOG.warn(String.format(RETRY_ATTEMPT_LOG, attempt), ex);
        }
      }
    }
  } finally {
    batch.clear();
  }
}
 

private void flushBatch() throws IOException, InterruptedException {
  if (batch.isEmpty()) {
    return;
  }

  try {
    // Since each element is a KV<tableName, writeRequest> in the batch, we need to group them
    // by tableName
    Map<String, List<WriteRequest>> mapTableRequest =
        batch.stream()
            .collect(
                Collectors.groupingBy(
                    KV::getKey, Collectors.mapping(KV::getValue, Collectors.toList())));

    BatchWriteItemRequest batchRequest = new BatchWriteItemRequest();
    mapTableRequest
        .entrySet()
        .forEach(
            entry -> batchRequest.addRequestItemsEntry(entry.getKey(), entry.getValue()));

    Sleeper sleeper = Sleeper.DEFAULT;
    BackOff backoff = retryBackoff.backoff();
    int attempt = 0;
    while (true) {
      attempt++;
      try {
        client.batchWriteItem(batchRequest);
        break;
      } catch (Exception ex) {
        // Fail right away if there is no retry configuration
        if (spec.getRetryConfiguration() == null
            || !spec.getRetryConfiguration().getRetryPredicate().test(ex)) {
          DYNAMO_DB_WRITE_FAILURES.inc();
          LOG.info(
              "Unable to write batch items {} due to {} ",
              batchRequest.getRequestItems().entrySet(),
              ex);
          throw new IOException("Error writing to DynamoDB (no attempt made to retry)", ex);
        }

        if (!BackOffUtils.next(sleeper, backoff)) {
          throw new IOException(
              String.format(
                  "Error writing to DynamoDB after %d attempt(s). No more attempts allowed",
                  attempt),
              ex);
        } else {
          // Note: this used in test cases to verify behavior
          LOG.warn(String.format(RETRY_ATTEMPT_LOG, attempt), ex);
        }
      }
    }
  } finally {
    batch.clear();
  }
}
 

/** Write the Mutations to Spanner, handling DEADLINE_EXCEEDED with backoff/retries. */
private void writeMutations(Iterable<Mutation> mutations) throws SpannerException, IOException {
  BackOff backoff = bundleWriteBackoff.backoff();
  long mutationsSize = Iterables.size(mutations);

  while (true) {
    Stopwatch timer = Stopwatch.createStarted();
    // loop is broken on success, timeout backoff/retry attempts exceeded, or other failure.
    try {
      spannerWriteWithRetryIfSchemaChange(mutations);
      spannerWriteSuccess.inc();
      return;
    } catch (SpannerException exception) {
      if (exception.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED) {
        spannerWriteTimeouts.inc();

        // Potentially backoff/retry after DEADLINE_EXCEEDED.
        long sleepTimeMsecs = backoff.nextBackOffMillis();
        if (sleepTimeMsecs == BackOff.STOP) {
          LOG.error(
              "DEADLINE_EXCEEDED writing batch of {} mutations to Cloud Spanner. "
                  + "Aborting after too many retries.",
              mutationsSize);
          spannerWriteFail.inc();
          throw exception;
        }
        LOG.info(
            "DEADLINE_EXCEEDED writing batch of {} mutations to Cloud Spanner, "
                + "retrying after backoff of {}ms\n"
                + "({})",
            mutationsSize,
            sleepTimeMsecs,
            exception.getMessage());
        spannerWriteRetries.inc();
        try {
          sleeper.sleep(sleepTimeMsecs);
        } catch (InterruptedException e) {
          // ignore.
        }
      } else {
        // Some other failure: pass up the stack.
        spannerWriteFail.inc();
        throw exception;
      }
    } finally {
      spannerWriteLatency.update(timer.elapsed(TimeUnit.MILLISECONDS));
    }
  }
}
 

/**
 * Writes a batch of mutations to Cloud Datastore.
 *
 * <p>If a commit fails, it will be retried up to {@link #MAX_RETRIES} times. All mutations in
 * the batch will be committed again, even if the commit was partially successful. If the retry
 * limit is exceeded, the last exception from Cloud Datastore will be thrown.
 *
 * @throws DatastoreException if the commit fails or IOException or InterruptedException if
 *     backing off between retries fails.
 */
private void flushBatch() throws DatastoreException, IOException, InterruptedException {
  LOG.debug("Writing batch of {} mutations", mutations.size());
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff = BUNDLE_WRITE_BACKOFF.backoff();

  while (true) {
    // Batch upsert entities.
    CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
    commitRequest.addAllMutations(mutations);
    commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
    long startTime = System.currentTimeMillis(), endTime;

    if (throttler.throttleRequest(startTime)) {
      LOG.info("Delaying request due to previous failures");
      throttledSeconds.inc(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS / 1000);
      sleeper.sleep(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS);
      continue;
    }

    try {
      datastore.commit(commitRequest.build());
      endTime = System.currentTimeMillis();

      writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
      throttler.successfulRequest(startTime);
      rpcSuccesses.inc();

      // Break if the commit threw no exception.
      break;
    } catch (DatastoreException exception) {
      if (exception.getCode() == Code.DEADLINE_EXCEEDED) {
        /* Most errors are not related to request size, and should not change our expectation of
         * the latency of successful requests. DEADLINE_EXCEEDED can be taken into
         * consideration, though. */
        endTime = System.currentTimeMillis();
        writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
      }
      // Only log the code and message for potentially-transient errors. The entire exception
      // will be propagated upon the last retry.
      LOG.error(
          "Error writing batch of {} mutations to Datastore ({}): {}",
          mutations.size(),
          exception.getCode(),
          exception.getMessage());
      rpcErrors.inc();

      if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) {
        throw exception;
      }
      if (!BackOffUtils.next(sleeper, backoff)) {
        LOG.error("Aborting after {} retries.", MAX_RETRIES);
        throw exception;
      }
    }
  }
  LOG.debug("Successfully wrote {} mutations", mutations.size());
  mutations.clear();
  mutationsSize = 0;
}
 

/** Write the Mutations to Spanner, handling DEADLINE_EXCEEDED with backoff/retries. */
private void writeMutations(Iterable<Mutation> mutations) throws SpannerException, IOException {
  BackOff backoff = bundleWriteBackoff.backoff();
  long mutationsSize = Iterables.size(mutations);

  while (true) {
    Stopwatch timer = Stopwatch.createStarted();
    // loop is broken on success, timeout backoff/retry attempts exceeded, or other failure.
    try {
      spannerWriteWithRetryIfSchemaChange(mutations);
      spannerWriteSuccess.inc();
      return;
    } catch (SpannerException exception) {
      if (exception.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED) {
        spannerWriteTimeouts.inc();

        // Potentially backoff/retry after DEADLINE_EXCEEDED.
        long sleepTimeMsecs = backoff.nextBackOffMillis();
        if (sleepTimeMsecs == BackOff.STOP) {
          LOG.error(
              "DEADLINE_EXCEEDED writing batch of {} mutations to Cloud Spanner. "
                  + "Aborting after too many retries.",
              mutationsSize);
          spannerWriteFail.inc();
          throw exception;
        }
        LOG.info(
            "DEADLINE_EXCEEDED writing batch of {} mutations to Cloud Spanner, "
                + "retrying after backoff of {}ms\n"
                + "({})",
            mutationsSize,
            sleepTimeMsecs,
            exception.getMessage());
        spannerWriteRetries.inc();
        try {
          sleeper.sleep(sleepTimeMsecs);
        } catch (InterruptedException e) {
          // ignore.
        }
      } else {
        // Some other failure: pass up the stack.
        spannerWriteFail.inc();
        throw exception;
      }
    } finally {
      spannerWriteLatency.update(timer.elapsed(TimeUnit.MILLISECONDS));
    }
  }
}