org.apache.mesos.Protos Java Examples

@Test
public void testResourceOfferContainerForAuxServiceWithOutRole() {
  Protos.Offer offer = new OfferBuilder("test.com")
      .addScalarResource("cpus", 2.0)
      .addScalarResource("mem", 8000)
      .addRangeResource("ports", 3500, 3600)
      .build();
  Map<String, Long> ports = new HashMap<>(4);
  ports.put("test1.address", 0L);
  ports.put("test2.address", 0L);
  ports.put("test3.address", 0L);
  ports.put("test4.port", 3501L);
  ServiceResourceProfile profile = new ServiceResourceProfile("jobhistory", 2.0, 8000.0, ports);
  ResourceOfferContainer roc = new ResourceOfferContainer(offer, profile, null);
  assertTrue(roc.getHostName().equals("test.com"));
  assertTrue("Should be satisfied if offer contains request", roc.satisfies(profile));
}
 

/**
 * Test Mesos registration handling.
 */
@Test
public void testRegistered() throws Exception {
	new Context() {{
		startResourceManager();

		Protos.MasterInfo masterInfo = Protos.MasterInfo.newBuilder()
			.setId("master1").setIp(0).setPort(5050).build();
		resourceManager.registered(new Registered(framework1, masterInfo));

		verify(rmServices.workerStore).setFrameworkID(Option.apply(framework1));
		resourceManager.connectionMonitor.expectMsgClass(Registered.class);
		resourceManager.reconciliationCoordinator.expectMsgClass(Registered.class);
		resourceManager.launchCoordinator.expectMsgClass(Registered.class);
		resourceManager.taskRouter.expectMsgClass(Registered.class);
	}};
}
 

/**
 * Maps the provided status to the service that owns its task, then queries that service with the status.
 *
 * <p>This is an optimization which avoids querying services about task statuses that don't relate to them.
 * <p>In addition to reducing unnecessary queries, this also improves isolation between services. They only see
 * task statuses which relate to them.
 */
@Override
public TaskStatusResponse taskStatus(Protos.TaskStatus status) {
  return multiServiceManager
      .getMatchingService(status)
      .map(x -> x.taskStatus(status))
      .orElseGet(() -> multiServiceManager
          .getServiceSanitized(frameworkName)
          .map(x -> {
            LOGGER.info("Forwarding task status to default service: {}", frameworkName);
            return x.taskStatus(status);
          })
          .orElseGet(() -> {
            // Unrecognized service. Status for old task ?
            LOGGER.info("Received status for unknown task {}: {}",
                status.getTaskId().getValue(), TextFormat.shortDebugString(status));
            return TaskStatusResponse.unknownTask();
          })
      );
}
 

public static Protos.Resource getReservedRootVolume(double diskSize, String resourceId, String persistenceId, String frameworkId) {
    VolumeSpec volumeSpec = DefaultVolumeSpec.createRootVolume(
            diskSize,
            TestConstants.CONTAINER_PATH,
            TestConstants.ROLE,
            Constants.ANY_ROLE,
            TestConstants.PRINCIPAL);
    return ResourceBuilder.fromSpec(
            volumeSpec,
            Optional.of(resourceId),
            Optional.empty(),
            Optional.of(persistenceId),
            Optional.empty(),
            Optional.empty(),
            Optional.of(frameworkId))
            .build();
}
 

private void onFinish(SingularityExecutorTask task, Protos.TaskState taskState) {
  processKiller.cancelDestroyFuture(task.getTaskId());

  tasks.remove(task.getTaskId());
  processRunningTasks.remove(task.getTaskId());
  processBuildingTasks.remove(task.getTaskId());

  task.cleanup(taskState);

  ListeningExecutorService executorService = taskToShellCommandPool.remove(
    task.getTaskId()
  );

  if (executorService != null) {
    executorService.shutdownNow();
    try {
      executorService.awaitTermination(5, TimeUnit.MILLISECONDS);
    } catch (InterruptedException e) {
      LOG.warn("Awaiting shutdown of shell executor service", e);
    }
  }

  logging.stopTaskLogger(task.getTaskId(), task.getLogbackLog());

  checkIdleExecutorShutdown(task.getDriver());
}
 

/**
 * This method keeps the resources associated with tasks in the state store up to date, when a task which shares
 * their resource-set is launched.
 *
 * @param podInstance the parent pod associated with the task being launched
 * @param taskInfo    the task being launched
 */
