org.apache.flink.runtime.resourcemanager.ResourceManagerId Java Examples

/**
 * Starts the slot manager with the given leader id and resource manager actions.
 *
 * @param newResourceManagerId to use for communication with the task managers
 * @param newMainThreadExecutor to use to run code in the ResourceManager's main thread
 * @param newResourceActions to use for resource (de-)allocations
 */
@Override
public void start(ResourceManagerId newResourceManagerId, Executor newMainThreadExecutor, ResourceActions newResourceActions) {
	LOG.info("Starting the SlotManager.");

	this.resourceManagerId = Preconditions.checkNotNull(newResourceManagerId);
	mainThreadExecutor = Preconditions.checkNotNull(newMainThreadExecutor);
	resourceActions = Preconditions.checkNotNull(newResourceActions);

	started = true;

	taskManagerTimeoutCheck = scheduledExecutor.scheduleWithFixedDelay(
		() -> mainThreadExecutor.execute(
			() -> checkTaskManagerTimeouts()),
		0L,
		taskManagerTimeout.toMilliseconds(),
		TimeUnit.MILLISECONDS);

	slotRequestTimeoutCheck = scheduledExecutor.scheduleWithFixedDelay(
		() -> mainThreadExecutor.execute(
			() -> checkSlotRequestTimeouts()),
		0L,
		slotRequestTimeout.toMilliseconds(),
		TimeUnit.MILLISECONDS);
}
 

TestingTaskExecutorGateway(
		String address,
		String hostname,
		BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer,
		BiConsumer<JobID, Throwable> disconnectJobManagerConsumer,
		BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer,
		Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction,
		BiFunction<AllocationID, Throwable, CompletableFuture<Acknowledge>> freeSlotFunction,
		Consumer<ResourceID> heartbeatResourceManagerConsumer,
		Consumer<Exception> disconnectResourceManagerConsumer,
		Function<ExecutionAttemptID, CompletableFuture<Acknowledge>> cancelTaskFunction,
		Supplier<CompletableFuture<Boolean>> canBeReleasedSupplier,
		BiConsumer<JobID, Collection<ResultPartitionID>> releasePartitionsConsumer) {
	this.address = Preconditions.checkNotNull(address);
	this.hostname = Preconditions.checkNotNull(hostname);
	this.heartbeatJobManagerConsumer = Preconditions.checkNotNull(heartbeatJobManagerConsumer);
	this.disconnectJobManagerConsumer = Preconditions.checkNotNull(disconnectJobManagerConsumer);
	this.submitTaskConsumer = Preconditions.checkNotNull(submitTaskConsumer);
	this.requestSlotFunction = Preconditions.checkNotNull(requestSlotFunction);
	this.freeSlotFunction = Preconditions.checkNotNull(freeSlotFunction);
	this.heartbeatResourceManagerConsumer = heartbeatResourceManagerConsumer;
	this.disconnectResourceManagerConsumer = disconnectResourceManagerConsumer;
	this.cancelTaskFunction = cancelTaskFunction;
	this.canBeReleasedSupplier = canBeReleasedSupplier;
	this.releasePartitionsConsumer = releasePartitionsConsumer;
}
 

@Override
protected RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess> generateRegistration() {
	return new RetryingRegistration<ResourceManagerId, ResourceManagerGateway, JobMasterRegistrationSuccess>(
		log,
		getRpcService(),
		"ResourceManager",
		ResourceManagerGateway.class,
		getTargetAddress(),
		getTargetLeaderId(),
		jobMasterConfiguration.getRetryingRegistrationConfiguration()) {

		@Override
		protected CompletableFuture<RegistrationResponse> invokeRegistration(
				ResourceManagerGateway gateway, ResourceManagerId fencingToken, long timeoutMillis) {
			Time timeout = Time.milliseconds(timeoutMillis);

			return gateway.registerJobManager(
				jobMasterId,
				jobManagerResourceID,
				jobManagerRpcAddress,
				jobID,
				timeout);
		}
	};
}
 

