com.google.appengine.api.taskqueue.dev.QueueStateInfo Java Examples

/**
 * Executes mapreduce tasks, increment the clock between each task.
 *
 * <p>Incrementing the clock between tasks is important if tasks have transactions inside the
 * mapper or reducer, which don't have access to the fake clock.
 *
 * <p>The maxTasks parameter determines how many tasks (at most) will be run. If maxTasks is
 * absent(), all tasks are run until the queue is empty. If maxTasks is zero, no tasks are run.
 */
protected void executeTasks(
    String queueName, @Nullable FakeClock clock, Optional<Integer> maxTasks) throws Exception {
  for (int numTasksDeleted = 0;
      !maxTasks.isPresent() || (numTasksDeleted < maxTasks.get());
      numTasksDeleted++) {
    ofy().clearSessionCache();
    // We have to re-acquire task list every time, because local implementation returns a copy.
    List<QueueStateInfo.TaskStateInfo> taskInfo =
        taskQueue.getQueueStateInfo().get(queueName).getTaskInfo();
    if (taskInfo.isEmpty()) {
      break;
    }
    QueueStateInfo.TaskStateInfo taskStateInfo = taskInfo.get(0);
    taskQueue.deleteTask(queueName, taskStateInfo.getTaskName());
    executeTask(queueName, taskStateInfo);
    if (clock != null) {
      clock.advanceOneMilli();
    }
  }
}
 

protected void waitUntilTaskQueueIsEmpty() throws InterruptedException {
  boolean hasMoreTasks = true;
  while (hasMoreTasks) {
    Map<String, QueueStateInfo> taskInfoMap = taskQueue.getQueueStateInfo();
    hasMoreTasks = false;
    for (QueueStateInfo taskQueueInfo : taskInfoMap.values()) {
      if (taskQueueInfo.getCountTasks() > 0) {
        hasMoreTasks = true;
        break;
      }
    }
    if (hasMoreTasks) {
      Thread.sleep(100);
    }
  }
}
 

/**
 * Ensures that the only tasks in the named queue are exactly those that match the expected
 * matchers.
 */
public static void assertTasksEnqueued(String queueName, Collection<TaskMatcher> taskMatchers) {
  QueueStateInfo qsi = getQueueInfo(queueName);
  assertThat(qsi.getTaskInfo()).hasSize(taskMatchers.size());
  List<TaskStateInfo> taskInfos = new ArrayList<>(qsi.getTaskInfo());
  for (final TaskMatcher taskMatcher : taskMatchers) {
    try {
      taskInfos.remove(taskInfos.stream().filter(taskMatcher).findFirst().get());
    } catch (NoSuchElementException e) {
      final Map<String, Object> taskMatcherMap = taskMatcher.expected.toMap();
      assertWithMessage(
              "Task not found in queue %s:\n\n%s\n\nPotential candidate match diffs:\n\n%s",
              queueName,
              taskMatcher,
              taskInfos.stream()
                  .map(
                      input ->
                          prettyPrintEntityDeepDiff(
                              taskMatcherMap,
                              Maps.filterKeys(
                                  new MatchableTaskInfo(input).toMap(),
                                  in(taskMatcherMap.keySet()))))
                  .collect(joining("\n")))
          .fail();
    }
  }
}
 

private void doTest() throws InterruptedException {
  QueueFactory.getDefaultQueue().add(TaskOptions.Builder.withTaskName("task29"));
  // Give the task time to execute if tasks are actually enabled (which they
  // aren't, but that's part of the test).
  Thread.sleep(1000);
  LocalTaskQueue ltq = LocalTaskQueueTestConfig.getLocalTaskQueue();
  QueueStateInfo qsi = ltq.getQueueStateInfo().get(QueueFactory.getDefaultQueue().getQueueName());
  assertEquals(1, qsi.getTaskInfo().size());
  assertEquals("task29", qsi.getTaskInfo().get(0).getTaskName());
}
 

private void doTest() throws InterruptedException {
  // [START QueueFactory]
  QueueFactory.getQueue("my-queue-name").add(TaskOptions.Builder.withTaskName("task29"));
  // [END QueueFactory]
  // Give the task time to execute if tasks are actually enabled (which they
  // aren't, but that's part of the test).
  Thread.sleep(1000);
  LocalTaskQueue ltq = LocalTaskQueueTestConfig.getLocalTaskQueue();
  QueueStateInfo qsi =
      ltq.getQueueStateInfo().get(QueueFactory.getQueue("my-queue-name").getQueueName());
  assertEquals(1, qsi.getTaskInfo().size());
  assertEquals("task29", qsi.getTaskInfo().get(0).getTaskName());
}
 