@VisibleForTesting
void updateTaskResourcesWithinResourceSet(PodInstance podInstance, Protos.TaskInfo taskInfo) {
  Optional<TaskSpec> taskSpec = TaskUtils.getTaskSpec(podInstance, taskInfo.getName());
  if (!taskSpec.isPresent()) {
    return;
  }

  // Update any other TaskInfos in this resource set to have the same resources:
  Collection<Protos.TaskInfo> taskInfosWithSameResourceSet =
      getOtherTasksInResourceSet(podInstance, taskSpec.get());
  stateStore.storeTasks(updateTasksWithResources(
      taskInfosWithSameResourceSet,
      taskInfo.getResourcesList(),
      taskInfo.hasExecutor() ?
          Optional.of(taskInfo.getExecutor().getResourcesList())
          : Optional.empty()));
}
 

private Protos.CommandInfo nativeCommand(Configuration configuration, List<String> args, Long elasticSearchNodeId) {
    String address = configuration.getFrameworkFileServerAddress();
    if (address == null) {
        throw new NullPointerException("Webserver address is null");
    }
    String httpPath = address + "/get/" + Configuration.ES_TAR;
    String command = configuration.nativeCommand(args);
    final Protos.Environment environment = Protos.Environment.newBuilder().addAllVariables(new ExecutorEnvironmentalVariables(configuration, elasticSearchNodeId).getList()).build();
    final Protos.CommandInfo.Builder builder = Protos.CommandInfo.newBuilder()
            .setShell(true)
            .setValue(command)
            .setUser("root")
            .mergeEnvironment(environment);
    if (configuration.getElasticsearchBinary().isEmpty()) {
        builder.addUris(Protos.CommandInfo.URI.newBuilder().setValue(httpPath));
    } else {
        builder.addUris(Protos.CommandInfo.URI.newBuilder().setValue(configuration.getElasticsearchBinary()));
    }
    if (!configuration.getElasticsearchSettingsLocation().isEmpty()) {
        builder.addUris(Protos.CommandInfo.URI.newBuilder().setValue(configuration.getElasticsearchSettingsLocation()));
    }
    return builder
            .build();
}
 

@Test
public void testRepeatedStoreStatus() throws Exception {
    Protos.TaskInfo task = createTask(TestConstants.TASK_NAME);
    store.storeTasks(Arrays.asList(task));

    // taskstatus id must exactly match taskinfo id:
    Protos.TaskStatus status = TASK_STATUS.toBuilder().setTaskId(task.getTaskId()).build();
    store.storeStatus(TestConstants.TASK_NAME, status);
    assertEquals(status, store.fetchStatus(TestConstants.TASK_NAME).get());

    store.storeStatus(TestConstants.TASK_NAME, status);
    assertEquals(status, store.fetchStatus(TestConstants.TASK_NAME).get());

    assertEquals(Arrays.asList(TestConstants.TASK_NAME), store.fetchTaskNames());
    Collection<Protos.TaskStatus> statuses = store.fetchStatuses();
    assertEquals(1, statuses.size());
    assertEquals(status, statuses.iterator().next());
}
 

@Test
public void validPortResource() throws Exception {
    ClassLoader classLoader = getClass().getClassLoader();
    File file = new File(classLoader.getResource("valid-multiple-ports.yml").getFile());
    DefaultServiceSpec serviceSpec = DefaultServiceSpec.newGenerator(file, SCHEDULER_CONFIG).build();

    List<ResourceSpec> portsResources = serviceSpec.getPods().get(0).getTasks().get(0).getResourceSet()
            .getResources()
            .stream()
            .filter(r -> r.getName().equals("ports"))
            .collect(Collectors.toList());

    Assert.assertEquals(2, portsResources.size());

    Protos.Value.Ranges http = portsResources.get(0).getValue().getRanges();
    Protos.Value.Ranges another = portsResources.get(1).getValue().getRanges();
    Assert.assertEquals(1, http.getRangeCount());
    Assert.assertEquals(1, another.getRangeCount());
    Assert.assertEquals(8080, http.getRange(0).getBegin(), http.getRange(0).getEnd());
    Assert.assertEquals(8088, another.getRange(0).getBegin(), another.getRange(0).getEnd());
}
 

@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;
	}
}
 

