org.apache.mesos.v1.Protos Java Examples

@Override
public void drainForInverseOffer(Protos.InverseOffer offer) {
  // TaskStore does not allow for querying by agent id.
  Optional<String> hostname = getHostname(offer);

  if (hostname.isPresent()) {
    String host = hostname.get();
    storage.write(storeProvider -> {
      // Create a dummy maintenance request zero percent sla and timeout to force drain.
      recordMaintenanceRequests(storeProvider, ImmutableSet.of(host), ZERO_PERCENT_SLA, 0);
      return drainTasksOnHost(host, storeProvider);
    });
  } else {
    LOG.error("Unable to drain tasks on agent {} because "
        + "no hostname attached to inverse offer {}.", offer.getAgentId(), offer.getId());
  }
}
 

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()
  );
}
 

@Override
public SchedulerDriver create(
    Scheduler scheduler,
    Optional<Protos.Credential> credentials,
    Protos.FrameworkInfo frameworkInfo,
    String master) {

  FrameworkInfo convertedFrameworkInfo = convert(frameworkInfo);
  Optional<Credential> convertedCredentials = credentials.map(ProtosConversion::convert);

  if (credentials.isPresent()) {
    return new MesosSchedulerDriver(
        scheduler,
        convertedFrameworkInfo,
        master,
        false, // Disable implicit acknowledgements.
        convertedCredentials.get());
  } else {
    return new MesosSchedulerDriver(
        scheduler,
        convertedFrameworkInfo,
        master,
        false); // Disable implicit acknowledgements.
  }
}
 

@Subscribe
public void accumulate(TaskStatusReceived event) {
  if (event.getState() == Protos.TaskState.TASK_LOST && event.getSource().isPresent()) {
    lostSourceCounters.getUnchecked(event.getSource().get()).incrementAndGet();
  }

  if (event.getReason().isPresent()) {
    reasonCounters.getUnchecked(event.getReason().get()).incrementAndGet();
  }

  if (event.getSource().isPresent() && event.getEpochTimestampMicros().isPresent()) {
    long nowMicros = clock.nowMillis() * 1000;
    // Avoid distorting stats by recording zero or negative values.  This can result if delivery
    // is faster than the clock resolution (1 ms) or there is clock skew between the systems.
    // In reality, this value is likely to be inaccurate, especially at the resolution of millis.
    if (event.getEpochTimestampMicros().get() < nowMicros) {
      latencyTimers.getUnchecked(event.getSource().get())
          .requestComplete(nowMicros - event.getEpochTimestampMicros().get());
    } else {
      LOG.debug("Not recording stats for status update with timestamp <= now");
    }
  }
}
 

public SingularityOfferHolder(
  List<Protos.Offer> offers,
  int taskSizeHint,
  String rackId,
  String slaveId,
  String hostname,
  Map<String, String> textAttributes,
  Map<String, String> reservedSlaveAttributes
) {
  this.rackId = rackId;
  this.slaveId = slaveId;
  this.hostname = hostname;
  this.offers = offers;
  this.roles = MesosUtils.getRoles(offers.get(0));
  this.acceptedTasks = Lists.newArrayListWithExpectedSize(taskSizeHint);
  this.currentResources =
    offers.size() > 1
      ? MesosUtils.combineResources(
        offers.stream().map(Protos.Offer::getResourcesList).collect(Collectors.toList())
      )
      : offers.get(0).getResourcesList();
  this.sanitizedHost = JavaUtils.getReplaceHyphensWithUnderscores(hostname);
  this.sanitizedRackId = JavaUtils.getReplaceHyphensWithUnderscores(rackId);
  this.textAttributes = textAttributes;
  this.reservedSlaveAttributes = reservedSlaveAttributes;
}
 

/**
 * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE}.
 * <p>
 *
 * @param frameworkId               The frameworkId to set on the {@link Protos.FrameworkInfo FrameworkInfo} and
 *                                  {@link Call Call} messages.
 * @param user                      The user to set on the {@link Protos.FrameworkInfo FrameworkInfo} message.
 * @param frameworkName             The name to set on the {@link Protos.FrameworkInfo FrameworkInfo} message.
 * @param failoverTimeoutSeconds    The failoverTimeoutSeconds to set on the
 *                                  {@link Protos.FrameworkInfo FrameworkInfo} message.
 * @return An {@link Call Call} of type {@link Type#SUBSCRIBE SUBSCRIBE} with the configured
 * {@link Subscribe Subscribe} sub-message.
 */
@NotNull
public static Call subscribe(
    @NotNull final Protos.FrameworkID frameworkId,
    @NotNull final String user,
    @NotNull final String frameworkName,
    final long failoverTimeoutSeconds
) {
    final Protos.FrameworkInfo frameworkInfo = Protos.FrameworkInfo.newBuilder()
        .setId(frameworkId)
        .setUser(user)
        .setName(frameworkName)
        .setFailoverTimeout(failoverTimeoutSeconds)
        .build();
    return subscribe(frameworkId, frameworkInfo);
}
 

