org.apache.mesos.v1.scheduler.Protos.Call Java Examples

@NotNull
private static Call invokeTasks(
        @NotNull final FrameworkID frameworkId,
        @NotNull final List<OfferID> offerIds,
        @NotNull final List<TaskInfo> tasks
) {
    return Call.newBuilder()
            .setFrameworkId(frameworkId)
            .setType(Call.Type.ACCEPT)
            .setAccept(
                    Call.Accept.newBuilder()
                            .addAllOfferIds(offerIds)
                            .addOperations(
                                    Offer.Operation.newBuilder()
                                            .setType(Offer.Operation.Type.LAUNCH)
                                            .setLaunch(
                                                    Offer.Operation.Launch.newBuilder()
                                                            .addAllTaskInfos(tasks)
                                            )
                            )
            )
            .build();
}
 

@Test
public void testAcceptInverseOffer() {
  expectStart();

  expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));

  Capture<Call> acceptCapture = createCapture();
  mesos.send(capture(acceptCapture));
  expectLastCall().once();

  control.replay();
  driverService.startAsync().awaitRunning();
  driverService.registered(new PubsubEvent.DriverRegistered());

  driverService.acceptInverseOffer(OFFER_ID, FILTER);

  assertTrue(acceptCapture.hasCaptured());
  Call accept = acceptCapture.getValue();
  assertEquals(accept.getFrameworkId().getValue(), FRAMEWORK_ID);
  assertEquals(accept.getType(), Call.Type.ACCEPT_INVERSE_OFFERS);
  assertEquals(accept.getAcceptInverseOffers().getFilters(), FILTER);
  assertEquals(
      accept.getAcceptInverseOffers().getInverseOfferIdsList(),
      ImmutableList.of(OFFER_ID)
  );
}
 

@Test
public void testAccept() {
  expectStart();
  expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));

  Capture<Call> acceptCapture = createCapture();
  mesos.send(capture(acceptCapture));
  expectLastCall().once();

  control.replay();
  driverService.startAsync().awaitRunning();
  driverService.registered(new PubsubEvent.DriverRegistered());

  driverService.acceptOffers(OFFER_ID, ImmutableList.of(), FILTER);

  assertTrue(acceptCapture.hasCaptured());
  Call accept = acceptCapture.getValue();
  assertEquals(accept.getFrameworkId().getValue(), FRAMEWORK_ID);
  assertEquals(accept.getType(), Call.Type.ACCEPT);
  assertEquals(accept.getAccept().getFilters(), FILTER);
  assertEquals(accept.getAccept().getOfferIdsList(), ImmutableList.of(OFFER_ID));
  assertEquals(accept.getAccept().getOperationsList(), ImmutableList.of());
}
 

@Test
public void testDecline() {
  expectStart();
  expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));

  Capture<Call> declineCapture = createCapture();
  mesos.send(capture(declineCapture));
  expectLastCall().once();

  control.replay();
  driverService.startAsync().awaitRunning();
  driverService.registered(new PubsubEvent.DriverRegistered());

  driverService.declineOffer(OFFER_ID, FILTER);

  assertTrue(declineCapture.hasCaptured());
  Call decline = declineCapture.getValue();
  assertEquals(decline.getFrameworkId().getValue(), FRAMEWORK_ID);
  assertEquals(decline.getType(), Call.Type.DECLINE);
  assertEquals(decline.getDecline().getOfferIdsList(), ImmutableList.of(OFFER_ID));
  assertEquals(decline.getDecline().getFilters(), FILTER);
}
 

@Test
public void testKill() {
  expectStart();
  expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));

  Capture<Call> killCapture = createCapture();
  mesos.send(capture(killCapture));
  expectLastCall().once();

  control.replay();
  driverService.startAsync().awaitRunning();
  driverService.registered(new PubsubEvent.DriverRegistered());

  driverService.killTask("task-id");

  assertTrue(killCapture.hasCaptured());
  Call kill = killCapture.getValue();
  assertEquals(kill.getFrameworkId().getValue(), FRAMEWORK_ID);
  assertEquals(kill.getType(), Call.Type.KILL);
  assertEquals(kill.getKill().getTaskId().getValue(), "task-id");
}
 

@Override
public void reconcileTasks(Collection<TaskStatus> statuses) {
  whenRegistered(() -> {
    Collection<Call.Reconcile.Task> tasks = Collections2.transform(statuses, taskStatus ->
        Call.Reconcile.Task.newBuilder()
            .setTaskId(taskStatus.getTaskId())
            .build());

    Futures.getUnchecked(mesosFuture).send(
        Call.newBuilder()
            .setType(Call.Type.RECONCILE)
            .setFrameworkId(getFrameworkId())
            .setReconcile(
                Call.Reconcile.newBuilder()
                    .addAllTasks(tasks))
            .build()
    );
  });
}
 

