com.google.api.client.util.Sleeper Java Examples

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} fails. */
@Test
public void testPollJobFailed() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus().setState("DONE").setErrorResult(new ErrorProto()));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testJob));

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(testJob, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Test
public void testEmptyBuilder() {
  RetryConfig config = RetryConfig.builder().build();

  assertTrue(config.getRetryStatusCodes().isEmpty());
  assertEquals(0, config.getMaxRetries());
  assertEquals(2 * 60 * 1000, config.getMaxIntervalMillis());
  assertEquals(2.0, config.getBackOffMultiplier(), 0.01);
  assertSame(Sleeper.DEFAULT, config.getSleeper());

  ExponentialBackOff backOff = (ExponentialBackOff) config.newBackOff();
  assertEquals(2 * 60 * 1000, backOff.getMaxIntervalMillis());
  assertEquals(2.0, backOff.getMultiplier(), 0.01);
  assertEquals(500, backOff.getInitialIntervalMillis());
  assertEquals(0.0, backOff.getRandomizationFactor(), 0.01);
  assertNotSame(backOff, config.newBackOff());
}
 

@Test
public void testGetTableNotFound() throws IOException, InterruptedException {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(404);

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");
  Table table = datasetService.getTable(tableRef, null, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);

  assertNull(table);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

/**
 * Attempts to get the underlying {@link Job}. Uses exponential backoff on failure up to the
 * maximum number of passed in attempts.
 *
 * @param backoff the {@link BackOff} used to control retries.
 * @param sleeper Object used to do the sleeps between attempts.
 * @return The underlying {@link Job} object.
 * @throws IOException When the maximum number of retries is exhausted, the last exception is
 *     thrown.
 */
private Job getJobWithRetries(BackOff backoff, Sleeper sleeper) throws IOException {
  // Retry loop ends in return or throw
  while (true) {
    try {
      Job job = dataflowClient.getJob(getJobId());
      State currentState = MonitoringUtil.toState(job.getCurrentState());
      if (currentState.isTerminal()) {
        terminalState = currentState;
        replacedByJob =
            new DataflowPipelineJob(
                dataflowClient, job.getReplacedByJobId(), dataflowOptions, transformStepNames);
      }
      return job;
    } catch (IOException exn) {
      LOG.warn("There were problems getting current job status: {}.", exn.getMessage());
      LOG.debug("Exception information:", exn);

      if (!nextBackOff(sleeper, backoff)) {
        throw exn;
      }
    }
  }
}
 

@Test
public void testExecuteWithRetries() throws IOException, InterruptedException {
  Table testTable = new Table();

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testTable));

  Table table =
      BigQueryServicesImpl.executeWithRetries(
          bigquery.tables().get("projectId", "datasetId", "tableId"),
          "Failed to get table.",
          Sleeper.DEFAULT,
          BackOff.STOP_BACKOFF,
          BigQueryServicesImpl.ALWAYS_RETRY);

  assertEquals(testTable, table);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Test
public void testIsTableEmptyNoRetryForNotFound() throws IOException, InterruptedException {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(404);

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");

  thrown.expect(IOException.class);
  thrown.expectMessage(String.format("Unable to list table data: %s", tableRef.getTableId()));

  try {
    datasetService.isTableEmpty(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
  } finally {
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
  }
}
 

@Test
public void testIsTableEmptySucceeds() throws Exception {
  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");

  TableDataList testDataList = new TableDataList().setRows(ImmutableList.of(new TableRow()));

  // First response is 403 rate limited, second response has valid payload.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403)))
      .thenReturn(toStream(testDataList));

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  assertFalse(datasetService.isTableEmpty(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT));

  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
}
 

/** Synchronously stages a package, with retry and backoff for resiliency. */
private StagingResult stagePackageSynchronously(
    PackageAttributes attributes, Sleeper retrySleeper, CreateOptions createOptions)
    throws IOException, InterruptedException {
  String sourceDescription = attributes.getSourceDescription();
  String target = attributes.getDestination().getLocation();

  if (alreadyStaged(attributes)) {
    LOG.debug("Skipping file already staged: {} at {}", sourceDescription, target);
    return StagingResult.cached(attributes);
  }

  try {
    return tryStagePackageWithRetry(attributes, retrySleeper, createOptions);
  } catch (Exception miscException) {
    throw new RuntimeException(
        String.format("Could not stage %s to %s", sourceDescription, target), miscException);
  }
}
 

/**
 * Returns the project number or throws an error if the project does not exist or has other
 * access errors.
 */
private static long getProjectNumber(
    String projectId, CloudResourceManager crmClient, BackOff backoff, Sleeper sleeper)
    throws IOException {
  CloudResourceManager.Projects.Get getProject = crmClient.projects().get(projectId);
  try {
    Project project =
        ResilientOperation.retry(
            ResilientOperation.getGoogleRequestCallable(getProject),
            backoff,
            RetryDeterminer.SOCKET_ERRORS,
            IOException.class,
            sleeper);
    return project.getProjectNumber();
  } catch (Exception e) {
    throw new IOException("Unable to get project number", e);
  }
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} succeeds. */
@Test
public void testPollJobSucceeds() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus().setState("DONE"));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testJob));

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(testJob, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds with a retry. */
@Test
public void testStartLoadJobRetry() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

  // First response is 403 rate limited, second response has valid payload.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403)))
      .thenReturn(toStream(testJob));

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob,
      new ApiErrorExtractor(),
      bigquery,
      sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
}
 

