org.apache.mesos.SchedulerDriver Java Examples

public void reconcileTasks(final SchedulerDriver driver) {
    if (!mesosConfiguration.isReconcilerEnabled()) {
        logger.info("Task reconciliation is turned-off");
        return;
    }
    try {
        if (reconciliationTrial++ % 2 == 0) {
            reconcileTasksKnownToUs(driver);
        } else {
            reconcileAllMesosTasks(driver);
        }
    } catch (Exception e) {
        // we don't want to throw errors lest periodically scheduled reconciliation be cancelled
        logger.error("Unexpected error (continuing): {}", e.getMessage(), e);
    }
}
 

@Override
public void frameworkMessage(final SchedulerDriver driver, final ExecutorID executorId, final SlaveID slaveId, final byte[] data) {
    if (LOGGER.isDebugEnabled()) {
        LOGGER.debug("frameworkMessage(driver : {}, executorId : {}, slaveId : {}, data : {})", driver, protoToString(executorId), protoToString(slaveId), protoToString(data));
    }

    try {
        final SlaveStatusDetails statusDetails = SlaveStatusDetails.parseFrom(data);
        switch (statusDetails.getStatusDetailsType()) {
            case HEALTH_CHECK_DETAILS:
                cassandraCluster.recordHealthCheck(executorId.getValue(), statusDetails.getHealthCheckDetails());
                break;
            case NODE_JOB_STATUS:
                cassandraCluster.onNodeJobStatus(statusDetails);
                break;
        }
    } catch (final InvalidProtocolBufferException e) {
        final String msg = "Error deserializing task status data to type: " + SlaveStatusDetails.class.getName();
        LOGGER.error(msg, e);
    }
}
 

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

/**
 * Verifies that the specified task (by name) was killed during the test, the specified number of times.
 * <p>
 * This count is (only) against the most recent id of the task at the point the {@link Expect} is invoked.
 */
public static Expect taskNameKilled(String taskName, int totalTimes) {
  if (totalTimes <= 0) {
    throw new IllegalArgumentException("To verify zero kills, use taskNameNotKilled()");
  }

  return new Expect() {
    @Override
    public void expect(ClusterState state, SchedulerDriver mockDriver) {
      ArgumentCaptor<Protos.TaskID> taskIdCaptor = ArgumentCaptor.forClass(Protos.TaskID.class);
      Mockito.verify(mockDriver, Mockito.atLeastOnce()).killTask(taskIdCaptor.capture());
      Protos.TaskID taskId = state.getTaskId(taskName);
      long matchingTaskKills = taskIdCaptor.getAllValues().stream().filter(i -> taskId.equals(i)).count();
      Assert.assertEquals(String.format("Task with name %s (id %s) was killed %d time%s",
          taskName, taskId.getValue(), matchingTaskKills, matchingTaskKills == 1 ? "" : "s"),
          totalTimes, matchingTaskKills);
    }

    @Override
    public String getDescription() {
      return String.format("Task named %s was killed %d time%s",
          taskName, totalTimes, totalTimes == 1 ? "" : "s");
    }
  };
}
 

@Override
public void resourceOffers(SchedulerDriver driver, List<Protos.Offer> offers) {
  try {
    Metrics.incrementReceivedOffers(offers.size());

    if (!apiServerStarted.get()) {
      LOGGER.info("Declining {} offer{}: Waiting for API Server to start.",
          offers.size(), offers.size() == 1 ? "" : "s");
      OfferProcessor.declineShort(offers);
      return;
    }

    // Filter any bad resources from the offers before they even enter processing.
    offerProcessor.enqueue(offers.stream()
        .map(this::filterBadResources)
        .collect(Collectors.toList()));
  } catch (Throwable e) {
    logExceptionAndExit(e);
  }
}
 

@Before
public void before() {
    // Setup mocks
    driver = mock(SchedulerDriver.class);
    when(driver.killTask(any())).thenReturn(Protos.Status.DRIVER_RUNNING);

    config = mock(Configuration.class);
    when(config.getElasticsearchNodes()).thenReturn(2);

    state = mock(ClusterState.class);
    List<Protos.TaskInfo> list = Arrays.asList(ProtoTestUtil.getDefaultTaskInfo(), ProtoTestUtil.getDefaultTaskInfo(), ProtoTestUtil.getDefaultTaskInfo());
    when(state.getTaskList()).thenReturn(list);
}
 

