org.apache.flink.runtime.jobmaster.AllocatedSlotReport Java Examples

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

private void failNoLongerAllocatedSlots(AllocatedSlotReport allocatedSlotReport, JobMasterGateway jobMasterGateway) {
	for (AllocatedSlotInfo allocatedSlotInfo : allocatedSlotReport.getAllocatedSlotInfos()) {
		final AllocationID allocationId = allocatedSlotInfo.getAllocationId();
		if (!taskSlotTable.isAllocated(
				allocatedSlotInfo.getSlotIndex(),
				allocatedSlotReport.getJobId(),
				allocationId)) {
			jobMasterGateway.failSlot(
					getResourceID(),
					allocationId,
					new FlinkException(
						String.format(
							"Slot %s on TaskExecutor %s is not allocated by job %s.",
							allocatedSlotInfo.getSlotIndex(),
							getResourceID(),
							allocatedSlotReport.getJobId())));
		}
	}
}
 

private void failNoLongerAllocatedSlots(AllocatedSlotReport allocatedSlotReport, JobMasterGateway jobMasterGateway) {
	for (AllocatedSlotInfo allocatedSlotInfo : allocatedSlotReport.getAllocatedSlotInfos()) {
		final AllocationID allocationId = allocatedSlotInfo.getAllocationId();
		if (!taskSlotTable.isAllocated(
				allocatedSlotInfo.getSlotIndex(),
				allocatedSlotReport.getJobId(),
				allocationId)) {
			jobMasterGateway.failSlot(
					getResourceID(),
					allocationId,
					new FlinkException(
						String.format(
							"Slot %s on TaskExecutor %s is not allocated by job %s.",
							allocatedSlotInfo.getSlotIndex(),
							getResourceID(),
							allocatedSlotReport.getJobId())));
		}
	}
}
 

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<Boolean> canBeReleasedSupplier) {
	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;
}
 

private void failNoLongerAllocatedSlots(AllocatedSlotReport allocatedSlotReport, JobMasterGateway jobMasterGateway) {
	for (AllocatedSlotInfo allocatedSlotInfo : allocatedSlotReport.getAllocatedSlotInfos()) {
		final AllocationID allocationId = allocatedSlotInfo.getAllocationId();
		if (!taskSlotTable.isAllocated(
				allocatedSlotInfo.getSlotIndex(),
				allocatedSlotReport.getJobId(),
				allocationId)) {
			jobMasterGateway.failSlot(
					getResourceID(),
					allocationId,
					new FlinkException(
						String.format(
							"Slot %s on TaskExecutor %s is not allocated by job %s.",
							allocatedSlotInfo.getSlotIndex(),
							getResourceID(),
							allocatedSlotReport.getJobId())));
		}
	}
}
 

@Override
public AllocatedSlotReport createAllocatedSlotReport(ResourceID taskManagerId) {
	final Set<AllocatedSlot> availableSlotsForTaskManager = availableSlots.getSlotsForTaskManager(taskManagerId);
	final Set<AllocatedSlot> allocatedSlotsForTaskManager = allocatedSlots.getSlotsForTaskManager(taskManagerId);

	List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(
			availableSlotsForTaskManager.size() + allocatedSlotsForTaskManager.size());
	for (AllocatedSlot allocatedSlot : Iterables.concat(availableSlotsForTaskManager, allocatedSlotsForTaskManager)) {
		allocatedSlotInfos.add(
				new AllocatedSlotInfo(allocatedSlot.getPhysicalSlotNumber(), allocatedSlot.getAllocationId()));
	}
	return new AllocatedSlotReport(jobId, allocatedSlotInfos);
}
 

/**
 * Tests that create report of allocated slots on a {@link TaskExecutor}.
 */
