Java Code Examples for org.apache.mesos.Protos#TaskID

@Test
public void testRepairTaskIdsMismatchedIdsEmptyId() {
    Protos.TaskInfo taskInfo = newTaskInfo(TestConstants.TASK_NAME);
    stateStore.storeTasks(ImmutableList.of(taskInfo));
    Protos.TaskID taskID = Protos.TaskID.newBuilder().setValue("").build();
    Protos.TaskStatus taskStatus = newTaskStatus(taskID, Protos.TaskState.TASK_UNKNOWN);
    stateStore.storeStatus(TestConstants.TASK_NAME, taskStatus);

    StateStoreUtils.repairTaskIDs(stateStore);

    // Protos.TaskInfo was updated to match Status' id:
    assertThat(stateStore.fetchTask(TestConstants.TASK_NAME).get().getTaskId(), is(taskID));
    // Protos.TaskStatus state was not updated due to empty id:
    Protos.TaskStatus status = stateStore.fetchStatus(TestConstants.TASK_NAME).get();
    assertThat(status.getState(), is(Protos.TaskState.TASK_UNKNOWN));
    assertThat(status.getTaskId(), is(taskID));
}
 

@Override
public Protos.TaskInfo createTask(ResourceOfferContainer resourceOfferContainer, Protos.FrameworkID frameworkId, Protos.TaskID taskId, NodeTask nodeTask) {
  ServiceResourceProfile serviceProfile = nodeTask.getProfile();
  Double taskMemory = serviceProfile.getAggregateMemory();
  Double taskCpus = serviceProfile.getAggregateCpu();
  List<Protos.Resource> portResources = resourceOfferContainer.consumePorts(serviceProfile.getPorts().values());
  Protos.CommandInfo commandInfo = clGenerator.generateCommandLine(serviceProfile, null, rangesConverter(portResources));
  Protos.ExecutorInfo executorInfo = getExecutorInfoForSlave(resourceOfferContainer, frameworkId, commandInfo);
  Protos.TaskInfo.Builder taskBuilder = Protos.TaskInfo.newBuilder().setName(cfg.getFrameworkName() + "-" + taskId.getValue()).setTaskId(taskId).setSlaveId(
      resourceOfferContainer.getSlaveId());

  return taskBuilder
      .addAllResources(resourceOfferContainer.consumeCpus(taskCpus))
      .addAllResources(resourceOfferContainer.consumeMem(taskMemory))
      .addAllResources(portResources)
      .setExecutor(executorInfo)
      .build();
}
 

@Override
public void send(ClusterState state, SchedulerDriver mockDriver, Scheduler scheduler) {
  Protos.TaskID taskId = builder.taskId == null ?
      state.getTaskId(builder.taskName) :
      Protos.TaskID.newBuilder().setValue(builder.taskId).build();
  Protos.TaskStatus.Builder taskStatusBuilder = Protos.TaskStatus.newBuilder()
      .setTaskId(taskId)
      .setState(builder.taskState)
      .setMessage("This is a test status");

  if (builder.readinessCheckExitCode.isPresent()) {
    taskStatusBuilder
        .getCheckStatusBuilder()
        .getCommandBuilder()
        .setExitCode(builder.readinessCheckExitCode.get());
  }
  scheduler.statusUpdate(mockDriver, taskStatusBuilder.build());
}
 

@Override
public boolean removeWorker(Protos.TaskID taskID) throws Exception {
	checkNotNull(taskID, "taskID");
	String path = getPathForWorker(taskID);
	synchronized (startStopLock) {
		verifyIsRunning();

		if (workersInZooKeeper.exists(path) == -1) {
			LOG.debug("No such worker {} in ZooKeeper.", taskID);
			return false;
		}

		workersInZooKeeper.releaseAndTryRemove(path);
		LOG.debug("Removed worker {} from ZooKeeper.", taskID);
		return true;
	}
}
 

/**
 * Marks the TaskInfo for replacement if the agent has been decommissioned.
 */
public static Collection<Protos.TaskInfo> getTasksForReplacement(
    Collection<Protos.TaskStatus> alltaskStatuses,
    Collection<Protos.TaskInfo> allTaskInfos)
{
  Map<Protos.TaskID, Protos.TaskInfo> infoMap = new HashMap<>();
  for (Protos.TaskInfo taskInfo : allTaskInfos) {
    infoMap.put(taskInfo.getTaskId(), taskInfo);
  }

  List<Protos.TaskInfo> tasksNeedingReplace = new ArrayList<>();
  for (Protos.TaskStatus taskStatus: alltaskStatuses) {
    if (taskStatus.getState().equals(Protos.TaskState.TASK_GONE_BY_OPERATOR) &&
        !FailureUtils.isPermanentlyFailed(infoMap.get(taskStatus.getTaskId())))
    {
      LOGGER.info("{} needs replacement with state: {}",
          infoMap.get(taskStatus.getTaskId()).getName(),
          taskStatus.getState());
      tasksNeedingReplace.add(infoMap.get(taskStatus.getTaskId()));
    }
  }
  return tasksNeedingReplace;
}
 

