org.apache.mesos.Scheduler Java Examples

@Override
public void send(ClusterState state, SchedulerDriver mockDriver, Scheduler scheduler) {
  Protos.TaskID taskId = builder.taskId == null ?
      state.getTaskId(builder.taskName) :
      Protos.TaskID.newBuilder().setValue(builder.taskId).build();
  Protos.TaskStatus.Builder taskStatusBuilder = Protos.TaskStatus.newBuilder()
      .setTaskId(taskId)
      .setState(builder.taskState)
      .setMessage("This is a test status");

  if (builder.readinessCheckExitCode.isPresent()) {
    taskStatusBuilder
        .getCheckStatusBuilder()
        .getCommandBuilder()
        .setExitCode(builder.readinessCheckExitCode.get());
  }
  scheduler.statusUpdate(mockDriver, taskStatusBuilder.build());
}
 

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

/**
 * Initiates the replacement of a pod through a call to {@link PodResource#replacePod(String)}.
 */
public static Send replacePod(String podName) {
  return new Send() {
    @Override
    public void send(ClusterState state, SchedulerDriver mockDriver, Scheduler scheduler) {
      PodResource r = (PodResource) state.getHTTPEndpoints().stream()
          .filter(resource -> resource instanceof PodResource)
          .findAny().get();
      r.replace(podName);
    }

    @Override
    public String getDescription() {
      return String.format("Replace pod: %s", podName);
    }
  };
}
 

public static Send register() {
  return new Send() {
    @Override
    public void send(ClusterState state, SchedulerDriver mockDriver, Scheduler scheduler) {
      scheduler.registered(
          mockDriver,
          Protos.FrameworkID.newBuilder()
              .setValue("test-framework-id")
              .build(),
          Protos.MasterInfo.newBuilder()
              .setId("test-master-id")
              .setIp(1)
              .setPort(2)
              .build());
    }

    @Override
    public String getDescription() {
      return String.format("Framework registration completed");
    }
  };
}
 

/**
 * Avoid calls to the MesosSchedulerDriver constructor, which triggers errors about libmesos not
 * being present.
 */
@Override
protected MesosSchedulerDriver createInternal(
        final Scheduler scheduler,
        final FrameworkInfo frameworkInfo,
        final String masterUrl,
        final Credential credential,
        final String mesosAPIVersion) {
    createCalls++;
    if (credential != null) {
        lastCallHadCredential = true;
        lastCallHadSecret = credential.hasSecret();
    } else {
        lastCallHadCredential = false;
        lastCallHadSecret = false;
    }
    return null; // avoid requiring a NoOpSchedulerDriver
}
 

@Subscribe
public void frameworkLostSlave(final LostSlaveMessageEnvelope envelope)
{
    checkState(envelope.getRecipient().equals(context.getDriverUPID()), "Received a remote message for local delivery");

    final UPID sender = envelope.getSender();

    if (!driverIsConnected(sender)) {
        return;
    }

    final LostSlaveMessage lostSlaveMessage = envelope.getMessage();
    final SlaveID slaveId = lostSlaveMessage.getSlaveId();

    context.removeSlave(slaveId);

    eventBus.post(new SchedulerCallback() {
        @Override
        public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
        {
            return new Runnable() {
                @Override
                public void run()
                {
                    scheduler.slaveLost(schedulerDriver, slaveId);
                }
            };
        }
    });
}
 

public JesosSchedulerDriver(final Scheduler scheduler,
                            final FrameworkInfo frameworkInfo,
                            final String master,
                            final Credential credential)
                throws IOException
{
    super(scheduler, frameworkInfo, master, true, credential);
}
 

public JesosSchedulerDriver(final Scheduler scheduler,
                            final FrameworkInfo frameworkInfo,
                            final String master,
                            boolean implicitAcknowledges) throws IOException
{
    super(scheduler, frameworkInfo, master, implicitAcknowledges, null);
}
 

public JesosSchedulerDriver(final Scheduler scheduler,
                            final FrameworkInfo frameworkInfo,
                            final String master,
                            boolean implicitAcknowledges,
                            final Credential credential)
                throws IOException
{
    super(scheduler, frameworkInfo, master, implicitAcknowledges, credential);
}
 