@Test
public void testCreateAllocatedSlotReport() throws Exception {

	try (SlotPoolImpl slotPool = createSlotPoolImpl()) {

		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);
		resourceManagerGateway.setRequestSlotConsumer(
				slotRequest -> allocationIds.offer(slotRequest.getAllocationId()));

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final SlotRequestId slotRequestId = new SlotRequestId();
		final CompletableFuture<LogicalSlot> slotRequestFuture = allocateSlot(scheduler, slotRequestId);

		final List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(2);
		final List<SlotOffer> slotOffers = new ArrayList<>(2);

		final AllocationID allocatedId = allocationIds.take();
		slotOffers.add(new SlotOffer(allocatedId, 0, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(0, allocatedId));

		final AllocationID availableId = new AllocationID();
		slotOffers.add(new SlotOffer(availableId, 1, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(1, availableId));

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());
		slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

		// wait for the completion of slot future
		slotRequestFuture.get();

		final AllocatedSlotReport slotReport = slotPool.createAllocatedSlotReport(taskManagerLocation.getResourceID());
		assertThat(jobId, is(slotReport.getJobId()));
		assertThat(slotReport.getAllocatedSlotInfos(), containsInAnyOrder(isEachEqual(allocatedSlotInfos)));
	}
}
 

@Override
public AllocatedSlotReport createAllocatedSlotReport(ResourceID taskManagerId) {
	final Set<AllocatedSlot> availableSlotsForTaskManager = availableSlots.getSlotsForTaskManager(taskManagerId);
	final Set<AllocatedSlot> allocatedSlotsForTaskManager = allocatedSlots.getSlotsForTaskManager(taskManagerId);

	List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(
			availableSlotsForTaskManager.size() + allocatedSlotsForTaskManager.size());
	for (AllocatedSlot allocatedSlot : Iterables.concat(availableSlotsForTaskManager, allocatedSlotsForTaskManager)) {
		allocatedSlotInfos.add(
				new AllocatedSlotInfo(allocatedSlot.getPhysicalSlotNumber(), allocatedSlot.getAllocationId()));
	}
	return new AllocatedSlotReport(jobId, allocatedSlotInfos);
}
 

private HeartbeatManager<AllocatedSlotReport, AccumulatorReport> createJobManagerHeartbeatManager(HeartbeatServices heartbeatServices, ResourceID resourceId) {
	return heartbeatServices.createHeartbeatManager(
		resourceId,
		new JobManagerHeartbeatListener(),
		getMainThreadExecutor(),
		log);
}
 

private void freeNoLongerUsedSlots(AllocatedSlotReport allocatedSlotReport) {
	final Iterator<AllocationID> slotsTaskManagerSide = taskSlotTable.getActiveSlots(allocatedSlotReport.getJobId());
	final Set<AllocationID> activeSlots = Sets.newHashSet(slotsTaskManagerSide);
	final Set<AllocationID> reportedSlots = allocatedSlotReport.getAllocatedSlotInfos().stream()
			.map(AllocatedSlotInfo::getAllocationId).collect(Collectors.toSet());

	final Sets.SetView<AllocationID> difference = Sets.difference(activeSlots, reportedSlots);

	for (AllocationID allocationID : difference) {
		freeSlotInternal(
			allocationID,
			new FlinkException(
				String.format("%s is no longer allocated by job %s.", allocationID, allocatedSlotReport.getJobId())));
	}
}
 

/**
 * Syncs the TaskExecutor's view on its allocated slots with the JobMaster's view.
 * Slots which are no longer reported by the JobMaster are being freed.
 * Slots which the JobMaster thinks it still owns but which are no longer allocated to it
 * will be failed via {@link JobMasterGateway#failSlot}.
 *
 * @param allocatedSlotReport represents the JobMaster's view on the current slot allocation state
 */
private void syncSlotsWithSnapshotFromJobMaster(AllocatedSlotReport allocatedSlotReport) {
	JobManagerConnection jobManagerConnection = jobManagerTable.get(allocatedSlotReport.getJobId());
	if (jobManagerConnection != null) {
		final JobMasterGateway jobMasterGateway = jobManagerConnection.getJobManagerGateway();

		failNoLongerAllocatedSlots(allocatedSlotReport, jobMasterGateway);

		freeNoLongerUsedSlots(allocatedSlotReport);
	} else {
		log.debug("Ignoring allocated slot report from job {} because there is no active leader.",
				allocatedSlotReport.getJobId());
	}
}
 

