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

public CassandraDaemonTask updateConfig(CassandraConfig cassandraConfig,
                                        ExecutorConfig executorConfig,
                                        UUID targetConfigName) {
    LOGGER.info("Updating config for task: {} to config: {}", getTaskInfo().getName(), targetConfigName.toString());
    final Protos.Label label = Protos.Label.newBuilder()
            .setKey("config_target")
            .setValue(targetConfigName.toString())
            .build();
    return new CassandraDaemonTask(getBuilder()
        .setExecutor(getExecutor().update(executorConfig).getExecutorInfo())
        .setTaskId(createId(getName()))
        .setData(getData().withNewConfig(cassandraConfig).getBytes())
        .clearResources()
        .addAllResources(TaskUtils.updateResources(
            cassandraConfig.getCpus(),
            cassandraConfig.getMemoryMb(),
            getTaskInfo().getResourcesList()
        ))
        .clearLabels()
        .setLabels(Protos.Labels.newBuilder().addLabels(label).build()).build());
}
 

private Set<UUID> getTaskConfigs() {
    final Collection<Protos.TaskInfo> taskInfos = stateStore.fetchTasks();
    final Set<UUID> activeConfigs = new HashSet<>();
    try {
        for (Protos.TaskInfo taskInfo : taskInfos) {
            final Protos.Labels labels = taskInfo.getLabels();
            for(Protos.Label label : labels.getLabelsList()) {
                if (label.getKey().equals(PersistentOfferRequirementProvider.CONFIG_TARGET_KEY)) {
                    activeConfigs.add(UUID.fromString(label.getValue()));
                }
            }
        }
    } catch (Exception e) {
        LOGGER.error("Failed to fetch configurations from taskInfos.", e);
    }

    return activeConfigs;
}
 

@Test
public void testCreateVipLabelOnOverlay() {
    Protos.Port.Builder portBuilder = newPortBuilder();
    AuxLabelAccess.setVIPLabels(portBuilder, newVIPSpec("vip", 5, "dcos"));
    Collection<Protos.Label> labels = portBuilder.getLabels().getLabelsList();
    assertEquals(2, labels.size());
    assertEquals(1, labels.stream()
            .filter(label -> label.getKey().startsWith("VIP_") && label.getValue().equals("vip:5"))
            .collect(Collectors.toList()).size());
    assertEquals(1, labels.stream()
            .filter(label -> label.getKey().equals(LabelConstants.VIP_OVERLAY_FLAG_KEY) &&
                    label.getValue().equals(LabelConstants.VIP_OVERLAY_FLAG_VALUE))
            .collect(Collectors.toList()).size());

    Collection<EndpointUtils.VipInfo> vips =
            AuxLabelAccess.getVIPsFromLabels(TestConstants.TASK_NAME, portBuilder.build());
    assertEquals(1, vips.size());
    EndpointUtils.VipInfo vip = vips.iterator().next();
    assertEquals("vip", vip.getVipName());
    assertEquals(5, vip.getVipPort());
}
 

private static Protos.TaskInfo updateConfigLabel(String configName, Protos.TaskInfo taskInfo) {
    final Protos.Labels.Builder labelsBuilder = Protos.Labels.newBuilder();

    final Protos.Labels labels = taskInfo.getLabels();
    for (Protos.Label label : labels.getLabelsList()) {
        final String key = label.getKey();
        if (!CONFIG_TARGET_KEY.equals(key)) {
            labelsBuilder.addLabels(label);
        }
    }

    labelsBuilder.addLabels(Protos.Label.newBuilder()
            .setKey(CONFIG_TARGET_KEY)
            .setValue(configName));
    return Protos.TaskInfo.newBuilder(taskInfo)
            .clearLabels()
            .setLabels(labelsBuilder.build())
            .build();
}
 

@Test
public void testCreateVipLabel() {
    Protos.Port.Builder portBuilder = newPortBuilder();
    AuxLabelAccess.setVIPLabels(portBuilder, newVIPSpec("vip", 5));
    Collection<Protos.Label> labels = portBuilder.getLabels().getLabelsList();
    assertEquals(1, labels.size());
    Label label = labels.iterator().next();
    assertTrue(label.getKey().startsWith("VIP_"));
    assertEquals("vip:5", label.getValue());

    Collection<EndpointUtils.VipInfo> vips =
            AuxLabelAccess.getVIPsFromLabels(TestConstants.TASK_NAME, portBuilder.build());
    assertEquals(1, vips.size());
    EndpointUtils.VipInfo vip = vips.iterator().next();
    assertEquals("vip", vip.getVipName());
    assertEquals(5, vip.getVipPort());
}
 

