/** Leases and returns all tasks from the queue with the specified tag tld, in batches. */
List<TaskHandle> loadAllTasks(Queue queue, String tld) {
  ImmutableList.Builder<TaskHandle> allTasks = new ImmutableList.Builder<>();
  while (true) {
    List<TaskHandle> tasks =
        retrier.callWithRetry(
            () ->
                queue.leaseTasks(
                    LeaseOptions.Builder.withTag(tld)
                        .leasePeriod(LEASE_PERIOD.getMillis(), TimeUnit.MILLISECONDS)
                        .countLimit(TaskQueueUtils.getBatchSize())),
            TransientFailureException.class,
            DeadlineExceededException.class);
    if (tasks.isEmpty()) {
      return allTasks.build();
    }
    allTasks.addAll(tasks);
  }
}
 

/** Deletes a list of tasks from the given queue using a retrier. */
private static void deleteTasksWithRetry(
    final List<DnsRefreshRequest> refreshRequests,
    final Queue queue,
    AsyncTaskMetrics asyncTaskMetrics,
    Retrier retrier,
    OperationResult result) {
  if (refreshRequests.isEmpty()) {
    return;
  }
  final List<TaskHandle> tasks =
      refreshRequests.stream().map(DnsRefreshRequest::task).collect(toImmutableList());
  retrier.callWithRetry(() -> queue.deleteTask(tasks), TransientFailureException.class);
  refreshRequests.forEach(
      r -> asyncTaskMetrics.recordAsyncFlowResult(DNS_REFRESH, result, r.requestedTime()));
}
 

@Test
public void testTransact_transientFailureException_retries() {
  assertThat(
          tm().transact(
                  new Supplier<Integer>() {

                    int count = 0;

                    @Override
                    public Integer get() {
                      count++;
                      if (count == 3) {
                        return count;
                      }
                      throw new TransientFailureException("");
                    }
                  }))
      .isEqualTo(3);
}
 

void doReadOnlyRetryTest(final RuntimeException e) {
  assertThat(
          tm().transactNewReadOnly(
                  new Supplier<Integer>() {

                    int count = 0;

                    @Override
                    public Integer get() {
                      count++;
                      if (count == 3) {
                        return count;
                      }
                      throw new TransientFailureException("");
                    }
                  }))
      .isEqualTo(3);
}
 

/**
 * Adds tasks to an App Engine task queue in a reliable manner.
 *
 * <p>This is the same as {@link Queue#add(Iterable)} except it'll automatically retry with
 * exponential backoff if {@link TransientFailureException} is thrown.
 *
 * @throws TransientFailureException if retrying failed for the maximum period of time, or an
 *     {@link InterruptedException} told us to stop trying
 * @return successfully enqueued tasks
 */
public List<TaskHandle> enqueue(final Queue queue, final Iterable<TaskOptions> tasks) {
  return retrier.callWithRetry(
      () -> {
        for (TaskOptions task : tasks) {
          logger.atInfo().log(
              "Enqueuing queue='%s' endpoint='%s'", queue.getQueueName(), task.getUrl());
        }
        return queue.add(tasks);
      },
      TransientFailureException.class);
}
 

/** Deletes the specified tasks from the queue in batches, with retrying. */
public void deleteTasks(Queue queue, List<TaskHandle> tasks) {
  Lists.partition(tasks, BATCH_SIZE)
      .stream()
      .forEach(
          batch ->
              retrier.callWithRetry(
                  () -> queue.deleteTask(batch), TransientFailureException.class));
}
 

/**
 * Transact with commit logs and retry with exponential backoff.
 *
 * <p>This method is broken out from {@link #transactNew(Work)} for testing purposes.
 */
@VisibleForTesting
<R> R transactCommitLoggedWork(CommitLoggedWork<R> work) {
  long baseRetryMillis = getBaseOfyRetryDuration().getMillis();
  for (long attempt = 0, sleepMillis = baseRetryMillis;
      true;
      attempt++, sleepMillis *= 2) {
    try {
      ofy().transactNew(() -> {
        work.run();
        return null;
      });
      return work.getResult();
    } catch (TransientFailureException
        | TimestampInversionException
        | DatastoreTimeoutException
        | DatastoreFailureException e) {
      // TransientFailureExceptions come from task queues and always mean nothing committed.
      // TimestampInversionExceptions are thrown by our code and are always retryable as well.
      // However, Datastore exceptions might get thrown even if the transaction succeeded.
      if ((e instanceof DatastoreTimeoutException || e instanceof DatastoreFailureException)
          && checkIfAlreadySucceeded(work)) {
        return work.getResult();
      }
      if (attempt == NUM_RETRIES) {
        throw e;  // Give up.
      }
      sleeper.sleepUninterruptibly(Duration.millis(sleepMillis));
      logger.atInfo().withCause(e).log(
          "Retrying %s, attempt %d", e.getClass().getSimpleName(), attempt);
    }
  }
}
 