@Subscribe
public void frameworkRegistered(final FrameworkRegisteredMessageEnvelope envelope)
{
    checkState(envelope.getRecipient().equals(context.getDriverUPID()), "Received a remote message for local delivery");

    final FrameworkRegisteredMessage frameworkRegisteredMessage = envelope.getMessage();

    if (!masterIsValid(frameworkRegisteredMessage.getMasterInfo())) {
        return;
    }

    final FrameworkID frameworkId = frameworkRegisteredMessage.getFrameworkId();

    context.connected();
    context.setFailover(false);
    context.setFrameworkId(frameworkId);

    eventBus.post(new SchedulerCallback() {
        @Override
        public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
        {
            return new Runnable() {
                @Override
                public void run()
                {
                    scheduler.registered(schedulerDriver, frameworkId, context.getMaster());
                }
            };
        }
    });
}
 

@Subscribe
public void frameworkReregistered(final FrameworkReregisteredMessageEnvelope envelope)
{
    checkState(envelope.getRecipient().equals(context.getDriverUPID()), "Received a remote message for local delivery");

    final FrameworkReregisteredMessage frameworkReregisteredMessage = envelope.getMessage();

    if (!masterIsValid(frameworkReregisteredMessage.getMasterInfo())) {
        return;
    }

    final FrameworkID frameworkId = frameworkReregisteredMessage.getFrameworkId();

    checkState(frameworkId != null, "Received null framework reregistration message!");
    checkState(frameworkId.equals(context.getFrameworkId()), "Received framework reregistration for %s but expected %s", frameworkId.getValue(), context.getFrameworkId().getValue());

    context.connected();
    context.setFailover(false);

    eventBus.post(new SchedulerCallback() {
        @Override
        public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
        {
            return new Runnable() {
                @Override
                public void run()
                {
                    scheduler.reregistered(schedulerDriver, context.getMaster());
                }
            };
        }
    });
}
 

@Subscribe
public void frameworkResourceOffers(final ResourceOffersMessageEnvelope envelope)
{
    checkState(envelope.getRecipient().equals(context.getDriverUPID()), "Received a remote message for local delivery");

    final UPID sender = envelope.getSender();

    if (!driverIsConnected(sender)) {
        return;
    }

    final ResourceOffersMessage resourceOffersMessage = envelope.getMessage();
    final List<Offer> offers = resourceOffersMessage.getOffersList();
    final List<UPID> pids = ImmutableList.copyOf(Lists.transform(resourceOffersMessage.getPidsList(), UPID.getCreateFunction()));

    checkState(offers.size() == pids.size(), "Received %s offers but only %s pids!", offers.size(), pids.size());

    int pidIndex = 0;
    for (final Offer offer : offers) {
        context.addOffer(offer.getId(), offer.getSlaveId(), pids.get(pidIndex++));
    }

    eventBus.post(new SchedulerCallback() {
        @Override
        public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
        {
            return new Runnable() {
                @Override
                public void run()
                {
                    scheduler.resourceOffers(schedulerDriver, resourceOffersMessage.getOffersList());
                }
            };
        }
    });
}
 

@Subscribe
public void frameworkRescindOffer(final RescindResourceOfferMessageEnvelope envelope)
{
    checkState(envelope.getRecipient().equals(context.getDriverUPID()), "Received a remote message for local delivery");

    final UPID sender = envelope.getSender();

    if (!driverIsConnected(sender)) {
        return;
    }

    final RescindResourceOfferMessage rescindResourceOfferMessage = envelope.getMessage();
    context.removeAllOffers(rescindResourceOfferMessage.getOfferId());

    eventBus.post(new SchedulerCallback() {
        @Override
        public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
        {
            return new Runnable() {
                @Override
                public void run()
                {
                    scheduler.offerRescinded(schedulerDriver, rescindResourceOfferMessage.getOfferId());
                }
            };
        }
    });
}
 

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

@Subscribe
public void frameworkFrameworkMessage(final ExecutorToFrameworkMessageEnvelope envelope)
{
    checkState(envelope.getRecipient().equals(context.getDriverUPID()), "Received a remote message for local delivery");

    final UPID sender = envelope.getSender();

    if (!driverIsConnected(sender)) {
        return;
    }

    final ExecutorToFrameworkMessage executorToFrameworkMessage = envelope.getMessage();

    eventBus.post(new SchedulerCallback() {
        @Override
        public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
        {
            return new Runnable() {
                @Override
                public void run()
                {
                    scheduler.frameworkMessage(schedulerDriver,
                        executorToFrameworkMessage.getExecutorId(),
                        executorToFrameworkMessage.getSlaveId(),
                        executorToFrameworkMessage.getData().toByteArray());
                }
            };
        }
    });
}
 

