org.apache.flink.mesos.runtime.clusterframework.store.MesosWorkerStore Java Examples

/**
 * Fetches framework/worker information persisted by a prior incarnation of the RM.
 */
private CompletableFuture<List<MesosWorkerStore.Worker>> getWorkersAsync() {
	// if this resource manager is recovering from failure,
	// then some worker tasks are most likely still alive and we can re-obtain them
	return CompletableFuture.supplyAsync(() -> {
		try {
			final List<MesosWorkerStore.Worker> tasksFromPreviousAttempts = workerStore.recoverWorkers();
			for (final MesosWorkerStore.Worker worker : tasksFromPreviousAttempts) {
				if (worker.state() == MesosWorkerStore.WorkerState.New) {
					// remove new workers because allocation requests are transient
					workerStore.removeWorker(worker.taskID());
				}
			}
			return tasksFromPreviousAttempts;
		} catch (final Exception e) {
			throw new CompletionException(new ResourceManagerException(e));
		}
	}, getRpcService().getExecutor());
}
 

/**
 * Fetches framework/worker information persisted by a prior incarnation of the RM.
 */
private CompletableFuture<List<MesosWorkerStore.Worker>> getWorkersAsync() {
	// if this resource manager is recovering from failure,
	// then some worker tasks are most likely still alive and we can re-obtain them
	return CompletableFuture.supplyAsync(() -> {
		try {
			final List<MesosWorkerStore.Worker> tasksFromPreviousAttempts = workerStore.recoverWorkers();
			for (final MesosWorkerStore.Worker worker : tasksFromPreviousAttempts) {
				if (worker.state() == MesosWorkerStore.WorkerState.New) {
					// remove new workers because allocation requests are transient
					workerStore.removeWorker(worker.taskID());
				}
			}
			return tasksFromPreviousAttempts;
		} catch (final Exception e) {
			throw new CompletionException(new ResourceManagerException(e));
		}
	}, getRpcService().getExecutor());
}
 

@Test
public void testClearStateAfterRevokeLeadership() throws Exception {
	new Context() {{
		final MesosWorkerStore.Worker worker1 = MesosWorkerStore.Worker.newWorker(task1);
		final MesosWorkerStore.Worker worker2 = MesosWorkerStore.Worker.newWorker(task2).launchWorker(slave1, slave1host);
		final MesosWorkerStore.Worker worker3 = MesosWorkerStore.Worker.newWorker(task3).launchWorker(slave1, slave1host).releaseWorker();
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(Arrays.asList(worker1, worker2, worker3)).thenReturn(Collections.emptyList());

		startResourceManager();
		rmServices.rmLeaderElectionService.notLeader();
		rmServices.grantLeadership();

		assertThat(resourceManager.workersInNew.size(), equalTo(0));
		assertThat(resourceManager.workersInLaunch.size(), equalTo(0));
		assertThat(resourceManager.workersBeingReturned.size(), equalTo(0));
		verify(rmServices.schedulerDriver).stop(true);
	}};
}
 

/**
 * Test unplanned task failure of a pending worker.
 */
@Test
public void testWorkerFailed() throws Exception {
	new Context() {{
		// set the initial persistent state with a launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		when(rmServices.workerStore.newTaskID()).thenReturn(task2);
		startResourceManager();

		// tell the RM that a task failed
		when(rmServices.workerStore.removeWorker(task1)).thenReturn(true);
		resourceManager.taskTerminated(new TaskMonitor.TaskTerminated(task1, Protos.TaskStatus.newBuilder()
			.setTaskId(task1).setSlaveId(slave1).setState(Protos.TaskState.TASK_FAILED).build()));

		// verify that the instance state was updated
		verify(rmServices.workerStore).removeWorker(task1);
		assertThat(resourceManager.workersInLaunch.entrySet(), empty());
		assertThat(resourceManager.workersBeingReturned.entrySet(), empty());
		assertThat(resourceManager.workersInNew, hasKey(extractResourceID(task2)));

		// verify that `closeTaskManagerConnection` was called
		assertThat(resourceManager.closedTaskManagerConnections, hasItem(extractResourceID(task1)));
	}};
}
 

