Java Code Examples for org.apache.mesos.v1.Protos#TaskStatus

public CompletableFuture<StatusUpdateResult> processStatusUpdateAsync(
  Protos.TaskStatus status
) {
  return CompletableFuture.supplyAsync(
    () -> {
      final String taskId = status.getTaskId().getValue();
      final Optional<SingularityTaskId> maybeTaskId = getTaskId(taskId);

      if (!maybeTaskId.isPresent()) {
        return StatusUpdateResult.INVALID_TASK_ID;
      }

      return schedulerLock.runWithRequestLockAndReturn(
        () -> unsafeProcessStatusUpdate(status, maybeTaskId.get()),
        maybeTaskId.get().getRequestId(),
        getClass().getSimpleName()
      );
    },
    statusUpdatesExecutor.getExecutorService()
  );
}
 

@Benchmark
public boolean runBenchmark() throws InterruptedException {
  for (String taskId : org.apache.aurora.scheduler.base.Tasks.ids(tasks)) {
    Protos.TaskStatus status = Protos.TaskStatus.newBuilder()
        .setState(Protos.TaskState.TASK_RUNNING)
        .setSource(Protos.TaskStatus.Source.SOURCE_EXECUTOR)
        .setMessage("message")
        .setTimestamp(1D)
        .setTaskId(Protos.TaskID.newBuilder().setValue(taskId).build())
        .build();

    scheduler.statusUpdate(new FakeSchedulerDriver(), ProtosConversion.convert(status));
  }

  // Wait for all task transitions to complete.
  countDownLatch.await();

  // Return an unguessable value.
  return System.currentTimeMillis() % 5 == 0;
}
 

@Test
public void testUpdateNoSource() {
  Protos.TaskStatus status = STATUS.toBuilder().clearSource().build();

  eventSink.post(new PubsubEvent.TaskStatusReceived(
      status.getState(),
      Optional.empty(),
      Optional.ofNullable(status.getReason()),
      Optional.of(1000000L)
  ));
  statusHandler.statusUpdate(status);

  control.replay();

  handler.handleUpdate(status);
  assertEquals(1L, statsProvider.getLongValue("scheduler_status_update"));
}
 

@Test
public void testUpdateNoReason() {
  Protos.TaskStatus status = STATUS.toBuilder().clearReason().build();

  eventSink.post(new PubsubEvent.TaskStatusReceived(
      status.getState(),
      Optional.ofNullable(status.getSource()),
      Optional.empty(),
      Optional.of(1000000L)
  ));
  statusHandler.statusUpdate(status);

  control.replay();

  handler.handleUpdate(status);
  assertEquals(1L, statsProvider.getLongValue("scheduler_status_update"));
}
 

@Test
public void testUpdateNoMessage() {
  Protos.TaskStatus status = STATUS.toBuilder().clearMessage().build();

  eventSink.post(new PubsubEvent.TaskStatusReceived(
      status.getState(),
      Optional.ofNullable(status.getSource()),
      Optional.ofNullable(status.getReason()),
      Optional.of(1000000L)
  ));
  statusHandler.statusUpdate(status);

  control.replay();

  handler.handleUpdate(status);
  assertEquals(1L, statsProvider.getLongValue("scheduler_status_update"));
}
 

@Override
public void acknowledgeStatusUpdate(Protos.TaskStatus status) {
  ensureRunning();

  TaskStatus convertedStatus = ProtosConversion.convert(status);
  Futures.getUnchecked(driverFuture).acknowledgeStatusUpdate(convertedStatus);
}
 

@Override
public void reconcileTasks(Collection<Protos.TaskStatus> statuses) {
  ensureRunning();

  Collection<TaskStatus> convertedStatuses =
      Collections2.transform(statuses, ProtosConversion::convert);
  Futures.getUnchecked(driverFuture).reconcileTasks(convertedStatuses);
}
 