@Test
public void testFrameworkInfoAllowGpu() {
  Protos.FrameworkInfo info = CommandLineDriverSettingsModule.buildFrameworkInfo(
      "aurora",
      "user",
      Optional.empty(),
      Amount.of(1L, Time.MINUTES),
      false, // revocable
      true, // allow gpu
      false, // partition aware
      Optional.empty());
  assertEquals("", info.getPrincipal());
  assertEquals(1, info.getCapabilitiesCount());
  assertEquals(GPU_RESOURCES, info.getCapabilities(0).getType());
  assertFalse(info.hasRole());
}
 

/**
 * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#ACKNOWLEDGE ACKNOWLEDGE}.
 * <p>
 *
 * @param frameworkId    The {@link Protos.FrameworkID} to be set on the {@link Call}
 * @param uuid           The {@link Protos.TaskStatus#getUuid() uuid} from the
 *                       {@link org.apache.mesos.v1.scheduler.Protos.Event.Update#getStatus() TaskStatus} received from Mesos.
 * @param agentId        The {@link Protos.TaskStatus#getAgentId() agentId} from the
 *                       {@link org.apache.mesos.v1.scheduler.Protos.Event.Update#getStatus() TaskStatus} received from Mesos.
 * @param taskId         The {@link Protos.TaskStatus#getTaskId() taskId} from the
 *                       {@link org.apache.mesos.v1.scheduler.Protos.Event.Update#getStatus() TaskStatus} received from Mesos.
 * @return  A {@link Call} with a configured {@link Acknowledge}.
 */
@NotNull
public static Call ackUpdate(
    @NotNull final Protos.FrameworkID frameworkId,
    @NotNull final ByteString uuid,
    @NotNull final Protos.AgentID agentId,
    @NotNull final Protos.TaskID taskId
) {
    return newBuilder()
        .setFrameworkId(frameworkId)
        .setType(Type.ACKNOWLEDGE)
        .setAcknowledge(
            Acknowledge.newBuilder()
                .setUuid(uuid)
                .setAgentId(agentId)
                .setTaskId(taskId)
                .build()
        )
        .build();
}
 

private Protos.TaskInfo assign(
    MutableStoreProvider storeProvider,
    Protos.Offer offer,
    String taskId,
    boolean revocable) {

  String host = offer.getHostname();
  IAssignedTask assigned = stateManager.assignTask(
      storeProvider,
      taskId,
      host,
      offer.getAgentId(),
      task -> mapAndAssignResources(offer, task));
  LOG.info(
      "Offer on agent {} (id {}) is being assigned task for {}.",
      host, offer.getAgentId().getValue(), taskId);
  return taskFactory.createFrom(assigned, offer, revocable);
}
 

public @NotNull Protos.CommandInfo protos(boolean isShell) {

        Protos.CommandInfo.Builder builder = Protos.CommandInfo.newBuilder();
        if(isShell) {
            builder.setValue(value);
        } else {
            builder.setShell(false);
            if (null != args && !args.isEmpty()) {
                builder.addAllArguments(args);
            }
        }

        if(null != uris && !uris.isEmpty()) {
            builder.addAllUris(uris.stream().map(uri -> Protos.CommandInfo.URI.newBuilder().setValue(uri).build()).collect(Collectors.toList()));
        }

        if(null != envs && !envs.isEmpty()) {
            Protos.Environment.Builder envBuilder = Protos.Environment.newBuilder();
            envBuilder.addAllVariables(envs.stream().map(env -> Protos.Environment.Variable.newBuilder().setName(env.getName()).setValue(env.getValue()).build()).collect(Collectors.toList()));
            builder.setEnvironment(envBuilder).build();
        }

        return builder.build();
    }
 

public @NotNull Protos.ContainerInfo protos() {

        Protos.ContainerInfo.Builder builder
                = Protos.ContainerInfo.newBuilder()
                .setDocker(docker.protos())
                .setType(getType());

        if(null != volumes && !volumes.isEmpty()) {
            volumes.forEach(
                    volume -> {
                        builder.addVolumes(Protos.Volume.newBuilder()
                                .setContainerPath(volume.getContainerPath())
                                .setHostPath(volume.getHostPath())
                                .setMode(volume.getMode())
                                .build());
                    }
            );
        }

        return builder.build();
    }
 

private static Protos.Resource.Builder resourceBuilder(
    ResourceType type,
    Optional<String> role,
    boolean revocable) {

  Protos.Resource.Builder builder = Protos.Resource.newBuilder()
      .setType(type.equals(PORTS) ? Type.RANGES : Type.SCALAR)
      .setName(type.getMesosName());

  if (revocable) {
    builder.setRevocable(Protos.Resource.RevocableInfo.getDefaultInstance());
  }

  if (role.isPresent()) {
    builder.setRole(role.get());
  }

  return builder;
}
 

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