private static SlotManager createSlotManager(
		List<Tuple3<JobID, AllocationID, Exception>> notifiedAllocationFailures,
		boolean startNewTMs) {

	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction((resourceProfile) -> startNewTMs ?
						Collections.singleton(resourceProfile) :
						Collections.emptyList())
		.setNotifyAllocationFailureConsumer(tuple3 -> notifiedAllocationFailures.add(tuple3))
		.build();

	SlotManager slotManager = SlotManagerBuilder.newBuilder().build();
	slotManager.start(ResourceManagerId.generate(), Executors.directExecutor(), resourceManagerActions);

	return slotManager;
}
 

/**
 * Tests that a second pending slot request is detected as a duplicate if the allocation ids are
 * the same.
 */
@Test
public void testDuplicatePendingSlotRequest() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final AtomicInteger numberAllocateResourceFunctionCalls = new AtomicInteger(0);
	final ResourceActions resourceManagerActions = new TestingResourceActionsBuilder()
		.setAllocateResourceConsumer(resourceProfile -> numberAllocateResourceFunctionCalls.incrementAndGet())
		.build();
	final AllocationID allocationId = new AllocationID();
	final ResourceProfile resourceProfile1 = new ResourceProfile(1.0, 2);
	final ResourceProfile resourceProfile2 = new ResourceProfile(2.0, 1);
	final SlotRequest slotRequest1 = new SlotRequest(new JobID(), allocationId, resourceProfile1, "foobar");
	final SlotRequest slotRequest2 = new SlotRequest(new JobID(), allocationId, resourceProfile2, "barfoo");

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		assertTrue(slotManager.registerSlotRequest(slotRequest1));
		assertFalse(slotManager.registerSlotRequest(slotRequest2));
	}

	// check that we have only called the resource allocation only for the first slot request,
	// since the second request is a duplicate
	assertThat(numberAllocateResourceFunctionCalls.get(), is(1));
}
 

/**
 * Tests that only free slots can fulfill/complete a pending task manager slot.
 */
@Test
public void testOnlyFreeSlotsCanFulfillPendingTaskManagerSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));

		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final SlotID slotId = new SlotID(taskExecutorConnection.getResourceID(), 0);
		final SlotStatus slotStatus = new SlotStatus(slotId, ResourceProfile.UNKNOWN, jobId, new AllocationID());
		final SlotReport slotReport = new SlotReport(slotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(1));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));
	}
}
 

ResourceManagerRegistration(
		Logger log,
		RpcService rpcService,
		String targetAddress,
		ResourceManagerId resourceManagerId,
		RetryingRegistrationConfiguration retryingRegistrationConfiguration,
		String taskExecutorAddress,
		ResourceID resourceID,
		int dataPort,
		HardwareDescription hardwareDescription) {

	super(log, rpcService, "ResourceManager", ResourceManagerGateway.class, targetAddress, resourceManagerId, retryingRegistrationConfiguration);
	this.taskExecutorAddress = checkNotNull(taskExecutorAddress);
	this.resourceID = checkNotNull(resourceID);
	this.dataPort = dataPort;
	this.hardwareDescription = checkNotNull(hardwareDescription);
}
 

public TaskExecutorToResourceManagerConnection(
		Logger log,
		RpcService rpcService,
		String taskManagerAddress,
		ResourceID taskManagerResourceId,
		RetryingRegistrationConfiguration retryingRegistrationConfiguration,
		int dataPort,
		HardwareDescription hardwareDescription,
		String resourceManagerAddress,
		ResourceManagerId resourceManagerId,
		Executor executor,
		RegistrationConnectionListener<TaskExecutorToResourceManagerConnection, TaskExecutorRegistrationSuccess> registrationListener) {

	super(log, resourceManagerAddress, resourceManagerId, executor);

	this.rpcService = checkNotNull(rpcService);
	this.taskManagerAddress = checkNotNull(taskManagerAddress);
	this.taskManagerResourceId = checkNotNull(taskManagerResourceId);
	this.retryingRegistrationConfiguration = checkNotNull(retryingRegistrationConfiguration);
	this.dataPort = dataPort;
	this.hardwareDescription = checkNotNull(hardwareDescription);
	this.registrationListener = checkNotNull(registrationListener);
}
 