@Override
public void killTask(String taskId) {
  whenRegistered(() -> {
    LOG.info("Killing task {}", taskId);

    Futures.getUnchecked(mesosFuture).send(
        Call.newBuilder().setType(Call.Type.KILL)
            .setFrameworkId(getFrameworkId())
            .setKill(
                Call.Kill.newBuilder()
                    .setTaskId(TaskID.newBuilder().setValue(taskId)))
            .build()
    );

  });
}
 

@Override
public void declineOffer(OfferID offerId, Filters filter) {
  whenRegistered(() -> {
    LOG.info("Declining offer {}", offerId.getValue());

    Futures.getUnchecked(mesosFuture).send(
        Call.newBuilder().setType(Call.Type.DECLINE)
            .setFrameworkId(getFrameworkId())
            .setDecline(
                Call.Decline.newBuilder()
                    .setFilters(filter)
                    .addOfferIds(offerId))
            .build()
    );
  });
}
 

@Override
public void acceptInverseOffer(OfferID offerID, Filters filter) {
  whenRegistered(() -> {
    LOG.info("Accepting Inverse Offer {}", offerID.getValue());

    Futures.getUnchecked(mesosFuture).send(
        Call.newBuilder()
            .setFrameworkId(getFrameworkId())
            .setType(Call.Type.ACCEPT_INVERSE_OFFERS)
            .setAcceptInverseOffers(
                Call.AcceptInverseOffers.newBuilder()
                  .addInverseOfferIds(offerID)
                  .setFilters(filter)
            )
        .build());
  });
}
 

@Override
public void acceptOffers(OfferID offerId, Collection<Operation> operations, Filters filter) {
  whenRegistered(() -> {
    Collection<Operation.Type> opTypes = Collections2.transform(operations, Operation::getType);
    LOG.info("Accepting offer {} with ops {}", offerId.getValue(), opTypes);

    Futures.getUnchecked(mesosFuture).send(
        Call.newBuilder()
            .setFrameworkId(getFrameworkId())
            .setType(Call.Type.ACCEPT)
            .setAccept(
                Call.Accept.newBuilder()
                    .addOfferIds(offerId)
                    .addAllOperations(operations)
                    .setFilters(filter))
            .build());
  });

}
 

@NotNull
private static Call sleep(
    @NotNull final FrameworkID frameworkId,
    @NotNull final List<OfferID> offerIds,
    @NotNull final List<TaskInfo> tasks
) {
    return Call.newBuilder()
        .setFrameworkId(frameworkId)
        .setType(Call.Type.ACCEPT)
        .setAccept(
            Call.Accept.newBuilder()
                .addAllOfferIds(offerIds)
                .addOperations(
                    Offer.Operation.newBuilder()
                        .setType(Offer.Operation.Type.LAUNCH)
                        .setLaunch(
                            Offer.Operation.Launch.newBuilder()
                                .addAllTaskInfos(tasks)
                        )
                )
        )
        .build();
}
 