/**
 * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds with an already
 * exist job.
 */
@Test
public void testStartLoadJobSucceedsAlreadyExists() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

  when(response.getStatusCode()).thenReturn(409); // 409 means already exists

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob,
      new ApiErrorExtractor(),
      bigquery,
      sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
  expectedLogs.verifyNotLogged("Started BigQuery job");
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds. */
@Test
public void testStartLoadJobSucceeds() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testJob));

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob,
      new ApiErrorExtractor(),
      bigquery,
      sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
  expectedLogs.verifyInfo(String.format("Started BigQuery job: %s", jobRef));
}
 

@Override
public Job pollJob(JobReference jobRef, int maxAttempts) throws InterruptedException {
  BackOff backoff =
      BackOffAdapter.toGcpBackOff(
          FluentBackoff.DEFAULT
              .withMaxRetries(maxAttempts)
              .withInitialBackoff(Duration.millis(10))
              .withMaxBackoff(Duration.standardSeconds(1))
              .backoff());
  Sleeper sleeper = Sleeper.DEFAULT;
  try {
    do {
      Job job = getJob(jobRef);
      if (job != null) {
        JobStatus status = job.getStatus();
        if (status != null
            && ("DONE".equals(status.getState()) || "FAILED".equals(status.getState()))) {
          return job;
        }
      }
    } while (BackOffUtils.next(sleeper, backoff));
  } catch (IOException e) {
    return null;
  }
  return null;
}
 

@Override
public Table patchTableDescription(
    TableReference tableReference, @Nullable String tableDescription)
    throws IOException, InterruptedException {
  Table table = new Table();
  table.setDescription(tableDescription);

  return executeWithRetries(
      client
          .tables()
          .patch(
              tableReference.getProjectId(),
              tableReference.getDatasetId(),
              tableReference.getTableId(),
              table),
      String.format(
          "Unable to patch table description: %s, aborting after %d retries.",
          tableReference, MAX_RPC_RETRIES),
      Sleeper.DEFAULT,
      createDefaultBackoff(),
      ALWAYS_RETRY);
}
 