@Override
public AllocatedSlotReport createAllocatedSlotReport(ResourceID taskManagerId) {
	final Set<AllocatedSlot> availableSlotsForTaskManager = availableSlots.getSlotsForTaskManager(taskManagerId);
	final Set<AllocatedSlot> allocatedSlotsForTaskManager = allocatedSlots.getSlotsForTaskManager(taskManagerId);

	List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(
			availableSlotsForTaskManager.size() + allocatedSlotsForTaskManager.size());
	for (AllocatedSlot allocatedSlot : Iterables.concat(availableSlotsForTaskManager, allocatedSlotsForTaskManager)) {
		allocatedSlotInfos.add(
				new AllocatedSlotInfo(allocatedSlot.getPhysicalSlotNumber(), allocatedSlot.getAllocationId()));
	}
	return new AllocatedSlotReport(jobId, allocatedSlotInfos);
}
 

private void freeNoLongerUsedSlots(AllocatedSlotReport allocatedSlotReport) {
	final Iterator<AllocationID> slotsTaskManagerSide = taskSlotTable.getActiveSlots(allocatedSlotReport.getJobId());
	final Set<AllocationID> activeSlots = Sets.newHashSet(slotsTaskManagerSide);
	final Set<AllocationID> reportedSlots = allocatedSlotReport.getAllocatedSlotInfos().stream()
			.map(AllocatedSlotInfo::getAllocationId).collect(Collectors.toSet());

	final Sets.SetView<AllocationID> difference = Sets.difference(activeSlots, reportedSlots);

	for (AllocationID allocationID : difference) {
		freeSlotInternal(
			allocationID,
			new FlinkException(
				String.format("%s is no longer allocated by job %s.", allocationID, allocatedSlotReport.getJobId())));
	}
}
 

@Override
public void reportPayload(ResourceID resourceID, AllocatedSlotReport allocatedSlotReport) {
	validateRunsInMainThread();
	OptionalConsumer.of(jobTable.getConnection(allocatedSlotReport.getJobId()))
		.ifPresent(
			jobManagerConnection -> {
				syncSlotsWithSnapshotFromJobMaster(jobManagerConnection.getJobManagerGateway(), allocatedSlotReport);
			})
		.ifNotPresent(() -> log.debug("Ignoring allocated slot report from job {} because there is no active leader.", allocatedSlotReport.getJobId()));

}
 

/**
 * Tests that create report of allocated slots on a {@link TaskExecutor}.
 */
@Test
public void testCreateAllocatedSlotReport() throws Exception {

	try (SlotPoolImpl slotPool = createSlotPoolImpl()) {

		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);
		resourceManagerGateway.setRequestSlotConsumer(
				slotRequest -> allocationIds.offer(slotRequest.getAllocationId()));

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final SlotRequestId slotRequestId = new SlotRequestId();
		final CompletableFuture<LogicalSlot> slotRequestFuture = allocateSlot(scheduler, slotRequestId);

		final List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(2);
		final List<SlotOffer> slotOffers = new ArrayList<>(2);

		final AllocationID allocatedId = allocationIds.take();
		slotOffers.add(new SlotOffer(allocatedId, 0, ResourceProfile.ANY));
		allocatedSlotInfos.add(new AllocatedSlotInfo(0, allocatedId));

		final AllocationID availableId = new AllocationID();
		slotOffers.add(new SlotOffer(availableId, 1, ResourceProfile.ANY));
		allocatedSlotInfos.add(new AllocatedSlotInfo(1, availableId));

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());
		slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

		// wait for the completion of slot future
		slotRequestFuture.get();

		final AllocatedSlotReport slotReport = slotPool.createAllocatedSlotReport(taskManagerLocation.getResourceID());
		assertThat(jobId, is(slotReport.getJobId()));
		assertThat(slotReport.getAllocatedSlotInfos(), containsInAnyOrder(isEachEqual(allocatedSlotInfos)));
	}
}
 

/**
 * Tests that create report of allocated slots on a {@link TaskExecutor}.
 */