private Optional<String> getStatusMessage(
  Protos.TaskStatus status,
  Optional<SingularityTask> task
) {
  if (status.hasMessage() && !Strings.isNullOrEmpty(status.getMessage())) {
    return Optional.of(status.getMessage());
  } else if (
    status.hasReason() &&
    status.getReason() == Reason.REASON_CONTAINER_LIMITATION_MEMORY
  ) {
    if (
      task.isPresent() &&
      task.get().getTaskRequest().getDeploy().getResources().isPresent()
    ) {
      if (
        task.get().getTaskRequest().getDeploy().getResources().get().getDiskMb() > 0
      ) {
        return Optional.of(
          String.format(
            "Task exceeded one or more memory limits (%s MB mem, %s MB disk).",
            task.get().getTaskRequest().getDeploy().getResources().get().getMemoryMb(),
            task.get().getTaskRequest().getDeploy().getResources().get().getDiskMb()
          )
        );
      } else {
        return Optional.of(
          String.format(
            "Task exceeded memory limit (%s MB mem).",
            task.get().getTaskRequest().getDeploy().getResources().get().getMemoryMb()
          )
        );
      }
    }
    return Optional.of("Task exceeded memory limit.");
  } else if (
    status.hasReason() && status.getReason() == Reason.REASON_CONTAINER_LIMITATION_DISK
  ) {
    if (
      task.isPresent() &&
      task.get().getTaskRequest().getDeploy().getResources().isPresent()
    ) {
      return Optional.of(
        String.format(
          "Task exceeded disk limit (%s MB disk).",
          task.get().getTaskRequest().getDeploy().getResources().get().getDiskMb()
        )
      );
    } else {
      return Optional.of("Task exceeded disk limit.");
    }
  }

  return Optional.empty();
}
 

public void handleCompletedTask(
  Optional<SingularityTask> task,
  SingularityTaskId taskId,
  long timestamp,
  ExtendedTaskState state,
  SingularityCreateResult taskHistoryUpdateCreateResult,
  Protos.TaskStatus status
) {
  final SingularityDeployStatistics deployStatistics = getDeployStatistics(
    taskId.getRequestId(),
    taskId.getDeployId()
  );

  if (!task.isPresent() || task.get().getTaskRequest().getRequest().isLoadBalanced()) {
    taskManager.createLBCleanupTask(taskId);
  }

  if (requestManager.isBouncing(taskId.getRequestId())) {
    List<SingularityTaskId> activeTaskIds = taskManager.getActiveTaskIdsForRequest(
      taskId.getRequestId()
    );
    boolean foundBouncingTask = false;
    for (SingularityTaskId activeTaskId : activeTaskIds) {
      Optional<SingularityTaskHistoryUpdate> maybeCleaningUpdate = taskManager.getTaskHistoryUpdate(
        activeTaskId,
        ExtendedTaskState.TASK_CLEANING
      );
      if (maybeCleaningUpdate.isPresent()) {
        if (maybeCleaningUpdate.get().getStatusReason().orElse("").contains("BOUNCE")) { // TaskCleanupType enum is included in status message
          LOG.debug("Found task {} still waiting for bounce to complete", activeTaskId);
          foundBouncingTask = true;
          break;
        } else if (!maybeCleaningUpdate.get().getPrevious().isEmpty()) {
          for (SingularityTaskHistoryUpdate previousUpdate : maybeCleaningUpdate
            .get()
            .getPrevious()) {
            if (previousUpdate.getStatusMessage().orElse("").contains("BOUNCE")) {
              LOG.debug(
                "Found task {} still waiting for bounce to complete",
                activeTaskId
              );
              foundBouncingTask = true;
              break;
            }
          }
        }
      }
    }
    if (!foundBouncingTask) {
      LOG.info(
        "Bounce completed for request {}, no cleaning tasks due to bounce found",
        taskId.getRequestId()
      );
      Optional<SingularityExpiringBounce> expiringBounce = requestManager.getExpiringBounce(
        taskId.getRequestId()
      );

      if (
        expiringBounce.isPresent() &&
        expiringBounce.get().getDeployId().equals(taskId.getDeployId())
      ) {
        requestManager.removeExpiringBounce(taskId.getRequestId());
      }
      requestManager.markBounceComplete(taskId.getRequestId());
    }
  }

  final Optional<PendingType> scheduleResult = handleCompletedTaskWithStatistics(
    task,
    taskId,
    timestamp,
    state,
    deployStatistics,
    taskHistoryUpdateCreateResult,
    status
  );

  if (taskHistoryUpdateCreateResult == SingularityCreateResult.EXISTED) {
    return;
  }

  updateDeployStatistics(
    deployStatistics,
    taskId,
    task,
    timestamp,
    state,
    scheduleResult,
    status
  );
}
 

public static Protos.TaskStatus convert(org.apache.mesos.Protos.TaskStatus s) {
  return convert(s, Protos.TaskStatus.newBuilder());
}
 

public static org.apache.mesos.Protos.TaskStatus convert(Protos.TaskStatus f) {
  return convert(f, org.apache.mesos.Protos.TaskStatus.newBuilder());
}
 