/**
 * Fetches framework/worker information persisted by a prior incarnation of the RM.
 */
private CompletableFuture<List<MesosWorkerStore.Worker>> getWorkersAsync() {
	// if this resource manager is recovering from failure,
	// then some worker tasks are most likely still alive and we can re-obtain them
	return CompletableFuture.supplyAsync(() -> {
		try {
			final List<MesosWorkerStore.Worker> tasksFromPreviousAttempts = workerStore.recoverWorkers();
			for (final MesosWorkerStore.Worker worker : tasksFromPreviousAttempts) {
				if (worker.state() == MesosWorkerStore.WorkerState.New) {
					// remove new workers because allocation requests are transient
					workerStore.removeWorker(worker.taskID());
				}
			}
			return tasksFromPreviousAttempts;
		} catch (final Exception e) {
			throw new CompletionException(new ResourceManagerException(e));
		}
	}, getRpcService().getExecutor());
}
 

@Test
public void testClearStateAfterRevokeLeadership() throws Exception {
	new Context() {{
		final MesosWorkerStore.Worker worker1 = MesosWorkerStore.Worker.newWorker(task1);
		final MesosWorkerStore.Worker worker2 = MesosWorkerStore.Worker.newWorker(task2).launchWorker(slave1, slave1host);
		final MesosWorkerStore.Worker worker3 = MesosWorkerStore.Worker.newWorker(task3).launchWorker(slave1, slave1host).releaseWorker();
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(Arrays.asList(worker1, worker2, worker3)).thenReturn(Collections.emptyList());

		startResourceManager();
		rmServices.rmLeaderElectionService.notLeader();
		rmServices.grantLeadership();

		assertThat(resourceManager.workersInNew.size(), equalTo(0));
		assertThat(resourceManager.workersInLaunch.size(), equalTo(0));
		assertThat(resourceManager.workersBeingReturned.size(), equalTo(0));
		verify(rmServices.schedulerDriver).stop(true);
	}};
}
 

/**
 * Test unplanned task failure of a pending worker.
 */
@Test
public void testWorkerFailed() throws Exception {
	new Context() {{
		// set the initial persistent state with a launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		when(rmServices.workerStore.newTaskID()).thenReturn(task2);
		startResourceManager();

		// tell the RM that a task failed
		when(rmServices.workerStore.removeWorker(task1)).thenReturn(true);
		resourceManager.taskTerminated(new TaskMonitor.TaskTerminated(task1, Protos.TaskStatus.newBuilder()
			.setTaskId(task1).setSlaveId(slave1).setState(Protos.TaskState.TASK_FAILED).build()));

		// verify that the instance state was updated
		verify(rmServices.workerStore).removeWorker(task1);
		assertThat(resourceManager.workersInLaunch.entrySet(), empty());
		assertThat(resourceManager.workersBeingReturned.entrySet(), empty());
		assertThat(resourceManager.workersInNew, hasKey(extractResourceID(task2)));

		// verify that `closeTaskManagerConnection` was called
		assertThat(resourceManager.closedTaskManagerConnections, hasItem(extractResourceID(task1)));
	}};
}
 

/**
 * Allocate a worker using the RM.
 */
public MesosWorkerStore.Worker allocateWorker(Protos.TaskID taskID, ResourceProfile resourceProfile) throws Exception {
	when(rmServices.workerStore.newTaskID()).thenReturn(taskID);
	rmServices.slotManagerStarted.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

	CompletableFuture<Void> allocateResourceFuture = resourceManager.callAsync(
		() -> {
			rmServices.rmActions.allocateResource(resourceProfile);
			return null;
		},
		timeout);
	MesosWorkerStore.Worker expected = MesosWorkerStore.Worker.newWorker(taskID, resourceProfile);

	// check for exceptions
	allocateResourceFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

	// drain the probe messages
	verify(rmServices.workerStore, Mockito.timeout(timeout.toMilliseconds())).putWorker(expected);
	assertThat(resourceManager.workersInNew, hasEntry(extractResourceID(taskID), expected));
	resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
	resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Launch.class);
	return expected;
}
 