@Subscribe
public void frameworkError(final FrameworkErrorMessageEnvelope envelope)
{
    checkState(envelope.getRecipient().equals(context.getDriverUPID()), "Received a remote message for local delivery");

    final UPID sender = envelope.getSender();

    if (!driverIsConnected(sender)) {
        return;
    }

    final FrameworkErrorMessage frameworkErrorMessage = envelope.getMessage();

    eventBus.post(new SchedulerCallback() {
        @Override
        public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
        {
            return new Runnable() {
                @Override
                public void run()
                {
                    schedulerDriver.abort();
                    scheduler.error(schedulerDriver, frameworkErrorMessage.getMessage());
                }
            };
        }
    });
}
 

@Override
public Status start()
{
    if (!context.isStateMachine(DRIVER_NOT_STARTED)) {
        return context.getStateMachine();
    }

    try {
        driverStart();

        masterChanged(null);

        context.setStateMachine(DRIVER_RUNNING);
    }
    catch (final Exception e) {
        context.setStateMachine(DRIVER_ABORTED);
        LOG.error(e, "Failed to create scheduler process for '%s'", context.getDriverUPID());

        eventBus.post(new SchedulerCallback() {
            @Override
            public Runnable getCallback(final Scheduler scheduler, final SchedulerDriver schedulerDriver)
            {
                return new Runnable() {
                    @Override
                    public void run()
                    {
                        scheduler.error(schedulerDriver, format("Failed to create scheduler process for '%s': %s", context.getDriverUPID(), e.getMessage()));
                    }
                };
            }
        });
    }

    return context.getStateMachine();
}
 

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

public MockSchedulerDriver(@NotNull final Scheduler scheduler) {
    this.scheduler = scheduler;
    masterInfo = Protos.MasterInfo.newBuilder()
            .setHostname("127.0.0.1")
            .setId(UUID.randomUUID().toString())
            .setIp(0x7f000001)
            .setPort(42)
            .build();
}
 

@Override
protected void configure() {
  bind(Scheduler.class).to(MesosSchedulerImpl.class);
  bind(org.apache.mesos.v1.scheduler.Scheduler.class).to(VersionedMesosSchedulerImpl.class);
  bind(MesosSchedulerImpl.class).in(Singleton.class);
  bind(MesosCallbackHandler.class).to(MesosCallbackHandlerImpl.class);
  bind(MesosCallbackHandlerImpl.class).in(Singleton.class);
  // TODO(zmanji): Create singleThreadedExecutor (non-scheduled) variant.
  bind(Executor.class).annotatedWith(SchedulerExecutor.class)
      .toInstance(AsyncUtil.singleThreadLoggingScheduledExecutor("SchedulerImpl-%d", LOG));

  switch (kind) {
    case SCHEDULER_DRIVER:
      bind(Driver.class).to(SchedulerDriverService.class);
      bind(SchedulerDriverService.class).in(Singleton.class);
      break;
    case V0_DRIVER:
      bind(Driver.class).to(VersionedSchedulerDriverService.class);
      bind(VersionedSchedulerDriverService.class).in(Singleton.class);
      PubsubEventModule.bindRegisteredSubscriber(binder(), VersionedSchedulerDriverService.class);
      break;
    case V1_DRIVER:
      bind(Driver.class).to(VersionedSchedulerDriverService.class);
      bind(VersionedSchedulerDriverService.class).in(Singleton.class);
      PubsubEventModule.bindRegisteredSubscriber(binder(), VersionedSchedulerDriverService.class);
      break;
    default:
      checkState(false, "Unknown driver kind.");
      break;
  }

  PubsubEventModule.bindSubscriber(binder(), TaskStatusStats.class);
  bind(TaskStatusStats.class).in(Singleton.class);
}
 

@Inject
SchedulerDriverService(
    Scheduler scheduler,
    Storage storage,
    DriverSettings driverSettings,
    DriverFactory driverFactory,
    FrameworkInfoFactory infoFactory) {

  this.scheduler = requireNonNull(scheduler);
  this.storage = requireNonNull(storage);
  this.driverSettings = requireNonNull(driverSettings);
  this.driverFactory = requireNonNull(driverFactory);
  this.infoFactory = requireNonNull(infoFactory);
}
 