MockMesosResourceManagerRuntimeServices() throws Exception {
	schedulerDriver = mock(SchedulerDriver.class);

	mesosConfig = mock(MesosConfiguration.class);
	when(mesosConfig.frameworkInfo()).thenReturn(Protos.FrameworkInfo.newBuilder());
	when(mesosConfig.withFrameworkInfo(any(Protos.FrameworkInfo.Builder.class))).thenReturn(mesosConfig);
	when(mesosConfig.createDriver(any(Scheduler.class), anyBoolean())).thenReturn(schedulerDriver);

	workerStore = mock(MesosWorkerStore.class);
	when(workerStore.getFrameworkID()).thenReturn(Option.<Protos.FrameworkID>empty());

	artifactServer = mock(MesosArtifactServer.class);
}
 

@Test
public void testNMTaskFactory() {
  NMExecutorCommandLineGenerator clGenerator = new NMExecutorCommandLineGenerator(cfgWithDocker);
  TaskUtils taskUtils = new TaskUtils(cfgWithDocker);
  Protos.Offer offer = new OfferBuilder("test.com")
      .addScalarResource("cpus", 10.0)
      .addScalarResource("mem", 16000)
      .addRangeResource("ports", 3500, 3505)
      .build();
  ServiceResourceProfile profile = new ExtendedResourceProfile(new NMProfile("tooMuchCpu", 7L, 8000L), taskUtils.getNodeManagerCpus(),
      taskUtils.getNodeManagerMemory(), taskUtils.getNodeManagerPorts());
  NodeTask nodeTask = new NodeTask(profile, null);
  ResourceOfferContainer roc = new ResourceOfferContainer(offer, profile, null);
  NMTaskFactory taskFactory = new NMTaskFactory(cfgWithDocker, taskUtils, clGenerator);
  Protos.TaskInfo taskInfo = taskFactory.createTask(roc, frameworkId, makeTaskId("nm.zero"), nodeTask);
  assertFalse("taskInfo should not have a container", taskInfo.hasContainer());
  assertTrue("The container should have an executor", taskInfo.hasExecutor());
  Protos.ExecutorInfo executorInfo = taskInfo.getExecutor();
  assertTrue("executorInfo should have container", executorInfo.hasContainer());
  Protos.ContainerInfo containerInfo = executorInfo.getContainer();
  assertTrue("There should be two volumes", containerInfo.getVolumesCount() == 2);
  assertTrue("The first volume should be read only", containerInfo.getVolumes(0).getMode().equals(Protos.Volume.Mode.RO));
  assertTrue("The first volume should be read write", containerInfo.getVolumes(1).getMode().equals(Protos.Volume.Mode.RW));
  assertTrue("There should be a docker image", containerInfo.getDocker().hasImage());
  assertTrue("The docker image should be mesos/myraid", containerInfo.getDocker().getImage().equals("mesos/myriad"));
  assertTrue("Should be using host networking", containerInfo.getDocker().getNetwork().equals(Protos.ContainerInfo.DockerInfo.Network.HOST));
  assertTrue("There should be two parameters", containerInfo.getDocker().getParametersList().size() == 2);
  assertTrue("Privledged mode should be false", !containerInfo.getDocker().getPrivileged());
}
 

protected void executorTaskError(final Protos.TaskInfo taskInfo) {
    scheduler.statusUpdate(driver, Protos.TaskStatus.newBuilder()
        .setExecutorId(executorId(taskInfo))
        .setHealthy(true)
        .setSlaveId(taskInfo.getSlaveId())
        .setSource(Protos.TaskStatus.Source.SOURCE_EXECUTOR)
        .setTaskId(taskInfo.getTaskId())
        .setTimestamp(System.currentTimeMillis())
        .setState(Protos.TaskState.TASK_ERROR)
            .build());
}
 

@Override
public void offerRescinded(SchedulerDriver driver, final Protos.OfferID offerId) {
	runAsync(new Runnable() {
		@Override
		public void run() {
			MesosResourceManager.this.offerRescinded(new OfferRescinded(offerId));
		}
	});
}
 

@Override
public Collection<String> getKeys(Protos.Offer offer) {
  if (offer.hasDomain() && offer.getDomain().hasFaultDomain()) {
    return Arrays.asList(offer.getDomain().getFaultDomain().getZone().getName());
  } else {
    return Collections.emptyList();
  }
}
 

private Protos.Attribute getTextAttribute(String name, String value) {
  return Protos.Attribute.newBuilder()
    .setName(name)
    .setType(TEXT)
    .setText(Text.newBuilder()
    .setValue(value))
    .build();
}
 

@Test
public void shouldHandleGetNull() throws IOException {
    when(state.get(anyString())).thenReturn(null);
    List<Protos.TaskInfo> taskList = clusterState.getTaskList();
    verify(state, times(1)).get(anyString());
    assertEquals(0, taskList.size());
}
 