@Override
public MesosWorkerStore createMesosWorkerStore(Configuration configuration, Executor executor) throws Exception {
	RetrievableStateStorageHelper<MesosWorkerStore.Worker> stateStorageHelper =
		ZooKeeperUtils.createFileSystemStateStorage(configuration, "mesosWorkerStore");

	ZooKeeperStateHandleStore<MesosWorkerStore.Worker> zooKeeperStateHandleStore = zooKeeperUtilityFactory.createZooKeeperStateHandleStore(
		"/workers",
		stateStorageHelper);

	ZooKeeperSharedValue frameworkId = zooKeeperUtilityFactory.createSharedValue("/frameworkId", new byte[0]);
	ZooKeeperSharedCount totalTaskCount = zooKeeperUtilityFactory.createSharedCount("/taskCount", 0);

	return new ZooKeeperMesosWorkerStore(
		zooKeeperStateHandleStore,
		frameworkId,
		totalTaskCount);
}
 

/**
 * Allocate a worker using the RM.
 */
public MesosWorkerStore.Worker allocateWorker(Protos.TaskID taskID, WorkerResourceSpec workerResourceSpec) throws Exception {
	when(rmServices.workerStore.newTaskID()).thenReturn(taskID);
	rmServices.slotManagerStarted.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

	CompletableFuture<Void> allocateResourceFuture = resourceManager.callAsync(
		() -> {
			rmServices.rmActions.allocateResource(workerResourceSpec);
			return null;
		},
		timeout);
	MesosWorkerStore.Worker expected = MesosWorkerStore.Worker.newWorker(taskID, workerResourceSpec);

	// check for exceptions
	allocateResourceFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

	// drain the probe messages
	verify(rmServices.workerStore, Mockito.timeout(timeout.toMilliseconds())).putWorker(expected);
	assertThat(resourceManager.workersInNew, hasEntry(extractResourceID(taskID), expected));
	resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
	resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Launch.class);
	return expected;
}
 

@Override
public MesosWorkerStore createMesosWorkerStore(Configuration configuration, Executor executor) throws Exception {
	RetrievableStateStorageHelper<MesosWorkerStore.Worker> stateStorageHelper =
		ZooKeeperUtils.createFileSystemStateStorage(configuration, "mesosWorkerStore");

	ZooKeeperStateHandleStore<MesosWorkerStore.Worker> zooKeeperStateHandleStore = zooKeeperUtilityFactory.createZooKeeperStateHandleStore(
		"/workers",
		stateStorageHelper);

	ZooKeeperSharedValue frameworkId = zooKeeperUtilityFactory.createSharedValue("/frameworkId", new byte[0]);
	ZooKeeperSharedCount totalTaskCount = zooKeeperUtilityFactory.createSharedCount("/taskCount", 0);

	return new ZooKeeperMesosWorkerStore(
		zooKeeperStateHandleStore,
		frameworkId,
		totalTaskCount);
}
 

@Override
public Collection<ResourceProfile> startNewWorker(ResourceProfile resourceProfile) {
	LOG.info("Starting a new worker.");
	try {
		// generate new workers into persistent state and launch associated actors
		MesosWorkerStore.Worker worker = MesosWorkerStore.Worker.newWorker(workerStore.newTaskID(), resourceProfile);
		workerStore.putWorker(worker);
		workersInNew.put(extractResourceID(worker.taskID()), worker);

		LaunchableMesosWorker launchable = createLaunchableMesosWorker(worker.taskID());

		LOG.info("Scheduling Mesos task {} with ({} MB, {} cpus).",
			launchable.taskID().getValue(), launchable.taskRequest().getMemory(), launchable.taskRequest().getCPUs());

		// tell the task monitor about the new plans
		taskMonitor.tell(new TaskMonitor.TaskGoalStateUpdated(extractGoalState(worker)), selfActor);

		// tell the launch coordinator to launch the new tasks
		launchCoordinator.tell(new LaunchCoordinator.Launch(Collections.singletonList(launchable)), selfActor);

		return slotsPerWorker;
	} catch (Exception ex) {
		onFatalError(new ResourceManagerException("Unable to request new workers.", ex));
		return Collections.emptyList();
	}
}
 