/**
 * 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);
}
 

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

public void sendCall(Call call) {
  if (publisher == null) {
    throw new RuntimeException(
      "No publisher found, please call subscribe before sending anything."
    );
  }
  publisher.onNext(Optional.of(SinkOperations.create(call)));
}
 

/**
 * 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 Call Call} messages.
 * @param frameworkInfo The frameworkInfo to set on the {@link Subscribe Subscribe} sub-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 Protos.FrameworkInfo frameworkInfo
) {
    return newBuilder()
        .setFrameworkId(frameworkId)
        .setType(Type.SUBSCRIBE)
        .setSubscribe(
            Subscribe.newBuilder()
                .setFrameworkInfo(frameworkInfo)
        )
        .build();
}
 

@Override
public void acknowledgeStatusUpdate(TaskStatus status) {
  // The Mesos API says frameworks are only supposed to acknowledge status updates
  // with a UUID. The V0Driver accepts them just fine but the V1Driver logs every time
  // a status update is given without a uuid. To silence logs, we drop them here.

  whenRegistered(() -> {
    if (!status.hasUuid()) {
      return;
    }

    LOG.info("Acking status update for {} with uuid: {}",
        status.getTaskId().getValue(),
        status.getUuid());

    Futures.getUnchecked(mesosFuture).send(
        Call.newBuilder().setType(Call.Type.ACKNOWLEDGE)
            .setFrameworkId(getFrameworkId())
            .setAcknowledge(
                Call.Acknowledge.newBuilder()
                    .setAgentId(status.getAgentId())
                    .setTaskId(status.getTaskId())
                    .setUuid(status.getUuid()))
            .build()
    );
  });

}
 

@Test(timeout = 300000)
public void testConnected() throws Exception {
  // Once the V1 driver has connected, we need to establish a subscription to get events
  expectFrameworkInfoRead();

  Capture<Call> subscribeCapture = createCapture();

  driver.send(capture(subscribeCapture));
  expectLastCall().once();

  Capture<ExceptionalSupplier<Boolean, RuntimeException>> supplierCapture = createCapture();
  backoffHelper.doUntilSuccess(capture(supplierCapture));
  expectLastCall().once();

  control.replay();

  scheduler.connected(driver);

  waitUntilCaptured(supplierCapture);

  assertTrue(supplierCapture.hasCaptured());
  ExceptionalSupplier<Boolean, RuntimeException> supplier = supplierCapture.getValue();

  // Make one connection attempt
  supplier.get();

  assertTrue(subscribeCapture.hasCaptured());

  Call subscribe = subscribeCapture.getValue();

  assertEquals(subscribe.getType(), Call.Type.SUBSCRIBE);
  assertEquals(subscribe.getFrameworkId(), FRAMEWORK);
  assertEquals(
      subscribe.getSubscribe().getFrameworkInfo(),
      FRAMEWORK_INFO.toBuilder().setId(FRAMEWORK).build());
}
 

/**
 * Utility method to more succinctly construct a {@link Call Call} of type {@link Type#DECLINE DECLINE}.
 * <p>
 *
 * @param frameworkId    The {@link Protos.FrameworkID FrameworkID} to be set on the {@link Call}
 * @param offerIds       A list of {@link Protos.OfferID OfferID} from the
 *                       {@link Protos.Offer}s received from Mesos.
 * @return  A {@link Call} with a configured {@link Decline}.
 */
@NotNull
public static Call decline(@NotNull final Protos.FrameworkID frameworkId, @NotNull final List<Protos.OfferID> offerIds) {
    return newBuilder()
        .setFrameworkId(frameworkId)
        .setType(Type.DECLINE)
        .setDecline(
            Decline.newBuilder()
                .addAllOfferIds(offerIds)
        )
        .build();
}
 

@NotNull
private static Call shutdownTasks(
        @NotNull final FrameworkID frameworkId,
        @NotNull final AgentID agentID,
        @NotNull final ExecutorID executorID
) {
    return Call.newBuilder()
            .setFrameworkId(frameworkId)
            .setType(Call.Type.SHUTDOWN).setShutdown(
                    Call.Shutdown.newBuilder().setAgentId(agentID).setExecutorId(executorID).build()
            )
            .build();
}
 

private void sendCall(Call.Builder b, Type t) {
  Call call = b.setType(t).setFrameworkId(frameworkId).build();
  sendCall(call);
}
 

private static Builder build() {
  return Call.newBuilder();
}
 

@NotNull
private static SinkOperation<Call> handleOffer(@NotNull final Tuple2<Offer, State<FrameworkID, TaskID, TaskState>> t) {
    final Offer offer = t._1;
    final State<FrameworkID, TaskID, TaskState> state = t._2;
    final int offerCount = state.getOfferCounter().incrementAndGet();
    final String desiredRole = state.getResourceRole();

    final FrameworkID frameworkId = state.getFwId();
    final AgentID agentId = offer.getAgentId();
    final List<OfferID> ids = newArrayList(offer.getId());

    final Map<String, List<Resource>> resources = offer.getResourcesList()
        .stream()
        .collect(groupingBy(Resource::getName));
    final List<Resource> cpuList = resources.get("cpus");
    final List<Resource> memList = resources.get("mem");
    if (
        cpuList != null && !cpuList.isEmpty()
            && memList != null && !memList.isEmpty()
            && cpuList.size() == memList.size()
        ) {
        final List<TaskInfo> tasks = newArrayList();
        for (int i = 0; i < cpuList.size(); i++) {
            final Resource cpus = cpuList.get(i);
            final Resource mem = memList.get(i);

            if (desiredRole.equals(cpus.getRole()) && desiredRole.equals(mem.getRole())) {
                double availableCpu = cpus.getScalar().getValue();
                double availableMem = mem.getScalar().getValue();
                final double cpusPerTask = state.getCpusPerTask();
                final double memMbPerTask = state.getMemMbPerTask();
                while (availableCpu >= cpusPerTask && availableMem >= memMbPerTask) {
                    availableCpu -= cpusPerTask;
                    availableMem -= memMbPerTask;
                    final String taskId = String.format("task-%d-%d", offerCount, state.getTotalTaskCounter().incrementAndGet());
                    tasks.add(sleepTask(agentId, taskId, cpus.getRole(), cpusPerTask, mem.getRole(), memMbPerTask));
                }
            }
        }

        if (!tasks.isEmpty()) {
            LOGGER.info("Launching {} tasks", tasks.size());
            return sink(
                sleep(frameworkId, ids, tasks),
                () -> tasks.forEach(task -> state.put(task.getTaskId(), TaskState.TASK_STAGING)),
                (e) -> LOGGER.warn("", e)
            );
        } else {
            return sink(decline(frameworkId, ids));
        }
    } else {
        return sink(decline(frameworkId, ids));
    }
}
 