/**
 * Tests the completion of pending task manager slots by registering a TaskExecutor.
 */
@Test
public void testPendingTaskManagerSlotCompletion() throws Exception {
	final int numberSlots = 3;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));
		assertThat(slotManager.getNumberRegisteredSlots(), is(0));

		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final SlotReport slotReport = createSlotReport(taskExecutorConnection.getResourceID(), numberSlots - 1);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(numberSlots - 1));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(1));
	}
}
 

MockResourceManagerRuntimeServices() throws Exception {
	highAvailabilityServices = new TestingHighAvailabilityServices();
	rmLeaderElectionService = new TestingLeaderElectionService();
	highAvailabilityServices.setResourceManagerLeaderElectionService(rmLeaderElectionService);
	heartbeatServices = new HeartbeatServices(5L, 5L);
	metricRegistry = mock(MetricRegistryImpl.class);
	slotManager = mock(SlotManager.class);
	slotManagerStarted = new CompletableFuture<>();
	jobLeaderIdService = new JobLeaderIdService(
		highAvailabilityServices,
		rpcService.getScheduledExecutor(),
		Time.minutes(5L));

	doAnswer(new Answer<Object>() {
		@Override
		public Object answer(InvocationOnMock invocation) throws Throwable {
			rmActions = invocation.getArgument(2);
			slotManagerStarted.complete(true);
			return null;
		}
	}).when(slotManager).start(any(ResourceManagerId.class), any(Executor.class), any(ResourceActions.class));

	when(slotManager.registerSlotRequest(any(SlotRequest.class))).thenReturn(true);
}
 

/**
 * Starts the slot manager with the given leader id and resource manager actions.
 *
 * @param newResourceManagerId to use for communication with the task managers
 * @param newMainThreadExecutor to use to run code in the ResourceManager's main thread
 * @param newResourceActions to use for resource (de-)allocations
 */
public void start(ResourceManagerId newResourceManagerId, Executor newMainThreadExecutor, ResourceActions newResourceActions) {
	LOG.info("Starting the SlotManager.");

	this.resourceManagerId = Preconditions.checkNotNull(newResourceManagerId);
	mainThreadExecutor = Preconditions.checkNotNull(newMainThreadExecutor);
	resourceActions = Preconditions.checkNotNull(newResourceActions);

	started = true;

	taskManagerTimeoutCheck = scheduledExecutor.scheduleWithFixedDelay(
		() -> mainThreadExecutor.execute(
			() -> checkTaskManagerTimeouts()),
		0L,
		taskManagerTimeout.toMilliseconds(),
		TimeUnit.MILLISECONDS);

	slotRequestTimeoutCheck = scheduledExecutor.scheduleWithFixedDelay(
		() -> mainThreadExecutor.execute(
			() -> checkSlotRequestTimeouts()),
		0L,
		slotRequestTimeout.toMilliseconds(),
		TimeUnit.MILLISECONDS);
}
 

/**
 * Tests that only free slots can fulfill/complete a pending task manager slot.
 */
@Test
public void testOnlyFreeSlotsCanFulfillPendingTaskManagerSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManager slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));

		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final SlotID slotId = new SlotID(taskExecutorConnection.getResourceID(), 0);
		final SlotStatus slotStatus = new SlotStatus(slotId, ResourceProfile.UNKNOWN, jobId, new AllocationID());
		final SlotReport slotReport = new SlotReport(slotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(1));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));
	}
}
 

/**
 * Tests that a failing allocation/slot request will return the pending task manager slot.
 */
@Test
public void testFailingAllocationReturnsPendingTaskManagerSlot() throws Exception {
	final int numberSlots = 2;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();
	try (final SlotManager slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();

		final SlotRequest slotRequest = createSlotRequest(jobId);
		assertThat(slotManager.registerSlotRequest(slotRequest), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(1));

		slotManager.unregisterSlotRequest(slotRequest.getAllocationId());

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(0));
	}
}
 

@Override
public void notifyLeaderAddress(final String leaderAddress, final UUID leaderSessionID) {
	runAsync(
		() -> notifyOfNewResourceManagerLeader(
			leaderAddress,
			ResourceManagerId.fromUuidOrNull(leaderSessionID)));
}
 