@Test
public void testCreateAllocatedSlotReport() throws Exception {

	try (SlotPoolImpl slotPool = new SlotPoolImpl(jobId)) {

		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);
		resourceManagerGateway.setRequestSlotConsumer(
				slotRequest -> allocationIds.offer(slotRequest.getAllocationId()));

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final SlotRequestId slotRequestId = new SlotRequestId();
		final CompletableFuture<LogicalSlot> slotRequestFuture = allocateSlot(scheduler, slotRequestId);

		final List<AllocatedSlotInfo> allocatedSlotInfos = new ArrayList<>(2);
		final List<SlotOffer> slotOffers = new ArrayList<>(2);

		final AllocationID allocatedId = allocationIds.take();
		slotOffers.add(new SlotOffer(allocatedId, 0, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(0, allocatedId));

		final AllocationID availableId = new AllocationID();
		slotOffers.add(new SlotOffer(availableId, 1, ResourceProfile.UNKNOWN));
		allocatedSlotInfos.add(new AllocatedSlotInfo(1, availableId));

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());
		slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers);

		// wait for the completion of slot future
		slotRequestFuture.get();

		final AllocatedSlotReport slotReport = slotPool.createAllocatedSlotReport(taskManagerLocation.getResourceID());
		assertThat(jobId, is(slotReport.getJobId()));
		assertThat(slotReport.getAllocatedSlotInfos(), containsInAnyOrder(isEachEqual(allocatedSlotInfos)));
	}
}
 

TestingTaskExecutorGateway(
		String address,
		String hostname,
		BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer,
		BiConsumer<JobID, Throwable> disconnectJobManagerConsumer,
		BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer,
		Function<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, 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,
		TriConsumer<JobID, Set<ResultPartitionID>, Set<ResultPartitionID>> releaseOrPromotePartitionsConsumer,
		Consumer<Collection<IntermediateDataSetID>> releaseClusterPartitionsConsumer,
		TriFunction<ExecutionAttemptID, OperatorID, SerializedValue<OperatorEvent>, CompletableFuture<Acknowledge>> operatorEventHandler,
		Supplier<CompletableFuture<ThreadDumpInfo>> requestThreadDumpSupplier) {

	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.releaseOrPromotePartitionsConsumer = releaseOrPromotePartitionsConsumer;
	this.releaseClusterPartitionsConsumer = releaseClusterPartitionsConsumer;
	this.operatorEventHandler = operatorEventHandler;
	this.requestThreadDumpSupplier = requestThreadDumpSupplier;
}
 

private void freeNoLongerUsedSlots(AllocatedSlotReport allocatedSlotReport) {
	final Iterator<AllocationID> slotsTaskManagerSide = taskSlotTable.getActiveSlots(allocatedSlotReport.getJobId());
	final Set<AllocationID> activeSlots = Sets.newHashSet(slotsTaskManagerSide);
	final Set<AllocationID> reportedSlots = allocatedSlotReport.getAllocatedSlotInfos().stream()
			.map(AllocatedSlotInfo::getAllocationId).collect(Collectors.toSet());

	final Sets.SetView<AllocationID> difference = Sets.difference(activeSlots, reportedSlots);

	for (AllocationID allocationID : difference) {
		freeSlotInternal(
			allocationID,
			new FlinkException(
				String.format("%s is no longer allocated by job %s.", allocationID, allocatedSlotReport.getJobId())));
	}
}
 

/**
 * Syncs the TaskExecutor's view on its allocated slots with the JobMaster's view.
 * Slots which are no longer reported by the JobMaster are being freed.
 * Slots which the JobMaster thinks it still owns but which are no longer allocated to it
 * will be failed via {@link JobMasterGateway#failSlot}.
 *
 * @param allocatedSlotReport represents the JobMaster's view on the current slot allocation state
 */