@Override
public final void map(@Nullable final DomainBase domain) {
  if (domain == null) {
    // Emit a single value so that the reducer always runs.  The key and value don't matter.
    emit(true, true);
    return;
  }
  Key<HostResource> referencingHostKey = null;
  for (DnsRefreshRequest request : refreshRequests) {
    if (isActive(domain, request.lastUpdateTime())
        && domain.getNameservers().contains(VKey.from(request.hostKey()))) {
      referencingHostKey = request.hostKey();
      break;
    }
  }
  if (referencingHostKey != null) {
    retrier.callWithRetry(
        () -> dnsQueue.addDomainRefreshTask(domain.getDomainName()),
        TransientFailureException.class);
    logger.atInfo().log(
        "Enqueued DNS refresh for domain %s referenced by host %s.",
        domain.getDomainName(), referencingHostKey);
    getContext().incrementCounter("domains refreshed");
  } else {
    getContext().incrementCounter("domains not refreshed");
  }

  // Don't catch errors -- we allow the mapreduce to terminate on any errors that can't be
  // resolved by retrying the transaction. The reducer only fires if the mapper completes
  // without errors, meaning that it is acceptable to delete all tasks.
}
 

/**
 * Deletes a list of tasks associated with deletion requests from the async delete queue using a
 * retrier.
 */
private void deleteStaleTasksWithRetry(final List<DeletionRequest> deletionRequests) {
  if (deletionRequests.isEmpty()) {
    return;
  }
  final List<TaskHandle> tasks =
      deletionRequests.stream().map(DeletionRequest::task).collect(toImmutableList());
  retrier.callWithRetry(() -> queue.deleteTask(tasks), TransientFailureException.class);
  deletionRequests.forEach(
      deletionRequest ->
          asyncTaskMetrics.recordAsyncFlowResult(
              deletionRequest.getMetricOperationType(),
              OperationResult.STALE,
              deletionRequest.requestedTime()));
}
 

/** Returns handles for a batch of tasks, leased for the specified duration. */
public List<TaskHandle> leaseTasks(Duration leaseDuration) {
  try {
    rateLimiter.acquire();
    int numTasks = queue.fetchStatistics().getNumTasks();
    logger.at((numTasks >= leaseTasksBatchSize) ? Level.WARNING : Level.INFO).log(
        "There are %d tasks in the DNS queue '%s'.", numTasks, DNS_PULL_QUEUE_NAME);
    return queue.leaseTasks(leaseDuration.getMillis(), MILLISECONDS, leaseTasksBatchSize);
  } catch (TransientFailureException | DeadlineExceededException e) {
    logger.atSevere().withCause(e).log("Failed leasing tasks too fast");
    return ImmutableList.of();
  }
}
 

/** Delete a list of tasks, removing them from the queue permanently. */
public void deleteTasks(List<TaskHandle> tasks) {
  try {
    queue.deleteTask(tasks);
  } catch (TransientFailureException | DeadlineExceededException e) {
    logger.atSevere().withCause(e).log("Failed deleting tasks too fast");
  }
}
 

@SuppressWarnings("unchecked")
@Test
public void test_loadAllTasks_retryLogic_thirdTrysTheCharm() {
  Queue queue = mock(Queue.class);
  TaskHandle task = new TaskHandle(TaskOptions.Builder.withTaskName("blah"), "blah");
  when(queue.leaseTasks(any(LeaseOptions.class)))
      .thenThrow(TransientFailureException.class)
      .thenThrow(DeadlineExceededException.class)
      .thenReturn(ImmutableList.of(task), ImmutableList.of());
  assertThat(action.loadAllTasks(queue, "tld")).containsExactly(task);
}
 

@Test
public void test_loadAllTasks_retryLogic_allFailures() {
  Queue queue = mock(Queue.class);
  when(queue.leaseTasks(any(LeaseOptions.class)))
      .thenThrow(new TransientFailureException("some transient error"));
  RuntimeException thrown =
      assertThrows(TransientFailureException.class, () -> action.loadAllTasks(queue, "tld"));
  assertThat(thrown).hasMessageThat().isEqualTo("some transient error");
}
 

@Test
public void testEnqueue_twoTransientErrorsThenSuccess_stillWorksAfterSleeping() {
  when(queue.add(ImmutableList.of(task)))
      .thenThrow(new TransientFailureException(""))
      .thenThrow(new TransientFailureException(""))
      .thenReturn(ImmutableList.of(handle));
  assertThat(taskQueueUtils.enqueue(queue, task)).isSameInstanceAs(handle);
  verify(queue, times(3)).add(ImmutableList.of(task));
  assertThat(clock.nowUtc()).isEqualTo(DateTime.parse("2000-01-01T00:00:00.6Z"));  // 200 + 400ms
}
 

@Test
public void testEnqueue_maxRetries_givesUp() {
  when(queue.add(ImmutableList.of(task)))
      .thenThrow(new TransientFailureException("one"))
      .thenThrow(new TransientFailureException("two"))
      .thenThrow(new TransientFailureException("three"))
      .thenThrow(new TransientFailureException("four"));
  TransientFailureException thrown =
      assertThrows(TransientFailureException.class, () -> taskQueueUtils.enqueue(queue, task));
  assertThat(thrown).hasMessageThat().contains("three");
}
 

