org.apache.flink.runtime.instance.InstanceID Java Examples

/**
 * Tests that a TaskExecutor can disconnect from the {@link ResourceManager}.
 */
@Test
public void testDisconnectTaskExecutor() throws Exception {
	final RegistrationResponse registrationResponse = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress()).get();
	assertThat(registrationResponse, instanceOf(TaskExecutorRegistrationSuccess.class));

	final InstanceID registrationId = ((TaskExecutorRegistrationSuccess) registrationResponse).getRegistrationId();
	final int numberSlots = 10;
	final Collection<SlotStatus> slots = createSlots(numberSlots);
	final SlotReport slotReport = new SlotReport(slots);
	rmGateway.sendSlotReport(taskExecutorResourceID, registrationId, slotReport, TIMEOUT).get();

	final ResourceOverview resourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(resourceOverview.getNumberTaskManagers(), is(1));
	assertThat(resourceOverview.getNumberRegisteredSlots(), is(numberSlots));

	rmGateway.disconnectTaskManager(taskExecutorResourceID, new FlinkException("testDisconnectTaskExecutor"));

	final ResourceOverview afterDisconnectResourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(afterDisconnectResourceOverview.getNumberTaskManagers(), is(0));
	assertThat(afterDisconnectResourceOverview.getNumberRegisteredSlots(), is(0));
}
 

@Override
@VisibleForTesting
public void unregisterTaskManagersAndReleaseResources() {
	Iterator<Map.Entry<InstanceID, TaskManagerRegistration>> taskManagerRegistrationIterator =
			taskManagerRegistrations.entrySet().iterator();

	while (taskManagerRegistrationIterator.hasNext()) {
		TaskManagerRegistration taskManagerRegistration =
				taskManagerRegistrationIterator.next().getValue();

		taskManagerRegistrationIterator.remove();

		internalUnregisterTaskManager(taskManagerRegistration);

		resourceActions.releaseResource(taskManagerRegistration.getInstanceId(), new FlinkException("Triggering of SlotManager#unregisterTaskManagersAndReleaseResources."));
	}
}
 

/**
 * Reports the current slot allocations for a task manager identified by the given instance id.
 *
 * @param instanceId identifying the task manager for which to report the slot status
 * @param slotReport containing the status for all of its slots
 * @return true if the slot status has been updated successfully, otherwise false
 */
public boolean reportSlotStatus(InstanceID instanceId, SlotReport slotReport) {
	checkInit();

	LOG.debug("Received slot report from instance {}: {}.", instanceId, slotReport);

	TaskManagerRegistration taskManagerRegistration = taskManagerRegistrations.get(instanceId);

	if (null != taskManagerRegistration) {

		for (SlotStatus slotStatus : slotReport) {
			updateSlot(slotStatus.getSlotID(), slotStatus.getAllocationID(), slotStatus.getJobID());
		}

		return true;
	} else {
		LOG.debug("Received slot report for unknown task manager with instance id {}. Ignoring this report.", instanceId);

		return false;
	}
}
 

/**
 * Unregisters the task manager identified by the given instance id and its associated slots
 * from the slot manager.
 *
 * @param instanceId identifying the task manager to unregister
 * @return True if there existed a registered task manager with the given instance id
 */
public boolean unregisterTaskManager(InstanceID instanceId) {
	checkInit();

	LOG.debug("Unregister TaskManager {} from the SlotManager.", instanceId);

	TaskManagerRegistration taskManagerRegistration = taskManagerRegistrations.remove(instanceId);

	if (null != taskManagerRegistration) {
		internalUnregisterTaskManager(taskManagerRegistration);

		return true;
	} else {
		LOG.debug("There is no task manager registered with instance ID {}. Ignoring this message.", instanceId);

		return false;
	}
}
 

protected void releaseResource(InstanceID instanceId, Exception cause) {
	WorkerType worker = null;

	// TODO: Improve performance by having an index on the instanceId
	for (Map.Entry<ResourceID, WorkerRegistration<WorkerType>> entry : taskExecutors.entrySet()) {
		if (entry.getValue().getInstanceID().equals(instanceId)) {
			worker = entry.getValue().getWorker();
			break;
		}
	}

	if (worker != null) {
		if (stopWorker(worker)) {
			closeTaskManagerConnection(worker.getResourceID(), cause);
		} else {
			log.debug("Worker {} could not be stopped.", worker.getResourceID());
		}
	} else {
		// unregister in order to clean up potential left over state
		slotManager.unregisterTaskManager(instanceId);
	}
}
 