@Override
public <T> long insertAll(
    TableReference ref,
    List<ValueInSingleWindow<TableRow>> rowList,
    @Nullable List<String> insertIdList,
    InsertRetryPolicy retryPolicy,
    List<ValueInSingleWindow<T>> failedInserts,
    ErrorContainer<T> errorContainer,
    boolean skipInvalidRows,
    boolean ignoreUnknownValues,
    boolean ignoreInsertIds)
    throws IOException, InterruptedException {
  return insertAll(
      ref,
      rowList,
      insertIdList,
      BackOffAdapter.toGcpBackOff(INSERT_BACKOFF_FACTORY.backoff()),
      Sleeper.DEFAULT,
      retryPolicy,
      failedInserts,
      errorContainer,
      skipInvalidRows,
      ignoreUnknownValues,
      ignoreInsertIds);
}
 

/**
 * {@inheritDoc}
 *
 * <p>Tries executing the RPC for at most {@code MAX_RPC_RETRIES} times until it succeeds.
 *
 * @throws IOException if it exceeds {@code MAX_RPC_RETRIES} attempts.
 */
@Override
public void createDataset(
    String projectId,
    String datasetId,
    @Nullable String location,
    @Nullable String description,
    @Nullable Long defaultTableExpirationMs)
    throws IOException, InterruptedException {
  createDataset(
      projectId,
      datasetId,
      location,
      description,
      defaultTableExpirationMs,
      Sleeper.DEFAULT,
      createDefaultBackoff());
}
 

@VisibleForTesting
boolean isTableEmpty(TableReference tableRef, BackOff backoff, Sleeper sleeper)
    throws IOException, InterruptedException {
  TableDataList dataList =
      executeWithRetries(
          client
              .tabledata()
              .list(tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()),
          String.format(
              "Unable to list table data: %s, aborting after %d retries.",
              tableRef.getTableId(), MAX_RPC_RETRIES),
          sleeper,
          backoff,
          DONT_RETRY_NOT_FOUND);
  return dataList.getRows() == null || dataList.getRows().isEmpty();
}
 

protected <T extends GenericJson> T retryFailures(
    AbstractGoogleClientRequest<T> request, BackOff backOff) throws IOException {
  while (true) {
    try {
      return request.execute();
    } catch (GoogleJsonResponseException e) {
      try {
        long nextPause = backOff.nextBackOffMillis();
        if (nextPause == BackOff.STOP) {
          throw e;
        }
        System.out.printf("Operation failed, retrying in %f seconds.%n", nextPause / 1000.0);
        BackOffUtils.next(Sleeper.DEFAULT, backOff);
      } catch (InterruptedException ie) {
        // Just go straight into retry if interrupted.
      }
    }
  }
}
 

@VisibleForTesting
@Nullable
Table getTable(
    TableReference ref, @Nullable List<String> selectedFields, BackOff backoff, Sleeper sleeper)
    throws IOException, InterruptedException {
  Tables.Get get =
      client.tables().get(ref.getProjectId(), ref.getDatasetId(), ref.getTableId());
  if (selectedFields != null && !selectedFields.isEmpty()) {
    get.setSelectedFields(String.join(",", selectedFields));
  }
  try {
    return executeWithRetries(
        get,
        String.format(
            "Unable to get table: %s, aborting after %d retries.",
            ref.getTableId(), MAX_RPC_RETRIES),
        sleeper,
        backoff,
        DONT_RETRY_NOT_FOUND);
  } catch (IOException e) {
    if (errorExtractor.itemNotFound(e)) {
      return null;
    }
    throw e;
  }
}
 