@Test
public void testClearStateAfterRevokeLeadership() throws Exception {
	new Context() {{
		final MesosWorkerStore.Worker worker1 = MesosWorkerStore.Worker.newWorker(task1, workerResourceSpec);
		final MesosWorkerStore.Worker worker2 = MesosWorkerStore.Worker.newWorker(task2, workerResourceSpec).launchWorker(slave1, slave1host);
		final MesosWorkerStore.Worker worker3 = MesosWorkerStore.Worker.newWorker(task3, workerResourceSpec).launchWorker(slave1, slave1host).releaseWorker();
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(Arrays.asList(worker1, worker2, worker3)).thenReturn(Collections.emptyList());

		startResourceManager();
		rmServices.rmLeaderElectionService.notLeader();
		rmServices.grantLeadership();

		assertThat(resourceManager.workersInNew.size(), equalTo(0));
		assertThat(resourceManager.workersInLaunch.size(), equalTo(0));
		assertThat(resourceManager.workersBeingReturned.size(), equalTo(0));
		verify(rmServices.schedulerDriver).stop(true);
	}};
}
 

/**
 * Test unplanned task failure of a pending worker.
 */
@Test
public void testWorkerFailed() throws Exception {
	new Context() {{
		// set the initial persistent state with a launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1, workerResourceSpec).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		when(rmServices.workerStore.newTaskID()).thenReturn(task2);
		startResourceManager();

		// tell the RM that a task failed
		when(rmServices.workerStore.removeWorker(task1)).thenReturn(true);
		resourceManager.taskTerminated(new TaskMonitor.TaskTerminated(task1, Protos.TaskStatus.newBuilder()
			.setTaskId(task1).setSlaveId(slave1).setState(Protos.TaskState.TASK_FAILED).build()));

		// verify that the instance state was updated
		verify(rmServices.workerStore).removeWorker(task1);
		assertThat(resourceManager.workersInLaunch.entrySet(), empty());
		assertThat(resourceManager.workersBeingReturned.entrySet(), empty());
		assertThat(resourceManager.workersInNew, hasKey(extractResourceID(task2)));

		// verify that `closeTaskManagerConnection` was called
		assertThat(resourceManager.closedTaskManagerConnections, hasItem(extractResourceID(task1)));
	}};
}
 

/**
 * Allocate a worker using the RM.
 */
public MesosWorkerStore.Worker allocateWorker(Protos.TaskID taskID, ResourceProfile resourceProfile) throws Exception {
	when(rmServices.workerStore.newTaskID()).thenReturn(taskID);
	rmServices.slotManagerStarted.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

	CompletableFuture<Void> allocateResourceFuture = resourceManager.callAsync(
		() -> {
			rmServices.rmActions.allocateResource(resourceProfile);
			return null;
		},
		timeout);
	MesosWorkerStore.Worker expected = MesosWorkerStore.Worker.newWorker(taskID, resourceProfile);

	// check for exceptions
	allocateResourceFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

	// drain the probe messages
	verify(rmServices.workerStore, Mockito.timeout(timeout.toMilliseconds())).putWorker(expected);
	assertThat(resourceManager.workersInNew, hasEntry(extractResourceID(taskID), expected));
	resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
	resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Launch.class);
	return expected;
}
 