private static void addTask(@NotNull Map<Long, String> killMap, @NotNull Protos.AgentID agentId, @NotNull Task task, @NotNull List<Protos.TaskInfo> tasks, boolean isRetryTask, Address address) {
    boolean isToKilled = false;
    try {
        isToKilled = killMap.containsKey(task.getTaskId());
        //update db set taskAgentRel
        AuxiliaryService.updateTask(killMap, task.getTaskId(), agentId.getValue(), isToKilled, address);
        if(!isToKilled){
            tasks.add(task.getTask(agentId));
            log.info("resourceAllocation --> add task : " + task.getTaskId());
        }
    } catch (Exception e) {
        if(!isToKilled) {
            taskReserve(task, isRetryTask);
        }
    }
}
 

@Test
public void testFrameworkInfoNoRevocableWithAnnouncedPrincipal() {
  Protos.FrameworkInfo info = CommandLineDriverSettingsModule.buildFrameworkInfo(
      "aurora",
      "user",
      Optional.of("auroraprincipal"),
      Amount.of(1L, Time.MINUTES),
      false, // revocable
      false, // allow gpu
      false, // partition aware
      Optional.empty());
  assertEquals("auroraprincipal", info.getPrincipal());
  assertEquals(0, info.getCapabilitiesCount());
  assertFalse(info.hasRole());
}
 

public
@NotNull
List<Protos.Resource> protos() {
    List<Protos.Resource> resources = new ArrayList<>();
    resources.add(addResource(CPUS, cpu));
    resources.add(addResource(MEMS, mem));
    return resources;
}
 

private Protos.ContainerInfo.DockerInfo.Network exchange() {
    if(null == net) {
        return Protos.ContainerInfo.DockerInfo.Network.BRIDGE;
    }
    switch (net) {
        case HOST: return Protos.ContainerInfo.DockerInfo.Network.HOST;
        case NONE: return Protos.ContainerInfo.DockerInfo.Network.NONE;
        case USER: return Protos.ContainerInfo.DockerInfo.Network.USER;
        default: return Protos.ContainerInfo.DockerInfo.Network.BRIDGE;
    }
}
 

public static Protos.Resource mesosRange(
    ResourceType type,
    Optional<String> role,
    Iterable<Integer> values) {

  return resourceBuilder(type, role, false)
      .setRanges(Protos.Value.Ranges.newBuilder().addAllRange(
          Iterables.transform(Numbers.toRanges(values), Numbers.RANGE_TRANSFORM)))
      .build();
}
 

public static Protos.Offer offer(String agentId, Protos.Resource... resources) {
  return Protos.Offer.newBuilder()
      .setId(Protos.OfferID.newBuilder().setValue("offer-id-" + agentId))
      .setFrameworkId(Protos.FrameworkID.newBuilder().setValue("framework-id"))
      .setAgentId(Protos.AgentID.newBuilder().setValue(agentId))
      .setHostname("hostname")
      .addAllResources(ImmutableSet.copyOf(resources)).build();
}
 

@Override
public SchedulerDriver create(
    Scheduler scheduler,
    Optional<Protos.Credential> credentials,
    Protos.FrameworkInfo frameworkInfo,
    String master) {

  schedulerFuture.set(scheduler);
  return this;
}
 

public TaskStatusReceived(
    Protos.TaskState state,
    Optional<TaskStatus.Source> source,
    Optional<TaskStatus.Reason> reason,
    Optional<Long> epochTimestampMicros) {

  this.state = requireNonNull(state);
  this.source = requireNonNull(source);
  this.reason = requireNonNull(reason);
  this.epochTimestampMicros = requireNonNull(epochTimestampMicros);
}
 

private static Protos.Resource makeRange(ResourceType type, Iterable<Integer> values) {
  return Protos.Resource.newBuilder()
      .setType(Protos.Value.Type.RANGES)
      .setName(type.getMesosName())
      .setRanges(Protos.Value.Ranges.newBuilder().addAllRange(
          Iterables.transform(Numbers.toRanges(values), Numbers.RANGE_TRANSFORM)))
      .build();
}
 

@Override
public boolean cancel(final Protos.OfferID offerId) {
  boolean success = removeFromHostOffers(offerId);
  if (!success) {
    // This will happen rarely when we race to process this rescind against accepting the offer
    // to launch a task.
    // If it happens frequently, we are likely processing rescinds before the offer itself.
    LOG.warn("Failed to cancel offer: {}.", offerId.getValue());
    this.offerCancelFailures.incrementAndGet();
  }
  return success;
}
 