private void syncSlotsWithSnapshotFromJobMaster(AllocatedSlotReport allocatedSlotReport) {
	JobManagerConnection jobManagerConnection = jobManagerTable.get(allocatedSlotReport.getJobId());
	if (jobManagerConnection != null) {
		final JobMasterGateway jobMasterGateway = jobManagerConnection.getJobManagerGateway();

		failNoLongerAllocatedSlots(allocatedSlotReport, jobMasterGateway);

		freeNoLongerUsedSlots(allocatedSlotReport);
	} else {
		log.debug("Ignoring allocated slot report from job {} because there is no active leader.",
				allocatedSlotReport.getJobId());
	}
}
 

@Override
public void heartbeatFromJobManager(ResourceID resourceID, AllocatedSlotReport allocatedSlotReport) {
	jobManagerHeartbeatManager.requestHeartbeat(resourceID, allocatedSlotReport);
}
 

@Override
public void heartbeatFromJobManager(ResourceID heartbeatOrigin, AllocatedSlotReport allocatedSlotReport) {
	heartbeatJobManagerConsumer.accept(heartbeatOrigin, allocatedSlotReport);
}
 

public TestingTaskExecutorGatewayBuilder setHeartbeatJobManagerConsumer(BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer) {
	this.heartbeatJobManagerConsumer = heartbeatJobManagerConsumer;
	return this;
}
 

/**
 * Tests that the TaskExecutor syncs its slots view with the JobMaster's view
 * via the AllocatedSlotReport reported by the heartbeat (See FLINK-11059).
 */