/**
 * Verifies that the specified task (by exact id) was killed during the test.
 */
public static Expect taskIdKilled(String taskId) {
  return new Expect() {
    @Override
    public void expect(ClusterState state, SchedulerDriver mockDriver) {
      ArgumentCaptor<Protos.TaskID> taskIdCaptor = ArgumentCaptor.forClass(Protos.TaskID.class);
      Mockito.verify(mockDriver, Mockito.atLeastOnce()).killTask(taskIdCaptor.capture());
      Assert.assertEquals(taskId, taskIdCaptor.getValue().getValue());
    }

    @Override
    public String getDescription() {
      return String.format("Task with ID %s was killed", taskId);
    }
  };
}
 

@NotNull
static Protos.TaskStatus taskStatus(
    @NotNull final Protos.ExecutorID executorId,
    @NotNull final Protos.TaskID taskId,
    @NotNull final Protos.TaskState state,
    @NotNull final CassandraFrameworkProtos.SlaveStatusDetails details
) {
    return Protos.TaskStatus.newBuilder()
        .setExecutorId(executorId)
        .setTaskId(taskId)
        .setState(state)
        .setSource(Protos.TaskStatus.Source.SOURCE_EXECUTOR)
        .setData(ByteString.copyFrom(details.toByteArray()))
        .build();
}
 

public synchronized void makeTaskStaging(Protos.TaskID taskId) {
  Objects.requireNonNull(taskId, "taskId cannot be empty or null");
  String taskPrefix = taskIdPattern.split(taskId.getValue())[0];
  SchedulerStateForType taskTypeState = statesForTaskType.get(taskPrefix);
  if (taskTypeState == null) {
    taskTypeState = new SchedulerStateForType(taskPrefix);
    statesForTaskType.put(taskPrefix, taskTypeState);
  }
  taskTypeState.makeTaskStaging(taskId);
  updateStateStore();
}
 

public Set<Protos.TaskID> getStagingTaskIds() {
  Set<Protos.TaskID> returnSet = new HashSet<>();
  for (Map.Entry<String, SchedulerStateForType> entry : statesForTaskType.entrySet()) {
    returnSet.addAll(entry.getValue().getStagingTaskIds());
  }
  return returnSet;
}
 

/**
 * Construct a launchable Mesos worker.
 * @param resolver The resolver for retrieving artifacts (e.g. jars, configuration)
 * @param params the TM parameters such as memory, cpu to acquire.
 * @param containerSpec an abstract container specification for launch time.
 * @param taskID the taskID for this worker.
 */
public LaunchableMesosWorker(
		MesosArtifactResolver resolver,
		MesosTaskManagerParameters params,
		ContainerSpecification containerSpec,
		Protos.TaskID taskID,
		MesosConfiguration mesosConfiguration) {
	this.resolver = Preconditions.checkNotNull(resolver);
	this.containerSpec = Preconditions.checkNotNull(containerSpec);
	this.params = Preconditions.checkNotNull(params);
	this.taskID = Preconditions.checkNotNull(taskID);
	this.mesosConfiguration = Preconditions.checkNotNull(mesosConfiguration);

	this.taskRequest = new Request();
}
 

public Set<Protos.TaskID> getLostTaskIds() {
  Set<Protos.TaskID> returnSet = new HashSet<>();
  for (Map.Entry<String, SchedulerStateForType> entry : statesForTaskType.entrySet()) {
    returnSet.addAll(entry.getValue().getLostTaskIds());
  }
  return returnSet;
}
 

public synchronized NodeTask getTask(Protos.TaskID taskId) {
  return this.tasks.get(taskId);
}
 

@Override
public boolean clearDelay(Protos.TaskID taskID) {
  return false;
}
 

public Set<Protos.TaskID> getActiveTasks() {
  return toTaskSet(activeTasks);
}
 

private MockTaskExecutor mockTaskExecutor(Protos.TaskID taskID) {
	MockTaskExecutor task = new MockTaskExecutor(taskID);
	rpcService.registerGateway(task.address, task.gateway);
	return task;
}
 

public Set<Protos.TaskID> getStagingTaskIds(String taskPrefix) {
  SchedulerStateForType stateTask = statesForTaskType.get(taskPrefix);
  return (stateTask == null ? new HashSet<Protos.TaskID>() : stateTask.getStagingTaskIds());
}
 

/**
 * Get the ZK path for the given task ID (with leading slash).
 */
private static String getPathForWorker(Protos.TaskID taskID) {
	checkNotNull(taskID, "taskID");
	return String.format("/%s", taskID.getValue());
}
 

public void setStagingTasks(Set<Protos.TaskID> tasks) {
  if (tasks != null) {
    stagingTasks = new ArrayList<ByteBuffer>(tasks.size());
    toTaskBuffer(tasks, stagingTasks);
  }
}
 

public NodeCleanupJob(@NotNull final Protos.TaskID taskId, @NotNull final ExecutorService executorService) {
    super(taskId);

    this.executorService = executorService;
}
 

public synchronized Set<Protos.TaskID> getLostTaskIds() {
  return Collections.unmodifiableSet(this.lostTasks);
}