org.apache.mesos.Protos.TaskID Java Examples

public Status kill(final TaskID taskId) {
  LOGGER.info("Killing task {}", taskId);
  this.driverLock.lock();
  try {
    if (isRunning()) {
      this.driverStatus = driver.kill(taskId);
      LOGGER.info("Task {} killed with status: {}", taskId, this.driverStatus);
    } else {
      LOGGER.warn("Cannot kill task, driver is not running");
    }
  } finally {
    this.driverLock.unlock();
  }

  return driverStatus;
}
 

public void killTask(TaskID taskId) {
    synchronized (lock) {
        TaskExecutorHolder taskExecutorHolder = pendingTasks.remove(taskId);
        if (taskExecutorHolder == null) {
            logger.warn(slaveId + " is not running task " + taskId); // don't throw, treat it as no-op
            return;
        }
        Protos.TaskState taskState = taskExecutorHolder.getState();
        if (taskState == Protos.TaskState.TASK_FINISHED || taskState == Protos.TaskState.TASK_FAILED || taskState == Protos.TaskState.TASK_KILLED) {
            taskUpdates.onNext(Protos.TaskStatus.newBuilder()
                    .setTaskId(taskId)
                    .setState(Protos.TaskState.TASK_LOST)
                    .setMessage("Task already terminated: " + taskState)
                    .build()
            );
            return;
        }

        taskExecutorHolder.transitionTo(Protos.TaskState.TASK_KILLED);

        releaseResourcesAndReOffer(taskExecutorHolder);
    }
}
 

public Status kill(final TaskID taskId) {
	LOGGER.info("Killing task {}", taskId);
	this.driverLock.lock();
	try {
		if (isRunning()) {
			this.driverStatus = driver.kill(taskId);
			LOGGER.info("Task {} killed with status: {}", taskId,
					this.driverStatus);
		} else {
			LOGGER.warn("Cannot kill task, driver is not running");
		}
	} finally {
		this.driverLock.unlock();
	}

	return driverStatus;
}
 

/**
 * Setup testing target & sample data.
 */
public LaunchTaskTest()
{
  SlaveID slaveID = SlaveID.newBuilder().setValue("s1").build();
  this.sampleTaskInfo =
      TaskInfo.newBuilder()
              .setName("t1").setSlaveId(slaveID)
              .setTaskId(TaskID.newBuilder().setValue("id2"))
              .setExecutor(
                  ExecutorInfo.newBuilder()
                              .setExecutorId(ExecutorID.newBuilder().setValue("e1"))
                              .setCommand(CommandInfo.getDefaultInstance()))
              .build();
  this.sampleOffer =
      Offer.newBuilder()
           .setHostname("h1").setSlaveId(slaveID)
           .setId(OfferID.newBuilder().setValue("id1"))
           .setFrameworkId(FrameworkID.newBuilder().setValue("f1").build())
           .build();
  this.target = new LaunchTask(sampleTaskInfo, sampleOffer);
}
 

private synchronized boolean unAssign(TaskID taskId) {
    EniAndIpAssignment assignment = eniTaskAssignments.remove(taskId);
    if (assignment == null) {
        logger.warn("[ENI unassign] ENI unassign request for unknown task: {}", taskId.getValue());
        return false;
    }
    String eniLabel = assignment.getEni();
    Pair<Integer, String> eniState = eniAssignments.get(eniLabel);
    int availableIpCount = eniState.getLeft() + 1;
    if (availableIpCount == ipsPerEni) {
        eniAssignments.remove(eniLabel);
    } else {
        eniAssignments.put(eniLabel, Pair.of(availableIpCount, eniState.getRight()));
    }
    logger.info("[ENI unassign] Unassigned ENI: taskId={}, eniIdx={}, availableIps={}", taskId.getValue(), eniLabel, availableIpCount);
    return true;
}
 

@Override
public void killedWorker(int port) {
  LOG.info("killedWorker: executor {} removing port {} assignment and sending " +
      "TASK_FINISHED update to Mesos", _executorId, port);
  TaskID taskId = _taskAssignments.deregister(port);
  if (taskId == null) {
    LOG.error("killedWorker: Executor {} failed to find TaskID for port {}, so not " +
        "issuing TaskStatus update to Mesos for this dead task.", _executorId, port);
    return;
  }
  TaskStatus status = TaskStatus.newBuilder()
      .setState(TaskState.TASK_FINISHED)
      .setTaskId(taskId)
      .build();
  _driver.sendStatusUpdate(status);
}
 