@Test
public void testGetJobNotFound() throws Exception {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(404);

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.getJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(null, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Test
public void testGetJobThrows() throws Exception {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(401);

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  thrown.expect(IOException.class);
  thrown.expectMessage(String.format("Unable to find BigQuery job: %s", jobRef));

  jobService.getJob(jobRef, Sleeper.DEFAULT, BackOff.STOP_BACKOFF);
}
 

@Override
public State getState() {
  if (terminalState != null) {
    return terminalState;
  }

  return getStateWithRetriesOrUnknownOnException(
      BackOffAdapter.toGcpBackOff(STATUS_BACKOFF_FACTORY.backoff()), Sleeper.DEFAULT);
}
 

/**
 * {@inheritDoc}
 *
 * <p>Tries executing the RPC for at most {@code MAX_RPC_RETRIES} times until it succeeds.
 *
 * @throws IOException if it exceeds {@code MAX_RPC_RETRIES} attempts.
 */
@Override
public void deleteDataset(String projectId, String datasetId)
    throws IOException, InterruptedException {
  executeWithRetries(
      client.datasets().delete(projectId, datasetId),
      String.format(
          "Unable to delete table: %s, aborting after %d retries.", datasetId, MAX_RPC_RETRIES),
      Sleeper.DEFAULT,
      createDefaultBackoff(),
      ALWAYS_RETRY);
}
 

/**
 * Attempts to get the state. Uses exponential backoff on failure up to the maximum number of
 * passed in attempts.
 *
 * @param attempts The amount of attempts to make.
 * @param sleeper Object used to do the sleeps between attempts.
 * @return The state of the job or State.UNKNOWN in case of failure.
 */
@VisibleForTesting
State getStateWithRetriesOrUnknownOnException(BackOff attempts, Sleeper sleeper) {
  try {
    return getStateWithRetries(attempts, sleeper);
  } catch (IOException exn) {
    // The only IOException that getJobWithRetries is permitted to throw is the final IOException
    // that caused the failure of retry. Other exceptions are wrapped in an unchecked exceptions
    // and will propagate.
    return State.UNKNOWN;
  }
}
 

State getStateWithRetries(BackOff attempts, Sleeper sleeper) throws IOException {
  if (terminalState != null) {
    return terminalState;
  }
  Job job = getJobWithRetries(attempts, sleeper);
  return MonitoringUtil.toState(job.getCurrentState());
}
 

/**
 * {@inheritDoc}
 *
 * <p>Tries executing the RPC for at most {@code MAX_RPC_RETRIES} times until it succeeds.
 *
 * @throws IOException if it exceeds {@code MAX_RPC_RETRIES} attempts.
 */
@Override
public void deleteTable(TableReference tableRef) throws IOException, InterruptedException {
  executeWithRetries(
      client
          .tables()
          .delete(tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()),
      String.format(
          "Unable to delete table: %s, aborting after %d retries.",
          tableRef.getTableId(), MAX_RPC_RETRIES),
      Sleeper.DEFAULT,
      createDefaultBackoff(),
      ALWAYS_RETRY);
}
 

/**
 * {@inheritDoc}
 *
 * <p>If a table with the same name already exists in the dataset, the function simply returns.
 * In such a case, the existing table doesn't necessarily have the same schema as specified by
 * the parameter.
 *
 * @throws IOException if other error than already existing table occurs.
 */
@Override
public void createTable(Table table) throws InterruptedException, IOException {
  LOG.info(
      "Trying to create BigQuery table: {}",
      BigQueryHelpers.toTableSpec(table.getTableReference()));
  BackOff backoff =
      new ExponentialBackOff.Builder()
          .setMaxElapsedTimeMillis(RETRY_CREATE_TABLE_DURATION_MILLIS)
          .build();

  tryCreateTable(table, backoff, Sleeper.DEFAULT);
}
 

/** Identical to {@link BackOffUtils#next} but without checked exceptions. */
private boolean nextBackOff(Sleeper sleeper, BackOff backoff) {
  try {
    return BackOffUtils.next(sleeper, backoff);
  } catch (InterruptedException | IOException e) {
    if (e instanceof InterruptedException) {
      Thread.currentThread().interrupt();
    }
    throw new RuntimeException(e);
  }
}
 

/** Identical to {@link BackOffUtils#next} but without checked IOException. */
private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) throws InterruptedException {
  try {
    return BackOffUtils.next(sleeper, backoff);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}