@Override
public MesosWorkerStore createMesosWorkerStore(Configuration configuration, Executor executor) throws Exception {
	RetrievableStateStorageHelper<MesosWorkerStore.Worker> stateStorageHelper =
		ZooKeeperUtils.createFileSystemStateStorage(configuration, "mesosWorkerStore");

	ZooKeeperStateHandleStore<MesosWorkerStore.Worker> zooKeeperStateHandleStore = zooKeeperUtilityFactory.createZooKeeperStateHandleStore(
		"/workers",
		stateStorageHelper);

	ZooKeeperSharedValue frameworkId = zooKeeperUtilityFactory.createSharedValue("/frameworkId", new byte[0]);
	ZooKeeperSharedCount totalTaskCount = zooKeeperUtilityFactory.createSharedCount("/taskCount", 0);

	return new ZooKeeperMesosWorkerStore(
		zooKeeperStateHandleStore,
		frameworkId,
		totalTaskCount);
}
 

/**
 * Test recovery of persistent workers.
 */
@Test
public void testRecoverWorkers() throws Exception {
	new Context() {{
		// set the initial persistent state then initialize the RM
		MesosWorkerStore.Worker worker1 = MesosWorkerStore.Worker.newWorker(task1);
		MesosWorkerStore.Worker worker2 = MesosWorkerStore.Worker.newWorker(task2).launchWorker(slave1, slave1host);
		MesosWorkerStore.Worker worker3 = MesosWorkerStore.Worker.newWorker(task3).launchWorker(slave1, slave1host).releaseWorker();
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(Arrays.asList(worker1, worker2, worker3));
		startResourceManager();

		// verify that the internal state was updated, the task router was notified,
		// and the launch coordinator was asked to launch a task.
		// note: "new" workers are discarded
		assertThat(resourceManager.workersInNew.entrySet(), empty());
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task2), worker2));
		assertThat(resourceManager.workersBeingReturned, hasEntry(extractResourceID(task3), worker3));
		resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
		LaunchCoordinator.Assign actualAssign =
			resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Assign.class);
		assertThat(actualAssign.tasks(), hasSize(1));
		assertThat(actualAssign.tasks().get(0).f0.getId(), equalTo(task2.getValue()));
		assertThat(actualAssign.tasks().get(0).f1, equalTo(slave1host));
		resourceManager.launchCoordinator.expectNoMsg();
	}};
}
 

/**
 * Test worker registration after launch.
 */
@Test
public void testWorkerStarted() throws Exception {
	new Context() {{
		// set the initial state with a (recovered) launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		startResourceManager();
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));

		final int dataPort = 1234;
		final HardwareDescription hardwareDescription = new HardwareDescription(1, 2L, 3L, 4L);
		// send registration message
		CompletableFuture<RegistrationResponse> successfulFuture =
			resourceManager.registerTaskExecutor(task1Executor.address, task1Executor.resourceID, dataPort, hardwareDescription, timeout);
		RegistrationResponse response = successfulFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);
		final TaskExecutorRegistrationSuccess registrationResponse = (TaskExecutorRegistrationSuccess) response;

		final CompletableFuture<Acknowledge> initialSlotReportFuture = resourceManager.sendSlotReport(task1Executor.resourceID, registrationResponse.getRegistrationId(), slotReport, timeout);

		// check for errors
		initialSlotReportFuture.get();

		// verify the internal state
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));
	}};
}
 

/**
 * Test planned stop of a launched worker.
 */
@Test
public void testStopWorker() throws Exception {
	new Context() {{
		// set the initial persistent state with a launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		startResourceManager();

		// drain the assign message
		resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Assign.class);

		// tell the RM to stop the worker
		resourceManager.stopWorker(new RegisteredMesosWorkerNode(worker1launched));

		// verify that the instance state was updated
		MesosWorkerStore.Worker worker1Released = worker1launched.releaseWorker();
		verify(rmServices.workerStore).putWorker(worker1Released);
		assertThat(resourceManager.workersInLaunch.entrySet(), empty());
		assertThat(resourceManager.workersBeingReturned, hasEntry(extractResourceID(task1), worker1Released));

		// verify that the monitor was notified
		resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
		resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Unassign.class);
	}};
}
 