private Protos.Labels getLabels(Configuration configuration) {
  Protos.Labels.Builder builder = Protos.Labels.newBuilder();
  Map<String, String> labels = configuration.getTaskLabels();
  for (Map.Entry<String, String> kvp : labels.entrySet()) {
    Protos.Label label = Protos.Label.newBuilder()
      .setKey(kvp.getKey())
      .setValue(kvp.getValue())
      .build();

    builder.addLabels(label);
  }

  return builder.build();
}
 

@Test
public void testDiscoveryInfoOnBridgeNetwork() throws Exception {
    Protos.Resource offeredPorts = ResourceTestUtils.getUnreservedPorts(10000, 10000);
    Protos.Offer offer = OfferTestUtils.getOffer(offeredPorts);

    Integer containerPort = 10000;  // non-offered port
    String bridgeNetwork = "mesos-bridge";

    PodInfoBuilder podInfoBuilder = getPodInfoBuilderOnNetwork(
            containerPort, Collections.emptyList(), Optional.of(bridgeNetwork));

    NamedVIPEvaluationStage vipEvaluationStage = getEvaluationStageOnNetwork(
            containerPort, Optional.empty(), Optional.of(bridgeNetwork));

    EvaluationOutcome outcome = vipEvaluationStage.evaluate(
            new MesosResourcePool(offer, Optional.of(Constants.ANY_ROLE)), podInfoBuilder);
    Assert.assertTrue(outcome.isPassing());

    Protos.DiscoveryInfo discoveryInfo = podInfoBuilder.getTaskBuilder(TestConstants.TASK_NAME).getDiscovery();
    String expectedName = TestConstants.POD_TYPE + "-0-" + TestConstants.TASK_NAME;
    String observedName = discoveryInfo.getName();
    Assert.assertEquals(expectedName, observedName);
    Assert.assertEquals(DiscoveryInfo.Visibility.CLUSTER, discoveryInfo.getVisibility());
    Protos.TaskInfo.Builder taskBuilder = podInfoBuilder.getTaskBuilder(TestConstants.TASK_NAME);
    Assert.assertEquals(1, taskBuilder.getResourcesCount());  // expect that bridge uses ports
    Protos.Port port = discoveryInfo.getPorts().getPorts(0);
    Assert.assertEquals(port.getNumber(), containerPort.longValue());
    Assert.assertEquals(port.getProtocol(), "sctp");

    Assert.assertEquals(2, port.getLabels().getLabelsCount());

    Protos.Label vipLabel = port.getLabels().getLabels(0);
    Assert.assertTrue(vipLabel.getKey().startsWith("VIP_"));
    Assert.assertEquals(vipLabel.getValue(), "test-vip:80");

    assertIsBridgeLabel(port.getLabels().getLabels(1));
}
 

@Test
public void testDiscoveryInfoPopulated() throws Exception {
    Protos.Resource offeredPorts = ResourceTestUtils.getUnreservedPorts(10000, 10000);
    Protos.Offer offer = OfferTestUtils.getOffer(offeredPorts);

    PodInfoBuilder podInfoBuilder = getPodInfoBuilderOnNetwork(10000, Collections.emptyList(), Optional.empty());

    // Evaluate stage
    NamedVIPEvaluationStage vipEvaluationStage = getEvaluationStageOnNetwork(10000, Optional.empty(), Optional.empty());
    EvaluationOutcome outcome = vipEvaluationStage.evaluate(
            new MesosResourcePool(offer, Optional.of(Constants.ANY_ROLE)), podInfoBuilder);
    Assert.assertTrue(outcome.isPassing());

    Protos.DiscoveryInfo discoveryInfo = podInfoBuilder.getTaskBuilder(TestConstants.TASK_NAME).getDiscovery();
    String expectedName = TestConstants.POD_TYPE + "-0-" + TestConstants.TASK_NAME;
    String observedName = discoveryInfo.getName();
    Assert.assertEquals(expectedName, observedName);
    Assert.assertEquals(DiscoveryInfo.Visibility.CLUSTER, discoveryInfo.getVisibility());

    Protos.Port port = discoveryInfo.getPorts().getPorts(0);
    Assert.assertEquals(port.getNumber(), 10000);
    Assert.assertEquals(port.getProtocol(), "sctp");
    Assert.assertEquals(1, port.getLabels().getLabelsCount());
    Protos.Label vipLabel = port.getLabels().getLabels(0);
    Assert.assertEquals("pod-type-0-test-task-name", discoveryInfo.getName());
    Assert.assertTrue(vipLabel.getKey().startsWith("VIP_"));
    Assert.assertEquals(vipLabel.getValue(), "test-vip:80");
}
 