/**
 * Asserts that the {@code numExpectedTasksInQueue} matches the actual number of tasks in the queue.
 */
private void assertNumTasksInQueue(int numExpectedTasksInQueue)
{
	LocalTaskQueue ltq = LocalTaskQueueTestConfig.getLocalTaskQueue();
	QueueStateInfo qsi = ltq.getQueueStateInfo()
		.get(QueueFactory.getQueue(DELETE_COUNTER_SHARD_QUEUE_NAME).getQueueName());
	assertEquals(numExpectedTasksInQueue, qsi.getTaskInfo().size());
}
 

/** Returns the info object for the provided queue name. */
public static QueueStateInfo getQueueInfo(String queueName) {
  return getLocalTaskQueue().getQueueStateInfo().get(queueName);
}
 

protected List<QueueStateInfo.TaskStateInfo> getTasks(String queueName) {
  return taskQueue.getQueueStateInfo().get(queueName).getTaskInfo();
}
 

protected void executeTask(String queueName, QueueStateInfo.TaskStateInfo taskStateInfo)
    throws Exception {
  logger.atFine().log(
      "Executing task %s with URL %s", taskStateInfo.getTaskName(), taskStateInfo.getUrl());
  // Hack to allow for deferred tasks. Exploits knowing how they work.
  if (taskStateInfo.getUrl().endsWith("__deferred__")) {
    ObjectInputStream oin =
        new ObjectInputStream(new ByteArrayInputStream(taskStateInfo.getBodyAsBytes()));
    Runnable object = (Runnable) oin.readObject();
    object.run();
    return;
  }
  HttpServletRequest request = mock(HttpServletRequest.class);
  HttpServletResponse response = mock(HttpServletResponse.class);

  // Strip off routing paths that are handled in web.xml in non-test scenarios.
  String pathInfo = taskStateInfo.getUrl();
  if (pathInfo.startsWith("/_dr/mapreduce/")) {
    pathInfo = pathInfo.replace("/_dr/mapreduce", "");
  } else if (pathInfo.startsWith("/mapreduce/")) {
      pathInfo = pathInfo.replace("/mapreduce", "");
  } else if (pathInfo.startsWith("/")) {
    pathInfo = pathInfo.replace("/_ah/", "");
    pathInfo = pathInfo.substring(pathInfo.indexOf('/'));
  } else {
    pathInfo = "/" + pathInfo;
  }
  when(request.getPathInfo()).thenReturn(pathInfo);
  when(request.getHeader("X-AppEngine-QueueName")).thenReturn(queueName);
  when(request.getHeader("X-AppEngine-TaskName")).thenReturn(taskStateInfo.getTaskName());
  // Pipeline looks at this header but uses the value only for diagnostic messages
  when(request.getIntHeader(TaskHandler.TASK_RETRY_COUNT_HEADER)).thenReturn(-1);
  for (HeaderWrapper header : taskStateInfo.getHeaders()) {
    int value = parseAsQuotedInt(header.getValue());
    when(request.getIntHeader(header.getKey())).thenReturn(value);
    logger.atFine().log("header: %s=%s", header.getKey(), header.getValue());
    when(request.getHeader(header.getKey())).thenReturn(header.getValue());
  }

  Map<String, String> parameters = decodeParameters(taskStateInfo.getBody());
  for (String name : parameters.keySet()) {
    when(request.getParameter(name)).thenReturn(parameters.get(name));
  }
  when(request.getParameterNames()).thenReturn(Collections.enumeration(parameters.keySet()));

  if (taskStateInfo.getMethod().equals("POST")) {
    if (taskStateInfo.getUrl().startsWith(PipelineServlet.BASE_URL)) {
      pipelineServlet.doPost(request, response);
    } else {
      mrServlet.doPost(request, response);
    }
  } else {
    throw new UnsupportedOperationException();
  }
}
 

private static int getCountTasks() {
    LocalTaskQueue localTaskQueue = LocalTaskQueueTestConfig.getLocalTaskQueue();

    Map<String, QueueStateInfo> queueStateInfo = localTaskQueue.getQueueStateInfo();

    int count = 0;

    for (String key : queueStateInfo.keySet()) {
        count += queueStateInfo.get(key).getCountTasks();
    }

    return count;
}