/**
 * Publishes a ResourceOffersEvent
 */
@Override
public void resourceOffers(final SchedulerDriver driver, final List<Protos.Offer> offers) {
  disruptorManager.getResourceOffersEventDisruptor().publishEvent(new EventTranslator<ResourceOffersEvent>() {
    @Override
    public void translateTo(ResourceOffersEvent event, long sequence) {
      event.setDriver(driver);
      event.setOffers(offers);
    }
  });
}
 

/**
 * Converts the provided list of zero or more attributes into a string suitable for comparison.
 *
 * @throws IllegalArgumentException if some part of the provided attributes couldn't be
 *                                  serialized
 */
public static String toString(List<Protos.Attribute> attributes) throws IllegalArgumentException {
  StringJoiner joiner = new StringJoiner(ATTRIBUTE_LIST_SEPARATOR);
  for (Protos.Attribute attribute : attributes) {
    joiner.add(toString(attribute));
  }
  return joiner.toString();
}
 

public ContainerStateRule build() {
    Protos.TaskState effectiveMesosTerminalState = mesosTerminalState;
    Protos.TaskStatus.Reason effectiveMesosReasonCode = mesosReasonCode;
    String effectiveReasonMessage = reasonMessage;

    if (mesosReasonCode == null && titusReasonCode != null) {
        switch (titusReasonCode) {
            case TaskStatus.REASON_INVALID_REQUEST:
            case TaskStatus.REASON_CRASHED:
            case TaskStatus.REASON_FAILED:
            case TaskStatus.REASON_LOCAL_SYSTEM_ERROR:
            case TaskStatus.REASON_TRANSIENT_SYSTEM_ERROR:
            case TaskStatus.REASON_UNKNOWN_SYSTEM_ERROR:
                effectiveMesosTerminalState = mesosTerminalState == null ? Protos.TaskState.TASK_FAILED : mesosTerminalState;
                effectiveMesosReasonCode = Protos.TaskStatus.Reason.REASON_COMMAND_EXECUTOR_FAILED;
                break;
            case TaskStatus.REASON_TASK_KILLED:
                effectiveMesosTerminalState = mesosTerminalState == null ? Protos.TaskState.TASK_KILLED : mesosTerminalState;
                effectiveMesosReasonCode = Protos.TaskStatus.Reason.REASON_TASK_KILLED_DURING_LAUNCH;
                break;
            case TaskStatus.REASON_TASK_LOST:
                effectiveMesosTerminalState = mesosTerminalState == null ? Protos.TaskState.TASK_LOST : mesosTerminalState;
                effectiveMesosReasonCode = Protos.TaskStatus.Reason.REASON_TASK_UNKNOWN;
                break;
            default:
                effectiveMesosTerminalState = mesosTerminalState == null ? Protos.TaskState.TASK_UNKNOWN : mesosTerminalState;
                effectiveMesosReasonCode = Protos.TaskStatus.Reason.REASON_TASK_INVALID;
        }
        effectiveReasonMessage = String.format("%s: (from titusReasonCode=%s)", StringExt.safeTrim(reasonMessage), titusReasonCode);
    }

    return new ContainerStateRule(delayInStateMs, action, Optional.ofNullable(effectiveMesosTerminalState),
            Optional.ofNullable(effectiveMesosReasonCode), Optional.ofNullable(effectiveReasonMessage)
    );
}
 

/**
 * Construct a Mesos environment variable.
 */
public static Protos.Environment.Variable variable(String name, String value) {
	checkNotNull(name);
	return Protos.Environment.Variable.newBuilder()
		.setName(name)
		.setValue(value)
		.build();
}
 

public static Optional<Protos.ResourceProviderID> getProviderId(Protos.Resource resource) {
  if (resource.hasProviderId()) {
    return Optional.of(resource.getProviderId());
  }

  return Optional.empty();
}
 

@Override
public <T> T afterImmediate(String injectionPoint, T result) {
    if (!(result instanceof Protos.TaskStatus)) {
        return result;
    }
    Protos.TaskStatus status = (Protos.TaskStatus) result;

    boolean taskMatches = taskState.map(expected -> status.getState() == expected).orElse(false);
    boolean reasonMatches = reasonCode.map(expected -> status.getReason() == expected).orElse(false);
    if (taskMatches || reasonMatches) {
        return (T) status.toBuilder().setMessage(override).build();
    }
    return result;
}
 