private void onRuntimeError() {
  // Shutdown REEF Evaluator
  if (this.evaluatorProcess != null) {
    this.evaluatorProcess.destroy();
  }
  mesosExecutorDriver.sendStatusUpdate(TaskStatus.newBuilder()
      .setTaskId(TaskID.newBuilder()
          .setValue(mesosExecutorId)
          .build())
      .setState(TaskState.TASK_FAILED)
      .setMessage("Evaluator Process exited with status " + String.valueOf(evaluatorProcessExitValue))
      .build());

  // Shutdown Mesos Executor
  this.executorService.shutdown();
  this.mesosExecutorDriver.stop();
}
 

@Override
@Before
public void setUp() throws Exception {
  super.setUp();
  sState = TestObjectFactory.getSchedulerState(this.cfg, "tmp/scheduler-utils-test");
  idOne = TaskID.newBuilder().setValue("Task1").build();
  idTwo = TaskID.newBuilder().setValue("Task2").build();
  idThree = TaskID.newBuilder().setValue("Task3").build();
  taskOne = TestObjectFactory.getNodeTask("zero", "server1", 0.0, 0.0, Long.valueOf(1), Long.valueOf(2));
  taskTwo = TestObjectFactory.getNodeTask("low", "localhost", 0.2, 1024.0, Long.valueOf(1), Long.valueOf(2));
  taskThree = TestObjectFactory.getNodeTask("medium", "localhost", 0.4, 2048.0, Long.valueOf(1), Long.valueOf(2));

  sState.addTask(idOne, taskOne);
  sState.addTask(idTwo, taskTwo);  
  sState.addTask(idThree, taskThree);
  
  this.baseStateStoreDirectory = "/tmp/scheduler-utils-test";
}
 

/**
 * De-serialize the internal ByteBuffers back into a Task map.
 *
 * @return
 */
public Map<Protos.TaskID, NodeTask> getTasks() {
  Map<Protos.TaskID, NodeTask> map = null;
  if (taskIds != null) {
    map = new HashMap<Protos.TaskID, NodeTask>(taskIds.size());
    int idx = 0;
    for (ByteBuffer bb : taskIds) {
      final Protos.TaskID taskId = ByteBufferSupport.toTaskId(bb);
      final NodeTask task = ByteBufferSupport.toNodeTask(taskNodes.get(idx++));
      if (task.getTaskPrefix() == null && taskId != null) {
        String taskPrefix = taskIdPattern.split(taskId.getValue())[0];
        task.setTaskPrefix(taskPrefix);
      }
      map.put(taskId, task);
    }
  } else {
    map = new HashMap<Protos.TaskID, NodeTask>(0);
  }
  return map;
}
 

@Override
public Status killTask(final TaskID taskId)
{
    checkNotNull(taskId, "taskId is null");

    if (!context.isStateMachine(DRIVER_RUNNING)) {
        return context.getStateMachine();
    }

    final KillTaskMessage message = KillTaskMessage.newBuilder()
        .setFrameworkId(context.getFrameworkId())
        .setTaskId(taskId)
        .build();
    eventBus.post(new RemoteMessageEnvelope(context.getDriverUPID(), context.getMasterUPID(), message));

    return context.getStateMachine();
}
 

private SchedulerState getSchedulerState() throws Exception {
  SchedulerState state = new SchedulerState(new MyriadFileSystemRMStateStore());
  idOne = Protos.TaskID.newBuilder().setValue("nt-1").build();
  idTwo = Protos.TaskID.newBuilder().setValue("nt-2").build();
  idThree = Protos.TaskID.newBuilder().setValue("nt-3").build();
  TreeMap<String, Long> ports = new TreeMap<>();

  state.addTask(idOne, new NodeTask(new ServiceResourceProfile("profile1", 0.2, 1024.0, ports), new LikeConstraint("localhost", "host-[0-9]*.example.com")));
  state.addTask(idTwo, new NodeTask(new ServiceResourceProfile("profile2", 0.4, 2048.0, ports), new LikeConstraint("localhost", "host-[0-9]*.example.com")));
  state.addTask(idThree, new NodeTask(new ServiceResourceProfile("profile3", 0.6, 3072.0, ports), new LikeConstraint("localhost", "host-[0-9]*.example.com")));

  state.setFrameworkId(FrameworkID.newBuilder().setValue("mock-framework").build());
  state.makeTaskActive(idOne);
  state.makeTaskPending(idTwo);
  state.makeTaskStaging(idThree);

  return state;
}
 