@Test
public void labelsAreCorrect() {
    stage.evaluate(new MesosResourcePool(offer, Optional.of(Constants.ANY_ROLE)), podInfoBuilder);
    Protos.TaskInfo.Builder taskBuilder = podInfoBuilder.getTaskBuilder(TestConstants.TASK_NAME);

    Assert.assertEquals(6, taskBuilder.getLabels().getLabelsCount());

    // labels are sorted alphabetically (see LabelUtils):
    Protos.Label label = taskBuilder.getLabels().getLabels(0);
    Assert.assertEquals("index", label.getKey());
    Assert.assertEquals(Integer.toString(TestConstants.TASK_INDEX), label.getValue());

    label = taskBuilder.getLabels().getLabels(1);
    Assert.assertEquals(label.getKey(), "label1");
    Assert.assertEquals("label1-value", label.getValue());

    label = taskBuilder.getLabels().getLabels(2);
    Assert.assertEquals("offer_attributes", label.getKey());
    Assert.assertEquals("", label.getValue());

    label = taskBuilder.getLabels().getLabels(3);
    Assert.assertEquals("offer_hostname", label.getKey());
    Assert.assertEquals(TestConstants.HOSTNAME, label.getValue());

    label = taskBuilder.getLabels().getLabels(4);
    Assert.assertEquals("target_configuration", label.getKey());
    Assert.assertEquals(36, label.getValue().length());

    label = taskBuilder.getLabels().getLabels(5);
    Assert.assertEquals(label.getKey(), "task_type");
    Assert.assertEquals(TestConstants.POD_TYPE, label.getValue());
}
 

private static String getConfigName(Protos.TaskInfo taskInfo) {
    for (Protos.Label label : taskInfo.getLabels().getLabelsList()) {
        if (label.getKey().equals("config_target")) {
            return label.getValue();
        }
    }

    return null;
}
 

private void replaceDuplicateConfig(Protos.TaskInfo taskInfo,
                                    StateStore stateStore,
                                    List<String> duplicateConfigs,
                                    UUID targetName) throws ConfigStoreException {
    try {
        final String taskConfigName = getConfigName(taskInfo);
        final String targetConfigName = targetName.toString();

        for(String duplicateConfig : duplicateConfigs) {
            if (duplicateConfig.equals(taskConfigName)) {
                final Protos.Label configTargetKeyLabel = Protos.Label.newBuilder()
                        .setKey(PersistentOfferRequirementProvider.CONFIG_TARGET_KEY)
                        .setValue(targetConfigName)
                        .build();

                final Protos.Labels labels = Protos.Labels.newBuilder()
                        .addLabels(configTargetKeyLabel)
                        .build();

                final Protos.TaskInfo updatedTaskInfo = Protos.TaskInfo.newBuilder(taskInfo)
                        .setLabels(labels).build();
                stateStore.storeTasks(Arrays.asList(updatedTaskInfo));
                LOGGER.info("Updated task: {} from duplicate config: {} to current target: {}",
                        updatedTaskInfo.getName(), taskConfigName, targetConfigName);
                return;
            }
        }
        LOGGER.info("Task: {} is update to date with target config: {}", taskInfo.getName(), targetConfigName);
    } catch (Exception e) {
        LOGGER.error("Failed to replace duplicate config for task: {} Reason: {}", taskInfo, e);
        throw new ConfigStoreException(e);
    }
}
 

private Optional<String> getTaskConfig(CassandraDaemonTask task) {
    final Protos.TaskInfo taskInfo = task.getTaskInfo();
    if (!taskInfo.hasLabels() || CollectionUtils.isEmpty(taskInfo.getLabels().getLabelsList())) {
        return Optional.empty();
    }

    for (Protos.Label label : taskInfo.getLabels().getLabelsList()) {
        final String key = label.getKey();
        if (PersistentOfferRequirementProvider.CONFIG_TARGET_KEY.equals(key)) {
            return Optional.ofNullable(label.getValue());
        }
    }

    return Optional.empty();
}
 