/**
 * Recovers given framework/worker information.
 *
 * @see #getWorkersAsync()
 */
private void recoverWorkers(final List<MesosWorkerStore.Worker> tasksFromPreviousAttempts) {
	assert(workersInNew.isEmpty());
	assert(workersInLaunch.isEmpty());
	assert(workersBeingReturned.isEmpty());

	if (!tasksFromPreviousAttempts.isEmpty()) {
		LOG.info("Retrieved {} TaskManagers from previous attempt", tasksFromPreviousAttempts.size());

		List<Tuple2<TaskRequest, String>> toAssign = new ArrayList<>(tasksFromPreviousAttempts.size());

		for (final MesosWorkerStore.Worker worker : tasksFromPreviousAttempts) {
			switch(worker.state()) {
				case Launched:
					workersInLaunch.put(extractResourceID(worker.taskID()), worker);
					final LaunchableMesosWorker launchable = createLaunchableMesosWorker(worker.taskID());
					toAssign.add(new Tuple2<>(launchable.taskRequest(), worker.hostname().get()));
					break;
				case Released:
					workersBeingReturned.put(extractResourceID(worker.taskID()), worker);
					break;
			}
			taskMonitor.tell(new TaskMonitor.TaskGoalStateUpdated(extractGoalState(worker)), selfActor);
		}

		// tell the launch coordinator about prior assignments
		if (toAssign.size() >= 1) {
			launchCoordinator.tell(new LaunchCoordinator.Assign(toAssign), selfActor);
		}
	}
}
 

/**
 * Test offer acceptance.
 */
@Test
public void testAcceptOffers() throws Exception {
	new Context() {{
		startResourceManager();

		// allocate a new worker
		MesosWorkerStore.Worker worker1 = allocateWorker(task1, resourceProfile1);

		// send an AcceptOffers message as the LaunchCoordinator would
		// to launch task1 onto slave1 with offer1
		Protos.TaskInfo task1info = Protos.TaskInfo.newBuilder()
			.setTaskId(task1).setName("").setSlaveId(slave1).build();
		AcceptOffers msg = new AcceptOffers(slave1host, singletonList(offer1), singletonList(launch(task1info)));
		resourceManager.acceptOffers(msg);

		// verify that the worker was persisted, the internal state was updated,
		// Mesos was asked to launch task1, and the task router was notified
		MesosWorkerStore.Worker worker1launched = worker1.launchWorker(slave1, slave1host);
		verify(rmServices.workerStore).putWorker(worker1launched);
		assertThat(resourceManager.workersInNew.entrySet(), empty());
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));
		resourceManager.taskRouter.expectMsg(
			new TaskMonitor.TaskGoalStateUpdated(extractGoalState(worker1launched)));
		verify(rmServices.schedulerDriver).acceptOffers(msg.offerIds(), msg.operations(), msg.filters());
	}};
}
 

/**
 * Recovers given framework/worker information.
 *
 * @see #getWorkersAsync()
 */
private void recoverWorkers(final List<MesosWorkerStore.Worker> tasksFromPreviousAttempts) {
	assert(workersInNew.isEmpty());
	assert(workersInLaunch.isEmpty());
	assert(workersBeingReturned.isEmpty());

	if (!tasksFromPreviousAttempts.isEmpty()) {
		LOG.info("Retrieved {} TaskManagers from previous attempt", tasksFromPreviousAttempts.size());

		List<Tuple2<TaskRequest, String>> toAssign = new ArrayList<>(tasksFromPreviousAttempts.size());

		for (final MesosWorkerStore.Worker worker : tasksFromPreviousAttempts) {
			switch(worker.state()) {
				case Launched:
					workersInLaunch.put(extractResourceID(worker.taskID()), worker);
					final LaunchableMesosWorker launchable = createLaunchableMesosWorker(worker.taskID());
					toAssign.add(new Tuple2<>(launchable.taskRequest(), worker.hostname().get()));
					break;
				case Released:
					workersBeingReturned.put(extractResourceID(worker.taskID()), worker);
					break;
			}
			taskMonitor.tell(new TaskMonitor.TaskGoalStateUpdated(extractGoalState(worker)), selfActor);
		}

		// tell the launch coordinator about prior assignments
		if (toAssign.size() >= 1) {
			launchCoordinator.tell(new LaunchCoordinator.Assign(toAssign), selfActor);
		}
	}
}
 