private void handleNonPendingTask(TaskID taskId) { 
  /*
   * Kill the task and decline additional offers for it, but hold off removing from SchedulerState. 
   * Removal of the killable task must be done following invocation of statusUpdate callback method
   * which constitutes acknowledgement from Mesos that the kill task request succeeded.
   */
  Status status = this.driverManager.kill(taskId);
  NodeTask task = schedulerState.getTask(taskId);

  if (task != null) {
    offerLifeCycleManager.declineOutstandingOffers(task.getHostname());
  } 
  if (status.equals(Status.DRIVER_RUNNING)) {
    LOGGER.info("Kill request for {} was submitted to a running SchedulerDriver", taskId);
  } else {
    LOGGER.warn("Kill task request for {} submitted to non-running SchedulerDriver, may fail", taskId);
  }
}
 

/**
 * Builds a TaskInfo from the given jobspec
 * @param job - JobSpec to TaskInfo-ify
 * @param offer - offer add extra data (SlaveId)
 * @return TaskInfo fully formed
 */
private TaskInfo getTaskInfo(JobSpec job,Offer offer) {
    TaskID taskId = TaskID.newBuilder().setValue(job.getJob().getId()).build();
    TaskInfo info = TaskInfo.newBuilder().setName("task " + taskId.getValue())
             .setTaskId(taskId)
             .setSlaveId(offer.getSlaveId())
             .addResources(Resource.newBuilder()
                           .setName("cpus")
                           .setType(Value.Type.SCALAR)
                           .setScalar(Value.Scalar.newBuilder().setValue(job.getJob().getLoadValue()*1.0)))
             .addResources(Resource.newBuilder()
                           .setName("mem")
                           .setType(Value.Type.SCALAR)
                           .setScalar(Value.Scalar.newBuilder().setValue(job.getJob().getLoadValue()*1024.0)))
             .setExecutor(ExecutorInfo.newBuilder(executor)).setData(MesosUtilities.jobSpecToByteString(job)).build();
    return info;
}
 

@Test 
public void testMakeTestLost() throws Exception {
  SchedulerState sState = initialize();
  TaskID idOne = TaskID.newBuilder().setValue("Task1").build();
  TaskID idTwo = TaskID.newBuilder().setValue("Task2").build();
  sState.makeTaskLost(idOne);
  sState.makeTaskLost(idTwo);
  assertEquals(2, sState.getLostTaskIds().size());
  assertTrue(sState.getLostTaskIds().contains(idOne));
  assertTrue(sState.getLostTaskIds().contains(idTwo));    
}
 

@Test
public void testRemoveTask() throws Exception {
  SchedulerState sState = initialize();
  TaskID idOne = TaskID.newBuilder().setValue("Task1").build();
  TaskID idTwo = TaskID.newBuilder().setValue("Task2").build();
  sState.removeTask(idOne);
  assertNull(sState.getTask(idOne));
  sState.removeTask(idTwo);
  assertNull(sState.getTask(idTwo));
}
 

private void onStop() {
  // Shutdown REEF Evaluator
  if (this.evaluatorProcess != null) {
    this.evaluatorProcess.destroy();
    mesosExecutorDriver.sendStatusUpdate(TaskStatus.newBuilder()
        .setTaskId(TaskID.newBuilder()
            .setValue(mesosExecutorId)
            .build())
        .setState(TaskState.TASK_FINISHED)
        .setMessage("Evaluator Process exited with status " + String.valueOf(evaluatorProcessExitValue))
        .build());
  } else {
    mesosExecutorDriver.sendStatusUpdate(TaskStatus.newBuilder()
        .setTaskId(TaskID.newBuilder()
            .setValue(mesosExecutorId)
            .build())
        .setState(TaskState.TASK_FINISHED)
        .setData(ByteString.copyFromUtf8("eval_not_run"))
        // TODO[JIRA REEF-102]: a hack to pass closeEvaluator test, replace this with a better interface
        .setMessage("Evaluator Process exited with status " + String.valueOf(evaluatorProcessExitValue))
        .build());
  }

  // Shutdown Mesos Executor
  this.executorService.shutdown();
  this.mesosExecutorDriver.stop();
}
 

/**
 * Invoked when a task running within this executor has been killed
 * (via SchedulerDriver::killTask). Note that no status update will
 * be sent on behalf of the executor, the executor is responsible
 * for creating a new TaskStatus (i.e., with TASK_KILLED) and
 * invoking ExecutorDriver::sendStatusUpdate.
 */