@Test
public void testEnqueue_transientErrorThenInterrupt_throwsTransientError() {
  when(queue.add(ImmutableList.of(task))).thenThrow(new TransientFailureException(""));
  try {
    Thread.currentThread().interrupt();
    assertThrows(TransientFailureException.class, () -> taskQueueUtils.enqueue(queue, task));
  } finally {
    Thread.interrupted();  // Clear interrupt state so it doesn't pwn other tests.
  }
}
 

RuntimeException constructApiException(String packageName, String methodName) {
  String message = "RCP Failure for API call: " + packageName + " " + methodName;

  switch (packageName) {
    case "taskqueue":
      return new TransientFailureException(message);
    case "app_identity_service":
      return new AppIdentityServiceFailureException(message);
    case "blobstore":
      return new BlobstoreFailureException(message);
    case "channel":
      return new ChannelFailureException(message);
    case "images":
      return new ImagesServiceFailureException(message);
    case "logservice":
      return constructException(
          LogServiceException.class.getName(), message, packageName, methodName);
    case "memcache":
      return new MemcacheServiceException(message);
    case "modules":
      return constructException(
          ModulesException.class.getName(), message, packageName, methodName);
    case "search":
      return new SearchException(message);
    case "user":
      return new UserServiceFailureException(message);
    case "xmpp":
      return new XMPPFailureException(message);
    default:

      // Cover all datastore versions:
      if (packageName.startsWith("datastore")) {
        return new DatastoreFailureException(message);
      } else {
        return new RPCFailedException(packageName, methodName);
      }
  }
}
 

@Test
public void testTransactNewReadOnly_transientFailureException_retries() {
  doReadOnlyRetryTest(new TransientFailureException(""));
}
 

/**
 * @return List<TaskHandle> the accumulated list of tasks up to count.
 */
private List<TaskHandle> leaseTasksByOptions(String groupTag, int count,
                                             boolean zeroSizeAcceptable, LeaseOptions options,
                                             int retry, int interval) {
    List<TaskHandle> handleList = null;
    List<TaskHandle> masterHandleList = new ArrayList<>();

    int retryCount = retry;
    int retryCounter = 0;
    int retryInterval = interval;
    while (masterHandleList.size() < count) {
        sync(retryInterval);  // first iteration gives time for tasks to activate.
        try {
            handleList = queue.leaseTasks(options);
        } catch (TransientFailureException tfe) {  // This is common.
            sync(retryInterval);
            log.warning(tfe.toString());
            handleList = null;
            continue;
        }

        if (handleList.size() > 0) {
            masterHandleList.addAll(handleList);
        }

        if (handleList.size() >= 0 && zeroSizeAcceptable) {
            return masterHandleList;  // even zero tasks okay, return what we got.
        }

        if (masterHandleList.size() >= count) {
            return masterHandleList;  // Success, got all tasks requested.
        }

        if (retryCounter++ > retryCount) {
            break;
        }
    }

    String errMsg = "Couldn't lease " + Integer.toString(count) + " tag:" +
        groupTag + " after " + retryCount + " attempts.";
    log.warning(errMsg);
    if (handleList == null) {  // Couldn't communicate with Task service.
        throw new TransientFailureException(errMsg);
    }

    return masterHandleList;  // Return what we've got, could be partial.
}
 

/**
 * @return List<TaskHandle> the accumulated list of tasks up to count.
 */
private List<TaskHandle> leaseTasksByOptions(String groupTag, int count,
                                             boolean zeroSizeAcceptable, LeaseOptions options,
                                             int retry, int interval) {
    List<TaskHandle> handleList = null;
    List<TaskHandle> masterHandleList = new ArrayList<TaskHandle>();

    int retryCount = retry;
    int retryCounter = 0;
    int retryInterval = interval;
    while (masterHandleList.size() < count) {
        sleep(retryInterval);  // first iteration gives time for tasks to activate.
        try {
            handleList = queue.leaseTasks(options);
        } catch (TransientFailureException tfe) {  // This is common.
            sleep(retryInterval);
            log.warning(tfe.toString());
            handleList = null;
            continue;
        }

        if (handleList.size() > 0) {
            masterHandleList.addAll(handleList);
        }

        if (handleList.size() >= 0 && zeroSizeAcceptable) {
            return masterHandleList;  // even zero tasks okay, return what we got.
        }

        if (masterHandleList.size() >= count) {
            return masterHandleList;  // Success, got all tasks requested.
        }

        if (retryCounter++ > retryCount) {
            break;
        }
    }

    String errMsg = "Couldn't lease " + Integer.toString(count) + " tag:" + groupTag + " after " + retryCount + " attempts.";
    log.warning(errMsg);
    if (handleList == null) {  // Couldn't communicate with Task service.
        throw new TransientFailureException(errMsg);
    }

    return masterHandleList;  // Return what we've got, could be partial.
}
 

/**
 * Adds a task to a queue with retrying, to avoid aborting the entire flow over a transient issue
 * enqueuing a task.
 */
private void addTaskToQueueWithRetry(final Queue queue, final TaskOptions task) {
  retrier.callWithRetry(() -> queue.add(task), TransientFailureException.class);
}