public CassandraDaemonTask move(
        CassandraDaemonTask currentDaemon, CassandraTaskExecutor executor) {
    final List<Protos.Label> labelsList =
            currentDaemon.getTaskInfo().getLabels().getLabelsList();
    String configName = "";
    for (Protos.Label label : labelsList) {
        if ("config_target".equals(label.getKey())) {
            configName = label.getValue();
        }
    }
    if (StringUtils.isBlank(configName)) {
        throw new IllegalStateException("Task should have 'config_target'");
    }
    CassandraDaemonTask replacementDaemon = new CassandraDaemonTask(
            currentDaemon.getName(),
            configName,
            executor,
            currentDaemon.getConfig(),
            capabilities,
            currentDaemon.getData().replacing(currentDaemon.getData().getHostname()));

    Protos.TaskInfo taskInfo = Protos.TaskInfo.newBuilder(replacementDaemon.getTaskInfo())
            .setTaskId(currentDaemon.getTaskInfo().getTaskId())
            .build();

    return new CassandraDaemonTask(taskInfo);
}
 

public static Protos.TaskInfo getTaskInfo(List<Protos.Resource> resources, Integer index) {
    Protos.TaskInfo.Builder builder = Protos.TaskInfo.newBuilder()
            .setTaskId(TestConstants.TASK_ID)
            .setName(TestConstants.TASK_NAME)
            .setSlaveId(TestConstants.AGENT_ID)
            .setCommand(TestConstants.COMMAND_INFO)
            .setContainer(TestConstants.CONTAINER_INFO);
    builder.setLabels(new TaskLabelWriter(builder)
            .setType(TestConstants.TASK_TYPE)
            .setIndex(index)
            .toProto());
    for (Protos.Resource r : resources) {
        String resourceId = "";
        String dynamicPortAssignment = null;
        String vipAssignment = null;
        for (Protos.Label l : r.getReservation().getLabels().getLabelsList()) {
            if (Objects.equals(l.getKey(), "resource_id")) {
               resourceId = l.getValue();
            } else if (Objects.equals(l.getKey(), TestConstants.HAS_DYNAMIC_PORT_ASSIGNMENT_LABEL)) {
                dynamicPortAssignment = l.getValue();
            } else if (Objects.equals(l.getKey(), TestConstants.HAS_VIP_LABEL)) {
                vipAssignment = l.getValue();
            }
        }

        if (Objects.equals(r.getName(), "ports")) {
            String portValue = dynamicPortAssignment == null ?
                    Long.toString(r.getRanges().getRange(0).getBegin()) : dynamicPortAssignment;

            if (!resourceId.isEmpty()) {
                builder.getCommandBuilder()
                        .getEnvironmentBuilder()
                        .addVariablesBuilder()
                        .setName(TestConstants.PORT_ENV_NAME)
                        .setValue(portValue);
            }
            if (!resourceId.isEmpty() && vipAssignment != null) {
                Protos.DiscoveryInfo.Builder discoveryBuilder = builder.getDiscoveryBuilder();
                discoveryBuilder.setVisibility(Protos.DiscoveryInfo.Visibility.CLUSTER);
                discoveryBuilder.setName(builder.getName());
                discoveryBuilder.getPortsBuilder()
                        .addPortsBuilder()
                        .setNumber(Integer.parseInt(portValue))
                        .getLabelsBuilder()
                        .addLabelsBuilder()
                        .setKey("VIP_" + UUID.randomUUID().toString())
                        .setValue(vipAssignment);
            }
        }
    }
    return builder.addAllResources(resources).build();
}
 