@Override
public void killTask(ExecutorDriver driver, TaskID taskId) {
	String tid = taskId.getValue();
	CmdInfo cmdInfo = cmdInfoById.get(tid);
	if (cmdInfo == null) return;
	cmdInfo.cmd.kill();
}
 

/**
 * Signal to the storm-core supervisor that this task should be started.
 */
public int register(TaskID taskId) throws IllegalArgumentException {
  int port = MesosCommon.portFromTaskId(taskId.getValue());
  AssignmentInfo existingAssignment = portAssignments.get(port);
  if (existingAssignment != null) {
    throw new IllegalArgumentException(
                  String.format("Refusing to register task %s because its port " +
                                "%d is already registered for task %s",
                                taskId.getValue(), port,
                                existingAssignment.taskId.getValue()));
  }
  portAssignments.put(port, new AssignmentInfo(taskId, TaskState.ACTIVE));
  return port;
}
 

public void doRegistration(final SchedulerDriver driver, Protos.FrameworkID id) {
  _driver = driver;
  // Now that we've set the driver, we can create our scheduler
  _stormScheduler = new StormSchedulerImpl(_driver);

  _state.put(FRAMEWORK_ID, id.getValue());
  _offers = new HashMap<Protos.OfferID, Protos.Offer>();

  if (_enabledLogviewerSidecar) {

    _timer.scheduleAtFixedRate(new TimerTask() {
      @Override
      public void run() {
        // performing "explicit" reconciliation; master will respond with the latest state for all logviewer tasks
        // in the framework scheduler's statusUpdate() method
        List<TaskStatus> taskStatuses = new ArrayList<TaskStatus>();
        List<String> logviewerPaths = _zkClient.getChildren(_logviewerZkDir);
        if (logviewerPaths == null) {
          _driver.reconcileTasks(taskStatuses);
          return;
        }
        for (String path : logviewerPaths) {
          TaskID logviewerTaskId = TaskID.newBuilder()
                                         .setValue(new String(_zkClient.getNodeData(String.format("%s/%s", _logviewerZkDir, path))))
                                         .build();
          TaskStatus logviewerTaskStatus = TaskStatus.newBuilder()
                                                     .setTaskId(logviewerTaskId)
                                                     .setState(TaskState.TASK_RUNNING)
                                                     .build();
          taskStatuses.add(logviewerTaskStatus);
        }
        _driver.reconcileTasks(taskStatuses);
        LOG.info("Performing task reconciliation between scheduler and master on following tasks: {}", taskStatusListToTaskIDsString(taskStatuses));
      }
    }, 0, TASK_RECONCILIATION_INTERVAL); // reconciliation performed every 5 minutes
  }
}
 

/**
 * We assume a long-running Mesos Task that manages a REEF Evaluator process, leveraging Mesos Executor's interface.
 */
@Override
public void launchTask(final ExecutorDriver driver, final TaskInfo task) {
  driver.sendStatusUpdate(TaskStatus.newBuilder()
      .setTaskId(TaskID.newBuilder().setValue(this.mesosExecutorId).build())
      .setState(TaskState.TASK_STARTING)
      .setSlaveId(task.getSlaveId())
      .setMessage(this.mesosRemoteManager.getMyIdentifier())
      .build());
}
 

@Test 
public void testMakeTestKillable() throws Exception {
  SchedulerState sState = initialize();
  TaskID idOne = TaskID.newBuilder().setValue("Task1").build();
  TaskID idTwo = TaskID.newBuilder().setValue("Task2").build();
  sState.makeTaskKillable(idOne);
  sState.makeTaskKillable(idTwo);
  assertEquals(2, sState.getKillableTaskIds().size());
  assertTrue(sState.getKillableTaskIds().contains(idOne));  
  assertTrue(sState.getKillableTaskIds().contains(idTwo));    
}
 

public void changeState(TaskID task, TaskState state) {
  assertNotStopped();

  checkState(activeTasks.containsKey(task), "Task " + task + " does not exist.");
  Futures.getUnchecked(schedulerFuture).statusUpdate(this, TaskStatus.newBuilder()
      .setTaskId(task)
      .setState(state)
      .build());
}
 

@Test
public void testMakeTestActive() throws Exception {
  SchedulerState sState = initialize();
  TaskID idOne = TaskID.newBuilder().setValue("Task1").build();
  TaskID idTwo = TaskID.newBuilder().setValue("Task2").build();
  sState.addTask(idOne, taskOne);
  sState.addTask(idTwo, taskTwo);
  sState.makeTaskActive(idOne);
  sState.makeTaskActive(idTwo);
  assertTrue(sState.getActiveTasks().contains(taskOne));
  assertTrue(sState.getActiveTasks().contains(taskTwo));
}
 