@Test
public void testSyncSlotsWithJobMasterByHeartbeat() throws Exception {
	final CountDownLatch activeSlots = new CountDownLatch(2);
	final TaskSlotTable<Task> taskSlotTable = new ActivateSlotNotifyingTaskSlotTable(
			2,
			activeSlots);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

	final BlockingQueue<AllocationID> allocationsNotifiedFree = new ArrayBlockingQueue<>(2);

	OneShotLatch initialSlotReporting = new OneShotLatch();
	testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReporting.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());

	});

	testingResourceManagerGateway.setNotifySlotAvailableConsumer(instanceIDSlotIDAllocationIDTuple3 ->
			allocationsNotifiedFree.offer(instanceIDSlotIDAllocationIDTuple3.f2));

	rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
	resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

	final BlockingQueue<AllocationID> failedSlotFutures = new ArrayBlockingQueue<>(2);
	final ResourceID jobManagerResourceId = ResourceID.generate();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
			.setFailSlotConsumer((resourceID, allocationID, throwable) ->
				failedSlotFutures.offer(allocationID))
			.setOfferSlotsFunction((resourceID, slotOffers) -> CompletableFuture.completedFuture(new ArrayList<>(slotOffers)))
			.setRegisterTaskManagerFunction((ignoredA, ignoredB) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jobManagerResourceId)))
			.build();
	final String jobManagerAddress = jobMasterGateway.getAddress();
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());

	taskExecutor.start();

	try {
		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		initialSlotReporting.await();

		final SlotID slotId1 = new SlotID(taskExecutor.getResourceID(), 0);
		final SlotID slotId2 = new SlotID(taskExecutor.getResourceID(), 1);
		final AllocationID allocationIdInBoth = new AllocationID();
		final AllocationID allocationIdOnlyInJM = new AllocationID();
		final AllocationID allocationIdOnlyInTM = new AllocationID();

		taskExecutorGateway.requestSlot(slotId1, jobId, allocationIdInBoth, ResourceProfile.ZERO, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
		taskExecutorGateway.requestSlot(slotId2, jobId, allocationIdOnlyInTM, ResourceProfile.ZERO, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		activeSlots.await();

		List<AllocatedSlotInfo> allocatedSlotInfos = Arrays.asList(
				new AllocatedSlotInfo(0, allocationIdInBoth),
				new AllocatedSlotInfo(1, allocationIdOnlyInJM)
		);
		AllocatedSlotReport allocatedSlotReport = new AllocatedSlotReport(jobId, allocatedSlotInfos);
		taskExecutorGateway.heartbeatFromJobManager(jobManagerResourceId, allocatedSlotReport);

		assertThat(failedSlotFutures.take(), is(allocationIdOnlyInJM));
		assertThat(allocationsNotifiedFree.take(), is(allocationIdOnlyInTM));
		assertThat(failedSlotFutures.poll(5L, TimeUnit.MILLISECONDS), nullValue());
		assertThat(allocationsNotifiedFree.poll(5L, TimeUnit.MILLISECONDS), nullValue());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

/**
 * Tests that the TaskExecutor syncs its slots view with the JobMaster's view
 * via the AllocatedSlotReport reported by the heartbeat (See FLINK-11059).
 */
@Test
public void testSyncSlotsWithJobMasterByHeartbeat() throws Exception {
	final CountDownLatch activeSlots = new CountDownLatch(2);
	final TaskSlotTable taskSlotTable = new ActivateSlotNotifyingTaskSlotTable(
			Arrays.asList(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN),
			timerService,
			activeSlots);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

	final BlockingQueue<AllocationID> allocationsNotifiedFree = new ArrayBlockingQueue<>(2);

	OneShotLatch initialSlotReporting = new OneShotLatch();
	testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReporting.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());

	});

	testingResourceManagerGateway.setNotifySlotAvailableConsumer(instanceIDSlotIDAllocationIDTuple3 ->
			allocationsNotifiedFree.offer(instanceIDSlotIDAllocationIDTuple3.f2));

	rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
	resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

	final BlockingQueue<AllocationID> failedSlotFutures = new ArrayBlockingQueue<>(2);
	final ResourceID jobManagerResourceId = ResourceID.generate();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
			.setFailSlotConsumer((resourceID, allocationID, throwable) ->
				failedSlotFutures.offer(allocationID))
			.setOfferSlotsFunction((resourceID, slotOffers) -> CompletableFuture.completedFuture(new ArrayList<>(slotOffers)))
			.setRegisterTaskManagerFunction((ignoredA, ignoredB) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jobManagerResourceId)))
			.build();
	final String jobManagerAddress = jobMasterGateway.getAddress();
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());

	taskExecutor.start();

	try {
		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		initialSlotReporting.await();

		final SlotID slotId1 = new SlotID(taskExecutor.getResourceID(), 0);
		final SlotID slotId2 = new SlotID(taskExecutor.getResourceID(), 1);
		final AllocationID allocationIdInBoth = new AllocationID();
		final AllocationID allocationIdOnlyInJM = new AllocationID();
		final AllocationID allocationIdOnlyInTM = new AllocationID();

		taskExecutorGateway.requestSlot(slotId1, jobId, allocationIdInBoth, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
		taskExecutorGateway.requestSlot(slotId2, jobId, allocationIdOnlyInTM, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		activeSlots.await();

		List<AllocatedSlotInfo> allocatedSlotInfos = Arrays.asList(
				new AllocatedSlotInfo(0, allocationIdInBoth),
				new AllocatedSlotInfo(1, allocationIdOnlyInJM)
		);
		AllocatedSlotReport allocatedSlotReport = new AllocatedSlotReport(jobId, allocatedSlotInfos);
		taskExecutorGateway.heartbeatFromJobManager(jobManagerResourceId, allocatedSlotReport);

		assertThat(failedSlotFutures.take(), is(allocationIdOnlyInJM));
		assertThat(allocationsNotifiedFree.take(), is(allocationIdOnlyInTM));
		assertThat(failedSlotFutures.poll(5L, TimeUnit.MILLISECONDS), nullValue());
		assertThat(allocationsNotifiedFree.poll(5L, TimeUnit.MILLISECONDS), nullValue());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

public TestingTaskExecutorGatewayBuilder setHeartbeatJobManagerConsumer(BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer) {
	this.heartbeatJobManagerConsumer = heartbeatJobManagerConsumer;
	return this;
}
 

@Override
public void heartbeatFromJobManager(ResourceID heartbeatOrigin, AllocatedSlotReport allocatedSlotReport) {
	heartbeatJobManagerConsumer.accept(heartbeatOrigin, allocatedSlotReport);
}
 

@Override
public void reportPayload(ResourceID resourceID, AllocatedSlotReport allocatedSlotReport) {
	validateRunsInMainThread();
	syncSlotsWithSnapshotFromJobMaster(allocatedSlotReport);
}
 

@Override
public void heartbeatFromJobManager(ResourceID resourceID, AllocatedSlotReport allocatedSlotReport) {
	jobManagerHeartbeatManager.requestHeartbeat(resourceID, allocatedSlotReport);
}
 

/**
 * Tests that the TaskExecutor syncs its slots view with the JobMaster's view
 * via the AllocatedSlotReport reported by the heartbeat (See FLINK-11059).
 */
@Test
public void testSyncSlotsWithJobMasterByHeartbeat() throws Exception {
	final CountDownLatch activeSlots = new CountDownLatch(2);
	final TaskSlotTable taskSlotTable = new ActivateSlotNotifyingTaskSlotTable(
			Arrays.asList(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN),
			timerService,
			activeSlots);
	final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();

	final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);

	final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();

	final BlockingQueue<AllocationID> allocationsNotifiedFree = new ArrayBlockingQueue<>(2);

	OneShotLatch initialSlotReporting = new OneShotLatch();
	testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
		initialSlotReporting.trigger();
		return CompletableFuture.completedFuture(Acknowledge.get());

	});

	testingResourceManagerGateway.setNotifySlotAvailableConsumer(instanceIDSlotIDAllocationIDTuple3 ->
			allocationsNotifiedFree.offer(instanceIDSlotIDAllocationIDTuple3.f2));

	rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
	resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());

	final BlockingQueue<AllocationID> failedSlotFutures = new ArrayBlockingQueue<>(2);
	final ResourceID jobManagerResourceId = ResourceID.generate();
	final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder()
			.setFailSlotConsumer((resourceID, allocationID, throwable) ->
				failedSlotFutures.offer(allocationID))
			.setOfferSlotsFunction((resourceID, slotOffers) -> CompletableFuture.completedFuture(new ArrayList<>(slotOffers)))
			.setRegisterTaskManagerFunction((ignoredA, ignoredB) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jobManagerResourceId)))
			.build();
	final String jobManagerAddress = jobMasterGateway.getAddress();
	rpc.registerGateway(jobManagerAddress, jobMasterGateway);
	jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterGateway.getFencingToken().toUUID());

	taskExecutor.start();

	try {
		final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);

		initialSlotReporting.await();

		final SlotID slotId1 = new SlotID(taskExecutor.getResourceID(), 0);
		final SlotID slotId2 = new SlotID(taskExecutor.getResourceID(), 1);
		final AllocationID allocationIdInBoth = new AllocationID();
		final AllocationID allocationIdOnlyInJM = new AllocationID();
		final AllocationID allocationIdOnlyInTM = new AllocationID();

		taskExecutorGateway.requestSlot(slotId1, jobId, allocationIdInBoth, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
		taskExecutorGateway.requestSlot(slotId2, jobId, allocationIdOnlyInTM, "foobar", testingResourceManagerGateway.getFencingToken(), timeout);

		activeSlots.await();

		List<AllocatedSlotInfo> allocatedSlotInfos = Arrays.asList(
				new AllocatedSlotInfo(0, allocationIdInBoth),
				new AllocatedSlotInfo(1, allocationIdOnlyInJM)
		);
		AllocatedSlotReport allocatedSlotReport = new AllocatedSlotReport(jobId, allocatedSlotInfos);
		taskExecutorGateway.heartbeatFromJobManager(jobManagerResourceId, allocatedSlotReport);

		assertThat(failedSlotFutures.take(), is(allocationIdOnlyInJM));
		assertThat(allocationsNotifiedFree.take(), is(allocationIdOnlyInTM));
		assertThat(failedSlotFutures.poll(5L, TimeUnit.MILLISECONDS), nullValue());
		assertThat(allocationsNotifiedFree.poll(5L, TimeUnit.MILLISECONDS), nullValue());
	} finally {
		RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
	}
}
 

public TestingTaskExecutorGatewayBuilder setHeartbeatJobManagerConsumer(BiConsumer<ResourceID, AllocatedSlotReport> heartbeatJobManagerConsumer) {
	this.heartbeatJobManagerConsumer = heartbeatJobManagerConsumer;
	return this;
}