/**
 * Test request for new workers.
 */
@Test
public void testRequestNewWorkers() throws Exception {
	new Context() {{
		startResourceManager();

		// allocate a worker
		when(rmServices.workerStore.newTaskID()).thenReturn(task1).thenThrow(new AssertionFailedError());
		rmServices.slotManagerStarted.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		CompletableFuture<Void> allocateResourceFuture = resourceManager.callAsync(
			() -> {
				rmServices.rmActions.allocateResource(resourceProfile1);
				return null;
			},
			timeout);

		// check for exceptions
		allocateResourceFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		// verify that a new worker was persisted, the internal state was updated, the task router was notified,
		// and the launch coordinator was asked to launch a task
		MesosWorkerStore.Worker expected = MesosWorkerStore.Worker.newWorker(task1, resourceProfile1);
		verify(rmServices.workerStore, Mockito.timeout(timeout.toMilliseconds())).putWorker(expected);
		assertThat(resourceManager.workersInNew, hasEntry(extractResourceID(task1), expected));
		resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
		resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Launch.class);
	}};
}
 

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

/**
 * Test offer acceptance.
 */
@Test
public void testAcceptOffers() throws Exception {
	new Context() {{
		startResourceManager();

		// allocate a new worker
		MesosWorkerStore.Worker worker1 = allocateWorker(task1, workerResourceSpec);

		// send an AcceptOffers message as the LaunchCoordinator would
		// to launch task1 onto slave1 with offer1
		Protos.TaskInfo task1info = Protos.TaskInfo.newBuilder()
			.setTaskId(task1).setName("").setSlaveId(slave1).build();
		AcceptOffers msg = new AcceptOffers(slave1host, singletonList(offer1), singletonList(launch(task1info)));
		resourceManager.acceptOffers(msg);

		// verify that the worker was persisted, the internal state was updated,
		// Mesos was asked to launch task1, and the task router was notified
		MesosWorkerStore.Worker worker1launched = worker1.launchWorker(slave1, slave1host);
		verify(rmServices.workerStore).putWorker(worker1launched);
		assertThat(resourceManager.workersInNew.entrySet(), empty());
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));
		resourceManager.taskRouter.expectMsg(
			new TaskMonitor.TaskGoalStateUpdated(extractGoalState(worker1launched)));
		verify(rmServices.schedulerDriver).acceptOffers(msg.offerIds(), msg.operations(), msg.filters());
	}};
}
 

/**
 * Test recovery of persistent workers.
 */
@Test
public void testRecoverWorkers() throws Exception {
	new Context() {{
		// set the initial persistent state then initialize the RM
		MesosWorkerStore.Worker worker1 = MesosWorkerStore.Worker.newWorker(task1, workerResourceSpec);
		MesosWorkerStore.Worker worker2 = MesosWorkerStore.Worker.newWorker(task2, workerResourceSpec).launchWorker(slave1, slave1host);
		MesosWorkerStore.Worker worker3 = MesosWorkerStore.Worker.newWorker(task3, workerResourceSpec).launchWorker(slave1, slave1host).releaseWorker();
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(Arrays.asList(worker1, worker2, worker3));
		startResourceManager();

		// verify that the internal state was updated, the task router was notified,
		// and the launch coordinator was asked to launch a task.
		// note: "new" workers are discarded
		assertThat(resourceManager.workersInNew.entrySet(), empty());
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task2), worker2));
		assertThat(resourceManager.workersBeingReturned, hasEntry(extractResourceID(task3), worker3));
		resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
		LaunchCoordinator.Assign actualAssign =
			resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Assign.class);
		assertThat(actualAssign.tasks(), hasSize(1));
		assertThat(actualAssign.tasks().get(0).f0.getId(), equalTo(task2.getValue()));
		assertThat(actualAssign.tasks().get(0).f1, equalTo(slave1host));
		resourceManager.launchCoordinator.expectNoMsg();
	}};
}
 