@Test
public void getRegionTaskKeys() {
    MaxPerRegionRule rule = new MaxPerRegionRule(2);
    Protos.TaskInfo taskInfo = TestConstants.TASK_INFO.toBuilder()
            .setLabels(
                    new TaskLabelWriter(TestConstants.TASK_INFO)
                            .setRegion(TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getRegion())
                            .toProto())
            .build();
    Collection<String> keys = rule.getKeys(taskInfo);
    Assert.assertEquals(1, keys.size());
    Assert.assertEquals(
            TestConstants.LOCAL_DOMAIN_INFO.getFaultDomain().getRegion().getName(),
            keys.stream().findFirst().get());
}
 

@Override
public void registered(SchedulerDriver driver, Protos.FrameworkID frameworkId, Protos.MasterInfo masterInfo) {
    LOGGER.info("Framework registered as " + frameworkId.getValue());

    List<Protos.Resource> resources = Resources.buildFrameworkResources(configuration);

    Protos.Request request = Protos.Request.newBuilder()
            .addAllResources(resources)
            .build();

    List<Protos.Request> requests = Collections.singletonList(request);
    driver.requestResources(requests);

    frameworkState.markRegistered(frameworkId, driver);
}
 

@Test
public void testTaskInfoWithNoStatusRequiresNoRecovery() throws Exception {
    ConfigStore<ServiceSpec> configStore = newConfigStore(persister);

    Protos.TaskInfo taskInfo = newTaskInfo("name-0-not-present", configStore);

    assertThat(TaskUtils.getTasksNeedingRecovery(configStore, Collections.singleton(taskInfo), Collections.emptyList()),
            is(empty()));
}
 

@Test
public void testFoldered2ToTaskId() throws Exception {
    Protos.TaskID taskId = CommonIdUtils.toTaskId(TEST_FOLDERED_SERVICE_NAME2, TEST_TASK_NAME);
    Assert.assertTrue(taskId.getValue().startsWith(TEST_FOLDERED_SANITIZED_NAME + "__" + TEST_TASK_NAME + "__"));
    Assert.assertNotNull(UUID.fromString(taskId.getValue().split("__")[2]));
    Assert.assertEquals(TEST_TASK_NAME, CommonIdUtils.toTaskName(taskId));
    Assert.assertEquals(TEST_FOLDERED_SANITIZED_NAME, CommonIdUtils.toSanitizedServiceName(taskId).get());
}
 

@Override
public void statusUpdate(SchedulerDriver driver, Protos.TaskStatus status) {
  try {
    LOGGER.info("Received status update for taskId={} state={} message={} protobuf={}",
        status.getTaskId().getValue(),
        status.getState().toString(),
        status.getMessage(),
        TextFormat.shortDebugString(status));
    Metrics.record(status);
    TaskStatusResponse response = mesosEventClient.taskStatus(status);
    boolean eligibleToKill = TaskKiller.update(status);
    switch (response.result) { // SUPPRESS CHECKSTYLE MissingSwitchDefaultCheck
      case UNKNOWN_TASK:
        if (eligibleToKill) {
          LOGGER.info("Received status update for unknown task, marking task to be killed: {}",
              status.getTaskId().getValue());
          TaskKiller.killTask(status.getTaskId());
        } else {
          // Special case: Mesos can send TASK_LOST+REASON_RECONCILIATION as a response to a
          // prior kill request against a task that is unknown to Mesos. When this happens, we
          // don't want to repeat the kill, because that would create a Kill -> Status -> Kill
          // -> ... loop
          LOGGER.warn(
              "Received status update for unknown task, but task should not be killed again: {}",
              status.getTaskId().getValue());
        }
        break;
      case PROCESSED:
        // No-op
        break;
    }
  } catch (Throwable e) {
    logExceptionAndExit(e);
  }
}
 

@Override
public CleanupStatus createStatus(
        Protos.TaskState state,
        Optional<String> message) {

    Protos.TaskStatus.Builder builder = getStatusBuilder();
    if (message.isPresent()) {
        builder.setMessage(message.get());
    }

    return CleanupStatus.create(builder
        .setData(CassandraData.createCleanupStatusData().getBytes())
        .setState(state)
        .build());
}
 

/**
 * Creates an allocation of resources for tasks to take.
 *
 * @param resources the resources to add to the allocation.
 */
public MesosResourceAllocation(Collection<Protos.Resource> resources) {
	this.resources = new ArrayList<>(checkNotNull(resources));

	// sort the resources to prefer reserved resources
	this.resources.sort(Comparator.comparing(r -> UNRESERVED_ROLE.equals(r.getRole())));
}