@Test
public void testCreateSucceeds() throws Exception {
    Protos.Resource offeredResource = ResourceTestUtils.getUnreservedMountVolume(2000, Optional.empty());
    Protos.Offer offer = OfferTestUtils.getCompleteOffer(offeredResource);

    MesosResourcePool mesosResourcePool = new MesosResourcePool(offer, Optional.of(Constants.ANY_ROLE));
    PodInstanceRequirement podInstanceRequirement =
            PodInstanceRequirementTestUtils.getMountVolumeRequirement(1.0, 1000);

    VolumeEvaluationStage volumeEvaluationStage = VolumeEvaluationStage.getNew(
            getVolumeSpec(podInstanceRequirement.getPodInstance()),
            Collections.singleton(getTaskName(podInstanceRequirement.getPodInstance())),
            Optional.empty(),
            Optional.empty());
    EvaluationOutcome outcome =
            volumeEvaluationStage.evaluate(
                    mesosResourcePool,
                    new PodInfoBuilder(
                            podInstanceRequirement,
                            TestConstants.SERVICE_NAME,
                            UUID.randomUUID(),
                            PodTestUtils.getTemplateUrlFactory(),
                            SchedulerConfigTestUtils.getTestSchedulerConfig(),
                            Collections.emptyList(),
                            TestConstants.FRAMEWORK_ID,
                            Collections.emptyMap()));
    Assert.assertTrue(outcome.isPassing());

    List<OfferRecommendation> recommendations = new ArrayList<>(outcome.getOfferRecommendations());
    Assert.assertEquals(2, outcome.getOfferRecommendations().size());

    OfferRecommendation reserveRecommendation = recommendations.get(0);
    Assert.assertEquals(Protos.Offer.Operation.Type.RESERVE, reserveRecommendation.getOperation().get().getType());

    Protos.Resource resource = reserveRecommendation.getOperation().get().getReserve().getResources(0);
    Assert.assertEquals("disk", resource.getName());
    Assert.assertEquals(resource.getScalar(), offeredResource.getScalar());
    Protos.Resource.ReservationInfo reservationInfo = ResourceUtils.getReservation(resource).get();
    Protos.Label reservationLabel = reservationInfo.getLabels().getLabels(0);
    Assert.assertEquals(reservationLabel.getKey(), "resource_id");
    Assert.assertNotEquals(reservationLabel.getValue(), "");

    OfferRecommendation createRecommendation = recommendations.get(1);
    resource = createRecommendation.getOperation().get().getCreate().getVolumes(0);
    Assert.assertEquals(Protos.Offer.Operation.Type.CREATE, createRecommendation.getOperation().get().getType());
    Assert.assertEquals("disk", resource.getName());
    Assert.assertEquals(resource.getScalar(), offeredResource.getScalar());
    reservationInfo = ResourceUtils.getReservation(resource).get();
    reservationLabel = reservationInfo.getLabels().getLabels(0);
    Assert.assertEquals(reservationLabel.getKey(), "resource_id");
    Assert.assertNotEquals(reservationLabel.getValue(), "");
    Assert.assertNotEquals(resource.getDisk().getPersistence().getId(), "");
}
 

private static void assertIsOverlayLabel(Protos.Label label) {
    Assert.assertEquals("network-scope", label.getKey());
    Assert.assertEquals("container", label.getValue());
}
 

private static void assertIsBridgeLabel(Protos.Label label) {
    Assert.assertEquals("network-scope", label.getKey());
    Assert.assertEquals("host", label.getValue());
}
 

protected CassandraTask(
    final String name,
    final String configName,
    final CassandraTaskExecutor executor,
    final double cpus,
    final int memoryMb,
    final int diskMb,
    final VolumeRequirement.VolumeMode volumeMode,
    final VolumeRequirement.VolumeType volumeType,
    final Collection<Integer> ports,
    @Nullable final DiscoveryInfo discoveryInfo,
    final CassandraData data) {

    String role = executor.getRole();
    String principal = executor.getPrincipal();

    Protos.TaskInfo.Builder builder = Protos.TaskInfo.newBuilder()
        .setTaskId(createId(name))
        .setName(name)
        .setSlaveId(EMPTY_SLAVE_ID)
        .setExecutor(executor.getExecutorInfo())
        .addAllResources(Arrays.asList(
            ResourceUtils.getDesiredScalar(role, principal, "cpus", cpus),
            ResourceUtils.getDesiredScalar(role, principal, "mem", memoryMb)))
        .setData(data.getBytes());
    final Protos.Label label = Protos.Label.newBuilder()
            .setKey("config_target")
            .setValue(configName)
            .build();
    builder.setLabels(Protos.Labels.newBuilder().addLabels(label));

    if (!volumeMode.equals(VolumeRequirement.VolumeMode.NONE)) {
        if (volumeType.equals(VolumeRequirement.VolumeType.MOUNT)) {
            builder.addResources(ResourceUtils.getDesiredMountVolume(role, principal, diskMb, CassandraConfig.VOLUME_PATH));
        } else {
            builder.addResources(ResourceUtils.getDesiredRootVolume(role, principal, diskMb, CassandraConfig.VOLUME_PATH));
        }
    }

    if (!ports.isEmpty()) {
        builder.addResources(ResourceUtils.getDesiredRanges(role, principal, "ports", Algorithms.createRanges(ports)));
    }

    if (discoveryInfo != null) {
        builder.setDiscovery(discoveryInfo);
    }

    info = builder.build();
}