/**
 * Test request for new workers.
 */
@Test
public void testRequestNewWorkers() throws Exception {
	new Context() {{
		startResourceManager();

		// allocate a worker
		when(rmServices.workerStore.newTaskID()).thenReturn(task1).thenThrow(new AssertionFailedError());
		rmServices.slotManagerStarted.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		CompletableFuture<Void> allocateResourceFuture = resourceManager.callAsync(
			() -> {
				rmServices.rmActions.allocateResource(workerResourceSpec);
				return null;
			},
			timeout);

		// check for exceptions
		allocateResourceFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		// verify that a new worker was persisted, the internal state was updated, the task router was notified,
		// and the launch coordinator was asked to launch a task
		MesosWorkerStore.Worker expected = MesosWorkerStore.Worker.newWorker(task1, workerResourceSpec);
		verify(rmServices.workerStore, Mockito.timeout(timeout.toMilliseconds())).putWorker(expected);
		assertThat(resourceManager.workersInNew, hasEntry(extractResourceID(task1), expected));
		resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
		resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Launch.class);
	}};
}
 

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

/**
 * Test worker registration after launch.
 */
@Test
public void testWorkerStarted() throws Exception {
	new Context() {{
		// set the initial state with a (recovered) launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1, workerResourceSpec).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		startResourceManager();
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));

		final int dataPort = 1234;
		final HardwareDescription hardwareDescription = new HardwareDescription(1, 2L, 3L, 4L);
		// send registration message
		TaskExecutorRegistration taskExecutorRegistration = new TaskExecutorRegistration(
			task1Executor.address,
			task1Executor.resourceID,
			dataPort,
			hardwareDescription,
			ResourceProfile.ZERO,
			ResourceProfile.ZERO);
		CompletableFuture<RegistrationResponse> successfulFuture =
			resourceManager.registerTaskExecutor(taskExecutorRegistration, timeout);
		RegistrationResponse response = successfulFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);
		final TaskExecutorRegistrationSuccess registrationResponse = (TaskExecutorRegistrationSuccess) response;

		final CompletableFuture<Acknowledge> initialSlotReportFuture = resourceManager.sendSlotReport(task1Executor.resourceID, registrationResponse.getRegistrationId(), slotReport, timeout);

		// check for errors
		initialSlotReportFuture.get();

		// verify the internal state
		assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));
	}};
}
 

/**
 * Test planned stop of a launched worker.
 */
@Test
public void testStopWorker() throws Exception {
	new Context() {{
		// set the initial persistent state with a launched worker
		MesosWorkerStore.Worker worker1launched = MesosWorkerStore.Worker.newWorker(task1, workerResourceSpec).launchWorker(slave1, slave1host);
		when(rmServices.workerStore.getFrameworkID()).thenReturn(Option.apply(framework1));
		when(rmServices.workerStore.recoverWorkers()).thenReturn(singletonList(worker1launched));
		startResourceManager();

		// drain the assign message
		resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Assign.class);

		// tell the RM to stop the worker
		resourceManager.stopWorker(new RegisteredMesosWorkerNode(worker1launched));

		// verify that the instance state was updated
		MesosWorkerStore.Worker worker1Released = worker1launched.releaseWorker();
		verify(rmServices.workerStore).putWorker(worker1Released);
		assertThat(resourceManager.workersInLaunch.entrySet(), empty());
		assertThat(resourceManager.workersBeingReturned, hasEntry(extractResourceID(task1), worker1Released));

		// verify that the monitor was notified
		resourceManager.taskRouter.expectMsgClass(TaskMonitor.TaskGoalStateUpdated.class);
		resourceManager.launchCoordinator.expectMsgClass(LaunchCoordinator.Unassign.class);
	}};
}