@Override
public void notifyLeaderAddress(final String leaderAddress, final UUID leaderSessionID) {
	runAsync(
		() -> notifyOfNewResourceManagerLeader(
			leaderAddress,
			ResourceManagerId.fromUuidOrNull(leaderSessionID)));
}
 

ResourceManagerConnection(
		final Logger log,
		final JobID jobID,
		final ResourceID jobManagerResourceID,
		final String jobManagerRpcAddress,
		final JobMasterId jobMasterId,
		final String resourceManagerAddress,
		final ResourceManagerId resourceManagerId,
		final Executor executor) {
	super(log, resourceManagerAddress, resourceManagerId, executor);
	this.jobID = checkNotNull(jobID);
	this.jobManagerResourceID = checkNotNull(jobManagerResourceID);
	this.jobManagerRpcAddress = checkNotNull(jobManagerRpcAddress);
	this.jobMasterId = checkNotNull(jobMasterId);
}
 

@Override
public void disconnectResourceManager(
		final ResourceManagerId resourceManagerId,
		final Exception cause) {

	if (isConnectingToResourceManager(resourceManagerId)) {
		reconnectToResourceManager(cause);
	}
}
 

@Nullable
private ResourceManagerAddress createResourceManagerAddress(@Nullable String newResourceManagerAddress, @Nullable ResourceManagerId resourceManagerId) {
	if (newResourceManagerAddress != null) {
		// the contract is: address == null <=> id == null
		checkNotNull(resourceManagerId);
		return new ResourceManagerAddress(newResourceManagerAddress, resourceManagerId);
	} else {
		return null;
	}
}
 

@Override
protected CompletableFuture<RegistrationResponse> invokeRegistration(
		ResourceManagerGateway resourceManager, ResourceManagerId fencingToken, long timeoutMillis) throws Exception {

	Time timeout = Time.milliseconds(timeoutMillis);
	return resourceManager.registerTaskExecutor(
		taskExecutorAddress,
		resourceID,
		dataPort,
		hardwareDescription,
		timeout);
}
 

/**
 * Checks that un-registering a pending slot request will cancel it, removing it from all
 * assigned task manager slots and then remove it from the slot manager.
 */
@Test
public void testUnregisterPendingSlotRequest() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple5 -> new CompletableFuture<>())
		.createTestingTaskExecutorGateway();

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

	final SlotRequest slotRequest = new SlotRequest(new JobID(), allocationId, resourceProfile, "foobar");

	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	try (SlotManagerImpl slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		slotManager.registerSlotRequest(slotRequest);

		assertNotNull(slotManager.getSlotRequest(allocationId));

		assertTrue(slot.getState() == TaskManagerSlot.State.PENDING);

		slotManager.unregisterSlotRequest(allocationId);

		assertNull(slotManager.getSlotRequest(allocationId));

		slot = slotManager.getSlot(slotId);
		assertTrue(slot.getState() == TaskManagerSlot.State.FREE);
	}
}
 

/**
 * Tests that a different slot can fulfill a pending slot request. If the
 * pending slot request has a pending task manager slot assigned, it should
 * be freed.
 */
@Test
public void testRegistrationOfDifferentSlot() throws Exception {
	final int numberSlots = 1;
	final TestingResourceActions resourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(convert(value -> numberSlots))
		.build();

	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(), resourceActions)) {
		final JobID jobId = new JobID();
		final ResourceProfile requestedSlotProfile = new ResourceProfile(1.0, 1);

		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId, requestedSlotProfile)), is(true));

		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));

		final int numberOfferedSlots = 1;
		final TaskExecutorConnection taskExecutorConnection = createTaskExecutorConnection();
		final ResourceProfile offeredSlotProfile = new ResourceProfile(2.0, 2);
		final SlotReport slotReport = createSlotReport(taskExecutorConnection.getResourceID(), numberOfferedSlots, offeredSlotProfile);

		slotManager.registerTaskManager(taskExecutorConnection, slotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(numberOfferedSlots));
		assertThat(slotManager.getNumberPendingTaskManagerSlots(), is(numberSlots));
		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(0));
	}
}
 