@Before
public void before() throws Exception {

  Config config = Mockito.mock(Config.class);
  Mockito.when(config.getStringValue(Key.TOPOLOGY_NAME)).thenReturn(TOPOLOGY_NAME);
  Mockito.when(config.getStringValue(Key.ROLE)).thenReturn(ROLE);
  Mockito.when(config.getStringValue(Key.CORE_PACKAGE_URI)).thenReturn(CORE_PACKAGE_URI);

  Config runtime = Mockito.mock(Config.class);
  Mockito.when(runtime.getLongValue(Key.NUM_CONTAINERS)).thenReturn(NUM_CONTAINER);
  Properties properties = new Properties();
  properties.put(Key.TOPOLOGY_PACKAGE_URI.value(), TOPOLOGY_PACKAGE_URI);
  Mockito.when(runtime.get(Key.SCHEDULER_PROPERTIES)).thenReturn(properties);

  mesosFramework = Mockito.mock(MesosFramework.class);
  SchedulerDriver driver = Mockito.mock(SchedulerDriver.class);
  baseContainer = Mockito.mock(BaseContainer.class);

  scheduler = Mockito.spy(MesosScheduler.class);
  Mockito.doReturn(mesosFramework).when(scheduler).getMesosFramework();
  Mockito.doReturn(driver).when(scheduler)
      .getSchedulerDriver(Mockito.anyString(), Mockito.eq(mesosFramework));
  Mockito.doNothing().when(scheduler)
      .startSchedulerDriver();

  scheduler.initialize(config, runtime);
}
 

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

@Override
public void frameworkMessage(SchedulerDriver driver, final Protos.ExecutorID executorId, final Protos.SlaveID slaveId, final byte[] data) {
	runAsync(new Runnable() {
		@Override
		public void run() {
			MesosResourceManager.this.frameworkMessage(new FrameworkMessage(executorId, slaveId, data));
		}
	});
}
 

/**
 * Publishes ErrorEvent
 */
@Override
public void error(final SchedulerDriver driver, final String message) {
  disruptorManager.getErrorEventDisruptor().publishEvent(new EventTranslator<ErrorEvent>() {
    @Override
    public void translateTo(ErrorEvent event, long sequence) {
      event.setDriver(driver);
      event.setMessage(message);
    }
  });
}
 

@Inject
REEFScheduler(final REEFEventHandlers reefEventHandlers,
              final MesosRemoteManager mesosRemoteManager,
              final REEFExecutors executors,
              final REEFFileNames fileNames,
              final EStage<SchedulerDriver> schedulerDriverEStage,
              final ClasspathProvider classpath,
              @Parameter(JobIdentifier.class) final String jobIdentifier,
              @Parameter(MesosMasterIp.class) final String masterIp,
              @Parameter(MesosSlavePort.class) final int slavePort,
              @Parameter(JobSubmissionDirectoryPrefix.class) final String jobSubmissionDirectoryPrefix) {
  this.mesosRemoteManager = mesosRemoteManager;
  this.reefEventHandlers = reefEventHandlers;
  this.executors = executors;
  this.fileNames = fileNames;
  this.jobSubmissionDirectoryPrefix = jobSubmissionDirectoryPrefix;
  this.reefTarUri = getReefTarUri(jobIdentifier);
  this.classpath = classpath;
  this.schedulerDriverEStage = schedulerDriverEStage;

  final Protos.FrameworkInfo frameworkInfo = Protos.FrameworkInfo.newBuilder()
      .setUser("")
      .setName(REEF_JOB_NAME_PREFIX + jobIdentifier)
      .build();
  this.mesosMaster = new MesosSchedulerDriver(this, frameworkInfo, masterIp);
  this.mesosSlavePort = slavePort;
}
 

@Override
public void offerRescinded(SchedulerDriver driver, Protos.OfferID offerId) {
  try {
    LOGGER.info("Rescinding offer: {}", offerId.getValue());
    offerProcessor.dequeue(offerId);
  } catch (Throwable e) {
    logExceptionAndExit(e);
  }
}
 

/**
 * Publishes a ReRegisteredEvent
 */
@Override
public void reregistered(final SchedulerDriver driver, final Protos.MasterInfo masterInfo) {
  disruptorManager.getReRegisteredEventDisruptor().publishEvent(new EventTranslator<ReRegisteredEvent>() {
    @Override
    public void translateTo(ReRegisteredEvent event, long sequence) {
      event.setDriver(driver);
      event.setMasterInfo(masterInfo);
    }
  });
}
 

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