@Test
public void testAddTask() throws Exception {
  SchedulerState sState = initialize();
  TaskID idOne = TaskID.newBuilder().setValue("Task1").build();
  TaskID idTwo = TaskID.newBuilder().setValue("Task2").build();
  sState.addTask(idOne, taskOne);
  sState.addTask(idTwo, taskTwo);
  assertEquals("zero", sState.getTask(idOne).getProfile().getName());
  assertEquals("low", sState.getTask(idTwo).getProfile().getName());
}
 

private void sendStatusUpdate(
  SingularityExecutorTask task,
  Protos.TaskState taskState,
  String message
) {
  executorUtils.sendStatusUpdate(
    task.getDriver(),
    TaskID.newBuilder().setValue(task.getTaskId()).build(),
    taskState,
    message,
    task.getLog()
  );
}
 

@Override
public boolean killJob(String jobId, ResourceNode node) {
    TaskID id = (TaskID)map.get(jobId);
    driver.killTask(id);
    Status status = driver.killTask(id);
    if (status != Status.DRIVER_RUNNING)
        throw new MesosFrameworkException("Mesos Schedule Driver is dead: "+status.toString());
    return true;
}
 

private void handlePendingTask(TaskID taskId) {
  /*
   * since task is pending and has not started, simply remove 
   * it from SchedulerState task collection
   */
  schedulerState.removeTask(taskId);
}
 

private void cleanupFailedTask(TaskID taskId, NodeTask task, String stopReason) {
  offerLifecycleManager.declineOutstandingOffers(task.getHostname());
  /*
   * Remove the task from SchedulerState if the task is killable.  Otherwise,
   * mark the task as pending to enable restart.
   */
  if (taskIsKillable(taskId)) {
    schedulerState.removeTask(taskId);
    LOGGER.info("Removed killable, {} task with id {}", stopReason, taskId);
  } else {
    schedulerState.makeTaskPending(taskId);        
    LOGGER.info("Marked as pending {} task with id {}", stopReason, taskId);
  }  
}
 

/**
 * Encapsulates logic that retrieves the collection of killable tasks from the
 * SchedulerState object. If a task is in pending state, the task is simply 
 * removed from SchedulerState. Any tasks in a running state were not successfully
 * killed by Mesos or the callback failed, so the another kill attempt is made.
 */
@Override
public void run() { 
  //If there are 1..n killable tasks, proceed; otherwise, simply return
  if (CollectionUtils.isNotEmpty(schedulerState.getKillableTaskIds())) {
    /*
     * Clone the killable task collection, iterate through all tasks, and 
     * process any pending and/or non-pending tasks
     */
    Set<TaskID> killableTasks = Sets.newHashSet(schedulerState.getKillableTaskIds());
    Status driverStatus = driverManager.getDriverStatus();

    //TODO (hokiegeek2) Can the DriverManager be restarted? If not, should the ResourceManager stop?
    if (Status.DRIVER_RUNNING != driverStatus) {
      LOGGER.warn("Cannot kill tasks because Mesos Driver is not running. Status: {}", driverStatus);
      return;
    }

    for (TaskID taskIdToKill : killableTasks) {
      LOGGER.info("Received task kill request for task: {}", taskIdToKill);
      if (isPendingTask(taskIdToKill)) {
        handlePendingTask(taskIdToKill);
      } else {
        handleNonPendingTask(taskIdToKill);
      }
    }
  }
}
 

@Override
public void run() {
	Set<String> killableTasks = schedulerState.getKillableTasks();

	if (CollectionUtils.isEmpty(killableTasks)) {
		return;
	}

	Status driverStatus = driverManager.getDriverStatus();
	if (Status.DRIVER_RUNNING != driverStatus) {
		LOGGER.warn(
				"Cannot kill tasks, as driver is not running. Status: {}",
				driverStatus);
		return;
	}

	Iterator<String> iterator = killableTasks.iterator();

	while (iterator.hasNext()) {
		String taskIdToKill = iterator.next();
		NodeTask task = this.schedulerState.getTask(taskIdToKill);
		TaskID mesosTaskId = task.getMesosTaskId();
		Status status = this.driverManager.kill(mesosTaskId);
		this.schedulerState.removeTask(taskIdToKill);
		Preconditions.checkState(status == Status.DRIVER_RUNNING);
	}
}