@Override
public void notifyLeaderAddress(final String leaderAddress, final UUID leaderSessionID) {
	runAsync(
		() -> notifyOfNewResourceManagerLeader(
			leaderAddress,
			ResourceManagerId.fromUuidOrNull(leaderSessionID)));
}
 

/**
 * Checks that un-registering a pending slot request will cancel it, removing it from all
 * assigned task manager slots and then remove it from the slot manager.
 */
@Test
public void testUnregisterPendingSlotRequest() throws Exception {
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceActions resourceManagerActions = mock(ResourceActions.class);
	final ResourceID resourceID = ResourceID.generate();
	final SlotID slotId = new SlotID(resourceID, 0);
	final AllocationID allocationId = new AllocationID();

	final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple5 -> new CompletableFuture<>())
		.createTestingTaskExecutorGateway();

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

	final SlotRequest slotRequest = new SlotRequest(new JobID(), allocationId, resourceProfile, "foobar");

	final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);

	try (SlotManager slotManager = createSlotManager(resourceManagerId, resourceManagerActions)) {
		slotManager.registerTaskManager(taskManagerConnection, slotReport);

		TaskManagerSlot slot = slotManager.getSlot(slotId);

		slotManager.registerSlotRequest(slotRequest);

		assertNotNull(slotManager.getSlotRequest(allocationId));

		assertTrue(slot.getState() == TaskManagerSlot.State.PENDING);

		slotManager.unregisterSlotRequest(allocationId);

		assertNull(slotManager.getSlotRequest(allocationId));

		slot = slotManager.getSlot(slotId);
		assertTrue(slot.getState() == TaskManagerSlot.State.FREE);
	}
}
 

/**
 * Tests that we continue reconnecting to the latest known RM after a disconnection
 * message.
 */
@Test
public void testReconnectionAfterDisconnect() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
		final BlockingQueue<JobMasterId> registrationsQueue = new ArrayBlockingQueue<>(1);

		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationsQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		final ResourceManagerId resourceManagerId = testingResourceManagerGateway.getFencingToken();
		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		// wait for first registration attempt
		final JobMasterId firstRegistrationAttempt = registrationsQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		assertThat(registrationsQueue.isEmpty(), is(true));
		jobMasterGateway.disconnectResourceManager(resourceManagerId, new FlinkException("Test exception"));

		// wait for the second registration attempt after the disconnect call
		assertThat(registrationsQueue.take(), equalTo(jobMasterId));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

/**
 * Tests that we only request new resources/containers once we have assigned
 * all pending task manager slots.
 */
@Test
public void testRequestNewResources() throws Exception {
	final int numberSlots = 2;
	final AtomicInteger resourceRequests = new AtomicInteger(0);
	final TestingResourceActions testingResourceActions = new TestingResourceActionsBuilder()
		.setAllocateResourceFunction(
			convert(ignored -> {
				resourceRequests.incrementAndGet();
				return numberSlots;
			}))
		.build();

	try (final SlotManagerImpl slotManager = createSlotManager(
		ResourceManagerId.generate(),
		testingResourceActions)) {

		final JobID jobId = new JobID();
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));
		assertThat(resourceRequests.get(), is(1));

		// the second slot request should not try to allocate a new resource because the
		// previous resource was started with 2 slots.
		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));
		assertThat(resourceRequests.get(), is(1));

		assertThat(slotManager.getNumberAssignedPendingTaskManagerSlots(), is(2));

		assertThat(slotManager.registerSlotRequest(createSlotRequest(jobId)), is(true));
		assertThat(resourceRequests.get(), is(2));
	}
}
 

/**
 * Tests that we continue reconnecting to the latest known RM after a disconnection
 * message.
 */