@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 void statusUpdate(SchedulerDriver driver, Protos.TaskStatus status) {
    LOGGER.info("Status update:" +
            " " + status.getSlaveId() +
            " " + status.getExecutorId() +
            " " + status.getHealthy() +
            " " + status.getState() +
            " " + status.getReason() +
            " " +
            " ");
    frameworkState.announceStatusUpdate(status);
}
 

@Override
public void error(SchedulerDriver driver, final String message) {
	runAsync(new Runnable() {
		@Override
		public void run() {
			onFatalError(new ResourceManagerException(message));
		}
	});
}
 

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

@Override
public void slaveLost(SchedulerDriver driver, Protos.SlaveID agentId) {
  try {
    LOGGER.warn("Agent lost: {}", agentId.getValue());
  } catch (Throwable e) {
    logExceptionAndExit(e);
  }
}
 

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

/**
 * Triggers offer processing against an empty list of offers. This is useful for bumping the plan status. In
 * practice this would happen automatically every few seconds, but we need to do it manually in these tests.
 * See OfferProcessor.
 */
public static Send emptyOffers() {
  return new Send() {
    @Override
    public void send(ClusterState state, SchedulerDriver mockDriver, Scheduler scheduler) {
      scheduler.resourceOffers(mockDriver, Collections.emptyList());
    }

    @Override
    public String getDescription() {
      return String.format("Nudge offer processing with empty list");
    }
  };
}
 

@Override
public void reregistered(SchedulerDriver driver, MasterInfo masterInfo) {
    logMesosCallbackWarn("Mesos re-registered: %s, masterInfo=%s", driver, masterInfo.getId());
    initializeNewDriver(driver);
    numMesosRegistered.increment();
    connected = true;
}
 

private static void initRestfulServer() {
    regCenter = mock(CoordinatorRegistryCenter.class);
    jobEventRdbSearch = mock(JobEventRdbSearch.class);
    SchedulerDriver schedulerDriver = mock(SchedulerDriver.class);
    ProducerManager producerManager = new ProducerManager(schedulerDriver, regCenter);
    producerManager.startup();
    restfulService = new RestfulService(regCenter, new RestfulServerConfiguration(19000), producerManager, new ReconcileService(schedulerDriver, new FacadeService(regCenter)));
    restfulService.start();
}
 

@Override
public void disconnected(SchedulerDriver driver) {
	runAsyncWithoutFencing(new Runnable() {
		@Override
		public void run() {
			MesosResourceManager.this.disconnected(new Disconnected());
		}
	});
}
 

@Override
public void executorLost(SchedulerDriver driver, Protos.ExecutorID executorId, Protos.SlaveID slaveId, int status) {
    // This is never called by Mesos, so we have to call it ourselves via a healthcheck
    // https://issues.apache.org/jira/browse/MESOS-313
    LOGGER.info("Executor lost: " + executorId.getValue() +
            " on slave " + slaveId.getValue() +
            " with status " + status);
    try {
        Protos.TaskInfo taskInfo = clusterState.getTask(executorId);
        statusUpdate(driver, Protos.TaskStatus.newBuilder().setExecutorId(executorId).setSlaveId(slaveId).setTaskId(taskInfo.getTaskId()).setState(Protos.TaskState.TASK_LOST).build());
        driver.killTask(taskInfo.getTaskId()); // It may not actually be lost, it may just have hanged. So Kill, just in case.
    } catch (IllegalArgumentException e) {
        LOGGER.warn("Unable to find TaskInfo with the given Executor ID", e);
    }
}
 

@Override
public void error(SchedulerDriver schedulerDriver, String message) {
  // TODO(mfu): TO handle this failure
  LOG.info("Received error: " + message);
}
 

private static void updateDriverAndDomain(SchedulerDriver driver, Protos.MasterInfo masterInfo) {
  Driver.setDriver(driver);
  if (masterInfo.hasDomain()) {
    IsLocalRegionRule.setLocalDomain(masterInfo.getDomain());
  }
}
 

@Override
protected ActorRef createLaunchCoordinator(SchedulerDriver schedulerDriver, ActorRef selfActorRef) {
	return launchCoordinator.ref();
}
 

@Override
public void disconnected(final SchedulerDriver driver) {
  this.onRuntimeError(new RuntimeException("Scheduler disconnected from MesosMaster"));
}
 

@Override
public void reregistered(final SchedulerDriver schedulerDriver, final Protos.MasterInfo masterInfo) {
    log.info("call reregistered");
    taskScheduler.expireAllLeases();
    MesosStateService.register(masterInfo.getHostname(), masterInfo.getPort());
}