@VisibleForTesting
public void unregisterTaskManagersAndReleaseResources() {
	Iterator<Map.Entry<InstanceID, TaskManagerRegistration>> taskManagerRegistrationIterator =
			taskManagerRegistrations.entrySet().iterator();

	while (taskManagerRegistrationIterator.hasNext()) {
		TaskManagerRegistration taskManagerRegistration =
				taskManagerRegistrationIterator.next().getValue();

		taskManagerRegistrationIterator.remove();

		internalUnregisterTaskManager(taskManagerRegistration);

		resourceActions.releaseResource(taskManagerRegistration.getInstanceId(), new FlinkException("Triggering of SlotManager#unregisterTaskManagersAndReleaseResources."));
	}
}
 

@Override
public CompletableFuture<TaskManagerInfo> requestTaskManagerInfo(ResourceID resourceId, Time timeout) {

	final WorkerRegistration<WorkerType> taskExecutor = taskExecutors.get(resourceId);

	if (taskExecutor == null) {
		return FutureUtils.completedExceptionally(new UnknownTaskExecutorException(resourceId));
	} else {
		final InstanceID instanceId = taskExecutor.getInstanceID();
		final TaskManagerInfo taskManagerInfo = new TaskManagerInfo(
			resourceId,
			taskExecutor.getTaskExecutorGateway().getAddress(),
			taskExecutor.getDataPort(),
			taskManagerHeartbeatManager.getLastHeartbeatFrom(resourceId),
			slotManager.getNumberRegisteredSlotsOf(instanceId),
			slotManager.getNumberFreeSlotsOf(instanceId),
			taskExecutor.getHardwareDescription());

		return CompletableFuture.completedFuture(taskManagerInfo);
	}
}
 

@Override
public void notifySlotAvailable(
		final InstanceID instanceID,
		final SlotID slotId,
		final AllocationID allocationId) {

	final ResourceID resourceId = slotId.getResourceID();
	WorkerRegistration<WorkerType> registration = taskExecutors.get(resourceId);

	if (registration != null) {
		InstanceID registrationId = registration.getInstanceID();

		if (Objects.equals(registrationId, instanceID)) {
			slotManager.freeSlot(slotId, allocationId);
		} else {
			log.debug("Invalid registration id for slot available message. This indicates an" +
				" outdated request.");
		}
	} else {
		log.debug("Could not find registration for resource id {}. Discarding the slot available" +
			"message {}.", resourceId, slotId);
	}
}
 

/**
 * Tests that the slot manager ignores slot reports of unknown origin (not registered
 * task managers).
 */
@Test
public void testReceivingUnknownSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);

	final InstanceID unknownInstanceID = new InstanceID();
	final SlotID unknownSlotId = new SlotID(ResourceID.generate(), 0);
	final ResourceProfile unknownResourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus unknownSlotStatus = new SlotStatus(unknownSlotId, unknownResourceProfile);
	final SlotReport unknownSlotReport = new SlotReport(unknownSlotStatus);

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		// check that we don't have any slots registered
		assertTrue(0 == slotManager.getNumberRegisteredSlots());

		// this should not update anything since the instance id is not known to the slot manager
		assertFalse(slotManager.reportSlotStatus(unknownInstanceID, unknownSlotReport));

		assertTrue(0 == slotManager.getNumberRegisteredSlots());
	}
}
 

protected void releaseResource(InstanceID instanceId, Exception cause) {
	WorkerType worker = null;

	// TODO: Improve performance by having an index on the instanceId
	for (Map.Entry<ResourceID, WorkerRegistration<WorkerType>> entry : taskExecutors.entrySet()) {
		if (entry.getValue().getInstanceID().equals(instanceId)) {
			worker = entry.getValue().getWorker();
			break;
		}
	}

	if (worker != null) {
		if (stopWorker(worker)) {
			closeTaskManagerConnection(worker.getResourceID(), cause);
		} else {
			log.debug("Worker {} could not be stopped.", worker.getResourceID());
		}
	} else {
		// unregister in order to clean up potential left over state
		slotManager.unregisterTaskManager(instanceId);
	}
}
 

@Override
public CompletableFuture<TaskManagerInfo> requestTaskManagerInfo(ResourceID resourceId, Time timeout) {

	final WorkerRegistration<WorkerType> taskExecutor = taskExecutors.get(resourceId);

	if (taskExecutor == null) {
		return FutureUtils.completedExceptionally(new UnknownTaskExecutorException(resourceId));
	} else {
		final InstanceID instanceId = taskExecutor.getInstanceID();
		final TaskManagerInfo taskManagerInfo = new TaskManagerInfo(
			resourceId,
			taskExecutor.getTaskExecutorGateway().getAddress(),
			taskExecutor.getDataPort(),
			taskManagerHeartbeatManager.getLastHeartbeatFrom(resourceId),
			slotManager.getNumberRegisteredSlotsOf(instanceId),
			slotManager.getNumberFreeSlotsOf(instanceId),
			taskExecutor.getHardwareDescription());

		return CompletableFuture.completedFuture(taskManagerInfo);
	}
}
 