@Override
public void connected(Mesos mesos) {
  LOG.info("Connected to Mesos master.");
  isConnected.set(true);

  Optional<String> frameworkId = storage.read(
      storeProvider -> storeProvider.getSchedulerStore().fetchFrameworkId());

  Protos.FrameworkInfo.Builder frameworkBuilder = infoFactory.getFrameworkInfo().toBuilder();

  Call.Builder call = Call.newBuilder().setType(Call.Type.SUBSCRIBE);

  if (frameworkId.isPresent()) {
    LOG.info("Found persisted framework ID: " + frameworkId);
    Protos.FrameworkID id = Protos.FrameworkID.newBuilder().setValue(frameworkId.get()).build();
    frameworkBuilder.setId(id);
    call.setFrameworkId(id);
  } else {
    frameworkBuilder.clearId();
    call.clearFrameworkId();
    LOG.warn("Did not find a persisted framework ID, connecting as a new framework.");
  }

  call.setSubscribe(Call.Subscribe.newBuilder().setFrameworkInfo(frameworkBuilder));

  executor.execute(() -> {
    LOG.info("Starting to subscribe to Mesos with backoff.");
    try {
      backoffHelper.doUntilSuccess(() -> {
        if (!isConnected.get())  {
          LOG.info("Disconnected while attempting to subscribe. Stopping attempt.");
          return true;
        }
        if (!isSubscribed.get()) {
          LOG.info("Sending subscribe call.");
          mesos.send(call.build());
          subcriptionCalls.incrementAndGet();
          return false;
        }
        LOG.info("Subscribed to Mesos");
        return true;
      });
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
  });
}
 

@NotNull
private static SinkOperation<Call> handleOffer(final Offer offer, final StateManager state) {
    final FrameworkID frameworkId = state.getFwId();
    final AgentID agentId = offer.getAgentId();
    final List<OfferID> ids = newArrayList(offer.getId());


    final Map<String, List<Resource>> resources = offer.getResourcesList()
            .stream()
            .collect(groupingBy(Resource::getName));
    final List<Resource> cpuList = resources.get("cpus");
    final List<Resource> memList = resources.get("mem");
    final List<DesiredTask> desired = state.currentDesiredTasks();

    if (null != desired &&
            !desired.isEmpty() &&
            cpuList != null && !cpuList.isEmpty()
            && memList != null && !memList.isEmpty()
            && cpuList.size() == memList.size()) {

        final List<TaskInfo> tasks = newArrayList();


        for (int i = 0; i < cpuList.size(); i++) {


            final Resource cpus = cpuList.get(i);
            final Resource mem = memList.get(i);
            double availableCpu = cpus.getScalar().getValue();
            double availableMem = mem.getScalar().getValue();
            for (DesiredTask dt : desired) {
                final String desiredRole = dt.role;
                if (!launched.contains(dt.id) && desiredRole.equals(cpus.getRole()) && desiredRole.equals(mem.getRole())) {

                    final double cpusPerTask = dt.cpu;
                    final double memMbPerTask = dt.memory;
                    if (availableCpu >= cpusPerTask && availableMem >= memMbPerTask) {
                        availableCpu -= cpusPerTask;
                        availableMem -= memMbPerTask;
                        launched.add(dt.id);
                        tasks.add(dt.createTaskInfo(agentId));
                    }
                }
            }
        }

        if (!tasks.isEmpty()) {
            LOGGER.info("Launching {} tasks", tasks.size());
            return sink(
                    invokeTasks(frameworkId, ids, tasks),
                    () -> tasks.forEach(task ->
                            state.update(UUID.fromString(task.getTaskId().getValue()),
                                    null,
                                    ContainerSubstrate.TaskState.Starting)),
                    (e) -> LOGGER.warn("", e)
            );
        } else {
            return sink(decline(frameworkId, ids));
        }
    } else {
        return sink(decline(frameworkId, ids));
    }
}
 

/**
 * 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 String frameworkId,
    @NotNull final String user,
    @NotNull final String frameworkName,
    final long failoverTimeoutSeconds
) {
    final Protos.FrameworkID frameworkID = Protos.FrameworkID.newBuilder().setValue(frameworkId).build();
    return subscribe(frameworkID, user, frameworkName, failoverTimeoutSeconds);
}