@TimedInterceptor.Timed("scheduler_received")
@Override
public void received(Mesos mesos, Event event) {
  countEventMetrics(event);
  switch(event.getType()) {
    case SUBSCRIBED:
      Event.Subscribed subscribed = event.getSubscribed();
      if (isRegistered.get()) {
        handler.handleReregistration(subscribed.getMasterInfo());
      } else {
        isRegistered.set(true);
        handler.handleRegistration(subscribed.getFrameworkId(), subscribed.getMasterInfo());
      }
      isSubscribed.set(true);
      break;

    case OFFERS:
      checkState(isSubscribed.get(), "Must be registered before receiving offers.");
      handler.handleOffers(event.getOffers().getOffersList());
      break;

    case RESCIND:
      handler.handleRescind(event.getRescind().getOfferId());
      break;

    case INVERSE_OFFERS:
      handler.handleInverseOffer(event.getInverseOffers().getInverseOffersList());
      break;

    case RESCIND_INVERSE_OFFER:
      Protos.OfferID id = event.getRescindInverseOffer().getInverseOfferId();
      LOG.warn("Ignoring rescinded inverse offer: {}", id);
      break;

    case UPDATE:
      Protos.TaskStatus status = event.getUpdate().getStatus();
      handler.handleUpdate(status);
      break;

    case MESSAGE:
      Event.Message m = event.getMessage();
      handler.handleMessage(m.getExecutorId(), m.getAgentId());
      break;

    case ERROR:
      handler.handleError(event.getError().getMessage());
      break;

    case FAILURE:
      Event.Failure failure = event.getFailure();
      if (failure.hasExecutorId()) {
        handler.handleLostExecutor(
            failure.getExecutorId(),
            failure.getAgentId(),
            failure.getStatus());
      } else {
        handler.handleLostAgent(failure.getAgentId());
      }
      break;

    // TODO(zmanji): handle HEARTBEAT in a graceful manner
    // For now it is ok to silently ignore heart beats because the driver wil
    // detect disconnections for us.
    case HEARTBEAT:
      break;

    default:
      LOG.warn("Unknown event from Mesos \n{}", event);
      break;
  }
}
 

@Override
public void acknowledgeStatusUpdate(Protos.TaskStatus status) {
  // no-op
}
 

@Override
public void reconcileTasks(Collection<Protos.TaskStatus> statuses) {
  // no-op
}
 

@Test
public void testExecution() {
  expect(statsProvider.makeCounter(EXPLICIT_STAT_NAME)).andReturn(explicitRuns);
  expect(statsProvider.makeCounter(IMPLICIT_STAT_NAME)).andReturn(implicitRuns);
  FakeScheduledExecutor clock =
      FakeScheduledExecutor.scheduleAtFixedRateExecutor(executorService, 2, 5);

  IScheduledTask task1 = makeTask("id1", TaskTestUtil.makeConfig(TaskTestUtil.JOB));
  IScheduledTask task2 = makeTask("id2", TaskTestUtil.makeConfig(TaskTestUtil.JOB));
  storageUtil.expectOperations();
  storageUtil.expectTaskFetch(
      Query.unscoped().byStatus(Tasks.SLAVE_ASSIGNED_STATES),
      task1,
      task2).times(7);

  List<List<Protos.TaskStatus>> batches = Lists.partition(ImmutableList.of(
      TASK_TO_PROTO.apply(task1),
      TASK_TO_PROTO.apply(task2)), BATCH_SIZE);

  driver.reconcileTasks(batches.get(0));
  expectLastCall().times(7);

  driver.reconcileTasks(batches.get(1));
  expectLastCall().times(7);

  driver.reconcileTasks(EasyMock.anyObject());
  expectLastCall().times(3);

  control.replay();

  TaskReconciler reconciler = new TaskReconciler(
      SETTINGS,
      storageUtil.storage,
      driver,
      executorService,
      statsProvider);

  reconciler.startAsync().awaitRunning();

  clock.advance(INITIAL_DELAY);
  assertEquals(1L, explicitRuns.get());
  assertEquals(0L, implicitRuns.get());

  clock.advance(SPREAD);
  assertEquals(1L, explicitRuns.get());
  assertEquals(1L, implicitRuns.get());

  clock.advance(EXPLICIT_SCHEDULE);
  assertEquals(2L, explicitRuns.get());
  assertEquals(1L, implicitRuns.get());

  clock.advance(IMPLICT_SCHEDULE);
  assertEquals(5L, explicitRuns.get());
  assertEquals(2L, implicitRuns.get());

  reconciler.triggerExplicitReconciliation(Optional.of(BATCH_SIZE));
  assertEquals(6L, explicitRuns.get());
  reconciler.triggerImplicitReconciliation();
  assertEquals(3L, implicitRuns.get());

  reconciler.triggerExplicitReconciliation(Optional.empty());
  assertEquals(7L, explicitRuns.get());
  assertEquals(3L, implicitRuns.get());
}