@Override
public void notifySlotAvailable(
		final InstanceID instanceID,
		final SlotID slotId,
		final AllocationID allocationId) {

	final ResourceID resourceId = slotId.getResourceID();
	WorkerRegistration<WorkerType> registration = taskExecutors.get(resourceId);

	if (registration != null) {
		InstanceID registrationId = registration.getInstanceID();

		if (Objects.equals(registrationId, instanceID)) {
			slotManager.freeSlot(slotId, allocationId);
		} else {
			log.debug("Invalid registration id for slot available message. This indicates an" +
				" outdated request.");
		}
	} else {
		log.debug("Could not find registration for resource id {}. Discarding the slot available" +
			"message {}.", resourceId, slotId);
	}
}
 

@Override
public void handleMessage(Object message) throws Exception {
	if (message instanceof RegistrationMessages.RegisterTaskManager) {
		final InstanceID iid = new InstanceID();
		final ActorRef self = getSelf();
		getSender().tell(
				decorateMessage(
						new RegistrationMessages.AcknowledgeRegistration(
							iid,
							12345)
				),
				self);
	}
	else if (message instanceof TaskMessages.UpdateTaskExecutionState){
		getSender().tell(true, getSelf());
	}
}
 

/**
 * Unregisters the task manager identified by the given instance id and its associated slots
 * from the slot manager.
 *
 * @param instanceId identifying the task manager to unregister
 * @return True if there existed a registered task manager with the given instance id
 */
@Override
public boolean unregisterTaskManager(InstanceID instanceId) {
	checkInit();

	LOG.debug("Unregister TaskManager {} from the SlotManager.", instanceId);

	TaskManagerRegistration taskManagerRegistration = taskManagerRegistrations.remove(instanceId);

	if (null != taskManagerRegistration) {
		internalUnregisterTaskManager(taskManagerRegistration);

		return true;
	} else {
		LOG.debug("There is no task manager registered with instance ID {}. Ignoring this message.", instanceId);

		return false;
	}
}
 

/**
 * Tests that a TaskExecutor can disconnect from the {@link ResourceManager}.
 */
@Test
public void testDisconnectTaskExecutor() throws Exception {
	final RegistrationResponse registrationResponse = registerTaskExecutor(rmGateway, taskExecutorGateway.getAddress()).get();
	assertThat(registrationResponse, instanceOf(TaskExecutorRegistrationSuccess.class));

	final InstanceID registrationId = ((TaskExecutorRegistrationSuccess) registrationResponse).getRegistrationId();
	final int numberSlots = 10;
	final Collection<SlotStatus> slots = createSlots(numberSlots);
	final SlotReport slotReport = new SlotReport(slots);
	rmGateway.sendSlotReport(taskExecutorResourceID, registrationId, slotReport, TIMEOUT).get();

	final ResourceOverview resourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(resourceOverview.getNumberTaskManagers(), is(1));
	assertThat(resourceOverview.getNumberRegisteredSlots(), is(numberSlots));

	rmGateway.disconnectTaskManager(taskExecutorResourceID, new FlinkException("testDisconnectTaskExecutor"));

	final ResourceOverview afterDisconnectResourceOverview = rmGateway.requestResourceOverview(TIMEOUT).get();
	assertThat(afterDisconnectResourceOverview.getNumberTaskManagers(), is(0));
	assertThat(afterDisconnectResourceOverview.getNumberRegisteredSlots(), is(0));
}
 

private void releaseTaskExecutorIfPossible(TaskManagerRegistration taskManagerRegistration) {
	long idleSince = taskManagerRegistration.getIdleSince();
	taskManagerRegistration
		.getTaskManagerConnection()
		.getTaskExecutorGateway()
		.canBeReleased()
		.thenAcceptAsync(
			canBeReleased -> {
				InstanceID timedOutTaskManagerId = taskManagerRegistration.getInstanceId();
				boolean stillIdle = idleSince == taskManagerRegistration.getIdleSince();
				if (stillIdle && canBeReleased) {
					releaseTaskExecutor(timedOutTaskManagerId);
				}
			},
			mainThreadExecutor);
}
 