@Test
public void testReconnectionAfterDisconnect() throws Exception {
	final JobMaster jobMaster = createJobMaster(
		configuration,
		jobGraph,
		haServices,
		new TestingJobManagerSharedServicesBuilder().build());

	final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);

	CompletableFuture<Acknowledge> startFuture = jobMaster.start(jobMasterId);

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);
		final TestingResourceManagerGateway testingResourceManagerGateway = createAndRegisterTestingResourceManagerGateway();
		final BlockingQueue<JobMasterId> registrationsQueue = new ArrayBlockingQueue<>(1);

		testingResourceManagerGateway.setRegisterJobManagerConsumer(
			jobMasterIdResourceIDStringJobIDTuple4 -> registrationsQueue.offer(jobMasterIdResourceIDStringJobIDTuple4.f0));

		final ResourceManagerId resourceManagerId = testingResourceManagerGateway.getFencingToken();
		notifyResourceManagerLeaderListeners(testingResourceManagerGateway);

		// wait for first registration attempt
		final JobMasterId firstRegistrationAttempt = registrationsQueue.take();

		assertThat(firstRegistrationAttempt, equalTo(jobMasterId));

		assertThat(registrationsQueue.isEmpty(), is(true));
		jobMasterGateway.disconnectResourceManager(resourceManagerId, new FlinkException("Test exception"));

		// wait for the second registration attempt after the disconnect call
		assertThat(registrationsQueue.take(), equalTo(jobMasterId));
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
	}
}
 

public TestingResourceManagerGateway(
		ResourceManagerId resourceManagerId,
		ResourceID resourceId,
		String address,
		String hostname) {
	this.resourceManagerId = Preconditions.checkNotNull(resourceManagerId);
	this.ownResourceId = Preconditions.checkNotNull(resourceId);
	this.address = Preconditions.checkNotNull(address);
	this.hostname = Preconditions.checkNotNull(hostname);
	this.slotFutureReference = new AtomicReference<>();
	this.cancelSlotConsumer = null;
	this.requestSlotConsumer = null;
}
 

public TestingResourceManagerGateway() {
	this(
		ResourceManagerId.generate(),
		ResourceID.generate(),
		"localhost/" + UUID.randomUUID(),
		"localhost");
}
 

@Nullable
private ResourceManagerAddress createResourceManagerAddress(@Nullable String newLeaderAddress, @Nullable ResourceManagerId newResourceManagerId) {
	if (newLeaderAddress == null) {
		return null;
	} else {
		assert(newResourceManagerId != null);
		return new ResourceManagerAddress(newLeaderAddress, newResourceManagerId);
	}
}
 

/**
 * Tests that a task manager timeout does not remove the slots from the SlotManager.
 * A timeout should only trigger the {@link ResourceActions#releaseResource(InstanceID, Exception)}
 * callback. The receiver of the callback can then decide what to do with the TaskManager.
 *
 * <p>See FLINK-7793
 */
@Test
public void testTaskManagerTimeoutDoesNotRemoveSlots() throws Exception {
	final Time taskManagerTimeout = Time.milliseconds(10L);
	final ResourceManagerId resourceManagerId = ResourceManagerId.generate();
	final ResourceID resourceID = ResourceID.generate();
	final ResourceActions resourceActions = mock(ResourceActions.class);
	final TaskExecutorGateway taskExecutorGateway = mock(TaskExecutorGateway.class);

	when(taskExecutorGateway.canBeReleased()).thenReturn(CompletableFuture.completedFuture(true));

	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(resourceID, taskExecutorGateway);
	final SlotStatus slotStatus = new SlotStatus(
		new SlotID(resourceID, 0),
		new ResourceProfile(1.0, 1));
	final SlotReport initialSlotReport = new SlotReport(slotStatus);

	try (final SlotManager slotManager = SlotManagerBuilder.newBuilder()
		.setTaskManagerTimeout(taskManagerTimeout)
		.build()) {

		slotManager.start(resourceManagerId, Executors.directExecutor(), resourceActions);

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		// wait for the timeout call to happen
		verify(resourceActions, timeout(taskManagerTimeout.toMilliseconds() * 20L).atLeast(1)).releaseResource(eq(taskExecutorConnection.getInstanceID()), any(Exception.class));

		assertEquals(1, slotManager.getNumberRegisteredSlots());

		slotManager.unregisterTaskManager(taskExecutorConnection.getInstanceID());

		assertEquals(0, slotManager.getNumberRegisteredSlots());
	}
}