private void runMultipleRoles(boolean revocable) {
  ResourceBag bag = bagFromResources(TASK.getTask().getResources());
  Protos.Offer offer = offer(
      // Make cpus come from two roles.
      mesosScalar(CPUS, TEST_ROLE, revocable, EXECUTOR_BAG.valueOf(CPUS)),
      mesosScalar(CPUS, ABSENT_ROLE, revocable, bag.valueOf(CPUS)),
      // Make ram come from default role
      mesosScalar(RAM_MB, ABSENT_ROLE, false, TOTAL_BAG.valueOf(RAM_MB)),
      // Make disk come from non-default role.
      mesosScalar(DISK_MB, TEST_ROLE, false, TOTAL_BAG.valueOf(DISK_MB)),
      mesosRange(PORTS, TEST_ROLE, TASK_PORTS));

  AcceptedOffer offerAllocation = AcceptedOffer.create(offer, TASK, EXECUTOR_BAG, revocable);

  Set<Resource> taskSet = ImmutableSet.<Resource>builder()
      .add(mesosScalar(CPUS, TEST_ROLE, revocable, EXECUTOR_BAG.valueOf(CPUS)))
      .add(mesosScalar(
          CPUS,
          ABSENT_ROLE,
          revocable,
          bag.subtract(EXECUTOR_BAG).valueOf(CPUS)))
      .add(mesosScalar(RAM_MB, ABSENT_ROLE, false, bag.valueOf(RAM_MB)))
      .add(mesosScalar(DISK_MB, TEST_ROLE, false, bag.valueOf(DISK_MB)))
      .add(mesosRange(PORTS, TEST_ROLE, TASK_PORTS))
      .build();
  assertEquals(taskSet, ImmutableSet.copyOf(offerAllocation.getTaskResources()));

  Set<Resource> executorSet = ImmutableSet.<Resource>builder()
      .add(mesosScalar(CPUS, ABSENT_ROLE, revocable, EXECUTOR_BAG.valueOf(CPUS)))
      .add(mesosScalar(RAM_MB, ABSENT_ROLE, false, EXECUTOR_BAG.valueOf(RAM_MB)))
      .add(mesosScalar(DISK_MB, TEST_ROLE, false, EXECUTOR_BAG.valueOf(DISK_MB)))
      .build();
  assertEquals(executorSet, ImmutableSet.copyOf(offerAllocation.getExecutorResources()));
}
 

@VisibleForTesting
IAssignedTask mapAndAssignResources(Protos.Offer offer, IAssignedTask task) {
  IAssignedTask assigned = task;
  for (ResourceType type : ResourceManager.getTaskResourceTypes(assigned)) {
    if (type.getMapper().isPresent()) {
      assigned = type.getMapper().get().mapAndAssign(offer, assigned);
    }
  }
  return assigned;
}
 

private static Protos.CommandInfo.URI toProtoURI(IMesosFetcherURI u) {
  Builder builder = Protos.CommandInfo.URI.newBuilder()
      .setValue(u.getValue())
      .setExecutable(false)
      .setExtract(u.isExtract())
      .setCache(u.isCache());

  if (u.isSetOutputFile()) {
    builder.setOutputFile(u.getOutputFile());
  }

  return builder.build();
}
 

@Test
public void testFrameworkInfoNoRevocable() {
  Protos.FrameworkInfo info = CommandLineDriverSettingsModule.buildFrameworkInfo(
      "aurora",
      "user",
      Optional.empty(),
      Amount.of(1L, Time.MINUTES),
      false, // revocable
      false, // allow gpu
      false, // partition aware
      Optional.empty());
  assertEquals("", info.getPrincipal());
  assertEquals(0, info.getCapabilitiesCount());
  assertFalse(info.hasRole());
}
 

@Test
public void testBagFromMesosResourcesUnsupportedResources() {
  Protos.Resource unsupported = Protos.Resource.newBuilder()
      .setName("unknown")
      .setType(SCALAR)
      .setScalar(Scalar.newBuilder().setValue(1.0).build()).build();
  assertEquals(
      new ResourceBag(ImmutableMap.of(CPUS, 3.0)),
      ResourceManager.bagFromMesosResources(
          ImmutableSet.of(mesosScalar(CPUS, 3.0), unsupported)));
}
 

synchronized Optional<HostOffer> get(Protos.AgentID slaveId) {
  HostOffer offer = offersBySlave.get(slaveId);
  if (offer == null || globallyBannedOffers.contains(offer.getOffer().getId())) {
    return Optional.empty();
  }

  return Optional.of(offer);
}
 

private boolean removeFromHostOffers(final Protos.OfferID offerId) {
  requireNonNull(offerId);

  // The small risk of inconsistency is acceptable here - if we have an accept/remove race
  // on an offer, the master will mark the task as LOST and it will be retried.
  return hostOffers.remove(offerId);
}