/**
 * Tests that the slot manager ignores slot reports of unknown origin (not registered
 * task managers).
 */
@Test
public void testReceivingUnknownSlotReport() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);

	final InstanceID unknownInstanceID = new InstanceID();
	final SlotID unknownSlotId = new SlotID(ResourceID.generate(), 0);
	final ResourceProfile unknownResourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus unknownSlotStatus = new SlotStatus(unknownSlotId, unknownResourceProfile);
	final SlotReport unknownSlotReport = new SlotReport(unknownSlotStatus);

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		// check that we don't have any slots registered
		assertTrue(0 == slotManager.getNumberRegisteredSlots());

		// this should not update anything since the instance id is not known to the slot manager
		assertFalse(slotManager.reportSlotStatus(unknownInstanceID, unknownSlotReport));

		assertTrue(0 == slotManager.getNumberRegisteredSlots());
	}
}
 

/**
 * Tests that idle task managers time out after the configured timeout. A timed out task manager
 * will be removed from the slot manager and the resource manager will be notified about the
 * timeout, if it can be released.
 */
@Test
public void testTaskManagerTimeout() throws Exception {
	final long tmTimeout = 10L;

	final CompletableFuture<InstanceID> releaseFuture = new CompletableFuture<>();
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setReleaseResourceConsumer((instanceID, e) -> releaseFuture.complete(instanceID))
		.build();
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	final SlotID slotId = new SlotID(resourceID, 0);
	final ResourceProfile resourceProfile = new ResourceProfile(1.0, 1);
	final SlotStatus slotStatus = new SlotStatus(slotId, resourceProfile);
	final SlotReport slotReport = new SlotReport(slotStatus);

	final Executor mainThreadExecutor = TestingUtils.defaultExecutor();

	try (SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(Time.milliseconds(tmTimeout))
		.build()) {

		slotManager.start(resourceManagerId, mainThreadExecutor, resourceManagerActions);

		mainThreadExecutor.execute(() -> slotManager.registerTaskManager(taskManagerConnection, slotReport));

		assertThat(releaseFuture.get(), is(equalTo(taskManagerConnection.getInstanceID())));
	}
}
 

/**
 * Suspends the component. This clears the internal state of the slot manager.
 */
@Override
public void suspend() {
	LOG.info("Suspending the SlotManager.");

	// stop the timeout checks for the TaskManagers and the SlotRequests
	if (taskManagerTimeoutCheck != null) {
		taskManagerTimeoutCheck.cancel(false);
		taskManagerTimeoutCheck = null;
	}

	if (slotRequestTimeoutCheck != null) {
		slotRequestTimeoutCheck.cancel(false);
		slotRequestTimeoutCheck = null;
	}

	for (PendingSlotRequest pendingSlotRequest : pendingSlotRequests.values()) {
		cancelPendingSlotRequest(pendingSlotRequest);
	}

	pendingSlotRequests.clear();

	ArrayList<InstanceID> registeredTaskManagers = new ArrayList<>(taskManagerRegistrations.keySet());

	for (InstanceID registeredTaskManager : registeredTaskManagers) {
		unregisterTaskManager(registeredTaskManager);
	}

	resourceManagerId = null;
	resourceActions = null;
	started = false;
}
 

/**
 * Create a new {@code TaskExecutorRegistrationSuccess} message.
 *
 * @param registrationId The ID that the ResourceManager assigned the registration.
 * @param resourceManagerResourceId The unique ID that identifies the ResourceManager.
 * @param clusterInformation information about the cluster
 */
public TaskExecutorRegistrationSuccess(
		InstanceID registrationId,
		ResourceID resourceManagerResourceId,
		ClusterInformation clusterInformation) {
	this.registrationId = Preconditions.checkNotNull(registrationId);
	this.resourceManagerResourceId = Preconditions.checkNotNull(resourceManagerResourceId);
	this.clusterInformation = Preconditions.checkNotNull(clusterInformation);
}
 

@Override
public int getNumberFreeSlotsOf(InstanceID instanceId) {
	TaskManagerRegistration taskManagerRegistration = taskManagerRegistrations.get(instanceId);

	if (taskManagerRegistration != null) {
		return taskManagerRegistration.getNumberFreeSlots();
	} else {
		return 0;
	}
}
 

/**
 * Suspends the component. This clears the internal state of the slot manager.
 */