@Before
public void setUp() {
  scheduler = createMock(Scheduler.class);
  storage = new StorageTestUtil(this);
  driverFactory = createMock(DriverFactory.class);
  schedulerDriver = createMock(SchedulerDriver.class);
  infoFactory = createMock(FrameworkInfoFactory.class);
  driverService = new SchedulerDriverService(
      scheduler,
      storage.storage,
      SETTINGS,
      driverFactory,
      infoFactory);
}
 

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

  schedulerFuture.set(scheduler);
  return this;
}
 

@Override
protected SchedulerDriver setup(SimulatedCloud cloud, Protos.FrameworkInfo framework, Scheduler callbackHandler) {
    int grpcPort = UnusedSocketPortAllocator.global().allocate();
    this.injector = RemoteConnectorUtil.createSimulatedCloudGrpcServer(cloud, grpcPort);

    this.channel = ManagedChannelBuilder.forAddress("localhost", grpcPort)
            .usePlaintext(true)
            .build();
    this.factory = new SimulatedRemoteMesosSchedulerDriverFactory(channel, TitusRuntimes.internal());
    return factory.createDriver(framework, "N/A", callbackHandler);
}
 

public SimulatedLocalMesosSchedulerDriver(SimulatedCloud simulatedCloud,
                                          Protos.FrameworkInfo framework,
                                          Scheduler scheduler) {
    this.simulatedCloud = simulatedCloud;
    this.framework = framework;
    this.masterInfo = Protos.MasterInfo.newBuilder()
            .setHostname("titus.embedded")
            .setId("titusEmbedded")
            .setIp(0)
            .setPort(5050)
            .build();
    this.scheduler = scheduler;
}
 

SimulatedRemoteMesosSchedulerDriver(Protos.MasterInfo masterInfo, Channel channel, Scheduler callbackHandler, TitusRuntime titusRuntime) {
    this.masterInfo = masterInfo;
    this.asyncClient = SimulatedMesosServiceGrpc.newStub(channel);
    this.blockingClient = SimulatedMesosServiceGrpc.newBlockingStub(channel);
    this.callbackHandler = callbackHandler;
    this.titusRuntime = titusRuntime;
}
 

private void runFramework(String mesosMaster, JobAPI.Job job) {

    Scheduler scheduler = new MesosScheduler(controller, config, job);
    driver = new MesosSchedulerDriver(scheduler, controller.getFrameworkInfo(),
        mesosMaster);
    int status = driver.run() == Protos.Status.DRIVER_STOPPED ? 0 : 1;

    LOG.warning("Job already terminated!");
    //driver.stop();
    //System.exit(status);
  }
 

private static void launchAll(
        final Collection<Protos.Offer.Operation> operations,
        final Scheduler scheduler,
        final SchedulerDriver driver) {
    operations.stream()
            .filter(op -> op.getType() == Protos.Offer.Operation.Type.LAUNCH)
            .flatMap(op -> op.getLaunch().getTaskInfosList().stream())
            .collect(Collectors.toList()).stream()
            .map(info -> CassandraDaemonTask.parse(info))
            .forEach(task -> scheduler.statusUpdate(driver, TestUtils.generateStatus(task.getTaskInfo().getTaskId(),
                    Protos.TaskState.TASK_RUNNING, CassandraMode.NORMAL)));
}
 

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

/**
 * Broken out into a separate function to allow testing with custom SchedulerDrivers.
 */
protected SchedulerDriver createInternal(
    final Scheduler scheduler,
    final FrameworkInfo frameworkInfo,
    final String masterUrl,
    @Nullable final Credential credential,
    final String mesosAPIVersion)
{
  Capabilities capabilities = Capabilities.getInstance();
  // TODO(DCOS-29172): This can be removed if/when we switch to using our own Mesos Client
  // Love to work around the fact that the MesosToSchedulerDriverAdapter both depends directly on the
  // process environment *and* uses two unrelated constructors for the case of credential being null
  return credential == null ?
      new MesosToSchedulerDriverAdapter(scheduler, frameworkInfo, masterUrl, true) {
        @Override
        protected Mesos startInternal() {
          return startInternalCustom(
              this,
              capabilities,
              frameworkInfo,
              masterUrl,
              null,
              mesosAPIVersion
          );
        }
      } :
      new MesosToSchedulerDriverAdapter(scheduler, frameworkInfo, masterUrl, true, credential) {
        @Override
        protected Mesos startInternal() {
          return startInternalCustom(
              this,
              capabilities,
              frameworkInfo,
              masterUrl,
              credential,
              mesosAPIVersion
          );
        }
      };
}