public void suspend() {
	LOG.info("Suspending the SlotManager.");

	// stop the timeout checks for the TaskManagers and the SlotRequests
	if (taskManagerTimeoutCheck != null) {
		taskManagerTimeoutCheck.cancel(false);
		taskManagerTimeoutCheck = null;
	}

	if (slotRequestTimeoutCheck != null) {
		slotRequestTimeoutCheck.cancel(false);
		slotRequestTimeoutCheck = null;
	}

	for (PendingSlotRequest pendingSlotRequest : pendingSlotRequests.values()) {
		cancelPendingSlotRequest(pendingSlotRequest);
	}

	pendingSlotRequests.clear();

	ArrayList<InstanceID> registeredTaskManagers = new ArrayList<>(taskManagerRegistrations.keySet());

	for (InstanceID registeredTaskManager : registeredTaskManagers) {
		unregisterTaskManager(registeredTaskManager);
	}

	resourceManagerId = null;
	resourceActions = null;
	started = false;
}
 

@Override
public void notifySlotAvailable(InstanceID instanceId, SlotID slotID, AllocationID oldAllocationId) {
	final Consumer<Tuple3<InstanceID, SlotID, AllocationID>> currentNotifySlotAvailableConsumer = notifySlotAvailableConsumer;

	if (currentNotifySlotAvailableConsumer != null) {
		currentNotifySlotAvailableConsumer.accept(Tuple3.of(instanceId, slotID, oldAllocationId));
	}
}
 

public static Instance getInstance(final TaskManagerGateway gateway, final int numberOfSlots) throws Exception {
	ResourceID resourceID = ResourceID.generate();
	HardwareDescription hardwareDescription = new HardwareDescription(4, 2L*1024*1024*1024, 1024*1024*1024, 512*1024*1024);
	InetAddress address = InetAddress.getByName("127.0.0.1");
	TaskManagerLocation connection = new TaskManagerLocation(resourceID, address, 10001);

	return new Instance(gateway, connection, new InstanceID(), hardwareDescription, numberOfSlots);
}
 

public int getNumberFreeSlotsOf(InstanceID instanceId) {
	TaskManagerRegistration taskManagerRegistration = taskManagerRegistrations.get(instanceId);

	if (taskManagerRegistration != null) {
		return taskManagerRegistration.getNumberFreeSlots();
	} else {
		return 0;
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
	final WorkerRegistration<WorkerType> workerTypeWorkerRegistration = taskExecutors.get(taskManagerResourceId);

	if (workerTypeWorkerRegistration.getInstanceID().equals(taskManagerRegistrationId)) {
		slotManager.registerTaskManager(workerTypeWorkerRegistration, slotReport);
		return CompletableFuture.completedFuture(Acknowledge.get());
	} else {
		return FutureUtils.completedExceptionally(new ResourceManagerException(String.format("Unknown TaskManager registration id %s.", taskManagerRegistrationId)));
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
	final Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> currentSendSlotReportFunction = sendSlotReportFunction;

	if (currentSendSlotReportFunction != null) {
		return currentSendSlotReportFunction.apply(Tuple3.of(taskManagerResourceId, taskManagerRegistrationId, slotReport));
	} else {
		return CompletableFuture.completedFuture(Acknowledge.get());
	}
}
 

@Override
public CompletableFuture<RegistrationResponse> registerTaskExecutor(String taskExecutorAddress, ResourceID resourceId, int dataPort, HardwareDescription hardwareDescription, Time timeout) {
	final Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> currentFunction = registerTaskExecutorFunction;

	if (currentFunction != null) {
		return currentFunction.apply(Tuple4.of(taskExecutorAddress, resourceId, dataPort, hardwareDescription));
	} else {
		return CompletableFuture.completedFuture(
			new TaskExecutorRegistrationSuccess(
				new InstanceID(),
				ownResourceId,
				new ClusterInformation("localhost", 1234)));
	}
}
 

@Override
public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
	final Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> currentSendSlotReportFunction = sendSlotReportFunction;

	if (currentSendSlotReportFunction != null) {
		return currentSendSlotReportFunction.apply(Tuple3.of(taskManagerResourceId, taskManagerRegistrationId, slotReport));
	} else {
		return CompletableFuture.completedFuture(Acknowledge.get());
	}
}
 

public TestingResourceActions(
		@Nonnull BiConsumer<InstanceID, Exception> releaseResourceConsumer,
		@Nonnull FunctionWithException<ResourceProfile, Collection<ResourceProfile>, ResourceManagerException> allocateResourceFunction,
		@Nonnull Consumer<Tuple3<JobID, AllocationID, Exception>> notifyAllocationFailureConsumer) {
	this.releaseResourceConsumer = releaseResourceConsumer;
	this.allocateResourceFunction = allocateResourceFunction;
	this.notifyAllocationFailureConsumer = notifyAllocationFailureConsumer;
}