Java Code Examples for org.apache.flink.runtime.clusterframework.types.ResourceProfile#UNKNOWN

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

private static ArchivedExecutionJobVertex createArchivedExecutionJobVertex(JobVertexID jobVertexID) {
	final StringifiedAccumulatorResult[] emptyAccumulators = new StringifiedAccumulatorResult[0];
	final long[] timestamps = new long[ExecutionState.values().length];
	final ExecutionState expectedState = ExecutionState.RUNNING;

	final LocalTaskManagerLocation assignedResourceLocation = new LocalTaskManagerLocation();
	final AllocationID allocationID = new AllocationID();

	final int subtaskIndex = 1;
	final int attempt = 2;
	return new ArchivedExecutionJobVertex(
		new ArchivedExecutionVertex[]{
			new ArchivedExecutionVertex(
				subtaskIndex,
				"test task",
				new ArchivedExecution(
					new StringifiedAccumulatorResult[0],
					null,
					new ExecutionAttemptID(),
					attempt,
					expectedState,
					"error",
					assignedResourceLocation,
					allocationID,
					subtaskIndex,
					timestamps),
				new EvictingBoundedList<>(0)
			)
		},
		jobVertexID,
		jobVertexID.toString(),
		1,
		1,
		ResourceProfile.UNKNOWN,
		emptyAccumulators);
}
 

/**
 * Create a new worker with the given taskID.
 * @return a new worker instance.
 */
public static Worker newWorker(Protos.TaskID taskID) {
	return new Worker(
		taskID,
		ResourceProfile.UNKNOWN,
		Option.<Protos.SlaveID>empty(), Option.<String>empty(),
		WorkerState.New);
}
 

/**
 * Create a new worker with the given taskID.
 * @return a new worker instance.
 */
public static Worker newWorker(Protos.TaskID taskID) {
	return new Worker(
		taskID,
		ResourceProfile.UNKNOWN,
		Option.<Protos.SlaveID>empty(), Option.<String>empty(),
		WorkerState.New);
}
 

public ArchivedExecutionJobVertex build() {
	Preconditions.checkNotNull(taskVertices);
	return new ArchivedExecutionJobVertex(
		taskVertices,
		id != null ? id : new JobVertexID(),
		name != null ? name : "task_" + RANDOM.nextInt(),
		parallelism,
		maxParallelism,
		ResourceProfile.UNKNOWN,
		archivedUserAccumulators != null ? archivedUserAccumulators : new StringifiedAccumulatorResult[0]
	);
}
 

public SimpleSlotContext(
		AllocationID allocationId,
		TaskManagerLocation taskManagerLocation,
		int physicalSlotNumber,
		TaskManagerGateway taskManagerGateway) {
	this(allocationId, taskManagerLocation, physicalSlotNumber, taskManagerGateway, ResourceProfile.UNKNOWN);
}
 

@Test
public void testCheckIdleSlot() throws Exception {
	final ManualClock clock = new ManualClock();

	try (SlotPoolImpl slotPool = new SlotPoolImpl(
		jobId,
		clock,
		TestingUtils.infiniteTime(),
		timeout)) {
		final BlockingQueue<AllocationID> freedSlots = new ArrayBlockingQueue<>(1);
		taskManagerGateway.setFreeSlotFunction(
			(AllocationID allocationId, Throwable cause) -> {
				try {
					freedSlots.put(allocationId);
					return CompletableFuture.completedFuture(Acknowledge.get());
				} catch (InterruptedException e) {
					return FutureUtils.completedExceptionally(e);
				}
			});

		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);

		final AllocationID expiredSlotID = new AllocationID();
		final AllocationID freshSlotID = new AllocationID();
		final SlotOffer slotToExpire = new SlotOffer(expiredSlotID, 0, ResourceProfile.UNKNOWN);
		final SlotOffer slotToNotExpire = new SlotOffer(freshSlotID, 1, ResourceProfile.UNKNOWN);

		assertThat(slotPool.registerTaskManager(taskManagerLocation.getResourceID()),
			Matchers.is(true));

		assertThat(
			slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotToExpire),
			Matchers.is(true));

		clock.advanceTime(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		assertThat(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotToNotExpire),
			Matchers.is(true));

		clock.advanceTime(1L, TimeUnit.MILLISECONDS);

		slotPool.triggerCheckIdleSlot();

		final AllocationID freedSlot = freedSlots.poll(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);

		assertThat(freedSlot, Matchers.is(expiredSlotID));
		assertThat(freedSlots.isEmpty(), Matchers.is(true));
	}
}
 

/**
 * Tests that idle slots which cannot be released will be discarded. See FLINK-11059.
 */
@Test
public void testDiscardIdleSlotIfReleasingFailed() throws Exception {
	final ManualClock clock = new ManualClock();

	try (SlotPoolImpl slotPool = new SlotPoolImpl(
		jobId,
		clock,
		TestingUtils.infiniteTime(),
		timeout)) {

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

		final AllocationID expiredAllocationId = new AllocationID();
		final SlotOffer slotToExpire = new SlotOffer(expiredAllocationId, 0, ResourceProfile.UNKNOWN);

		OneShotLatch freeSlotLatch = new OneShotLatch();
		taskManagerGateway.setFreeSlotFunction((AllocationID allocationId, Throwable cause) -> {
			freeSlotLatch.trigger();
			return FutureUtils.completedExceptionally(new TimeoutException("Test failure"));
		});

		assertThat(slotPool.registerTaskManager(taskManagerLocation.getResourceID()), Matchers.is(true));

		assertThat(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotToExpire), Matchers.is(true));

		clock.advanceTime(timeout.toMilliseconds() + 1, TimeUnit.MILLISECONDS);

		slotPool.triggerCheckIdleSlot();

		freeSlotLatch.await();

		CompletableFuture<LogicalSlot> allocatedSlotFuture = allocateSlot(scheduler, new SlotRequestId());

		try {
			// since the slot must have been discarded, we cannot fulfill the slot request
			allocatedSlotFuture.get(10L, TimeUnit.MILLISECONDS);
			fail("Expected to fail with a timeout.");
		} catch (TimeoutException ignored) {
			// expected
			assertEquals(0, slotPool.getAvailableSlots().size());
		}
	}
}
 

/**
 * Tests that the {@link AllocatedSlotReport} contains up to date information and not
 * stale information about the allocated slots on the {@link JobMaster}.
 *
 * <p>This is a probabilistic test case which only fails if executed repeatedly without
 * the fix for FLINK-12863.
 */
@Test
public void testAllocatedSlotReportDoesNotContainStaleInformation() throws Exception {
	final CompletableFuture<Void> assertionFuture = new CompletableFuture<>();
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	final AtomicBoolean terminateHeartbeatVerification = new AtomicBoolean(false);
	final OneShotLatch hasReceivedSlotOffers = new OneShotLatch();
	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
		.setHeartbeatJobManagerConsumer((taskManagerId, allocatedSlotReport) -> {
			try {
				if (hasReceivedSlotOffers.isTriggered()) {
					assertThat(allocatedSlotReport.getAllocatedSlotInfos(), hasSize(1));
				} else {
					assertThat(allocatedSlotReport.getAllocatedSlotInfos(), empty());
				}
			} catch (AssertionError e) {
				assertionFuture.completeExceptionally(e);
			}

			if (terminateHeartbeatVerification.get()) {
				assertionFuture.complete(null);
			}
		})
		.createTestingTaskExecutorGateway();

	rpcService.registerGateway(taskExecutorGateway.getAddress(), taskExecutorGateway);

	final JobManagerSharedServices jobManagerSharedServices = new TestingJobManagerSharedServicesBuilder().build();

	final JobMaster jobMaster = new JobMasterBuilder()
		.withJobGraph(createSingleVertexJobGraph())
		.withHeartbeatServices(new HeartbeatServices(5L, 1000L))
		.withSlotPoolFactory(new TestingSlotPoolFactory(hasReceivedSlotOffers))
		.createJobMaster();

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

	try {
		// wait for the start to complete
		startFuture.get(testingTimeout.toMilliseconds(), TimeUnit.MILLISECONDS);

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

		// register task manager will trigger monitor heartbeat target, schedule heartbeat request at interval time
		CompletableFuture<RegistrationResponse> registrationResponse = jobMasterGateway.registerTaskManager(
			taskExecutorGateway.getAddress(),
			taskManagerLocation,
			testingTimeout);

		// wait for the completion of the registration
		registrationResponse.get();

		final SlotOffer slotOffer = new SlotOffer(new AllocationID(), 0, ResourceProfile.UNKNOWN);

		final CompletableFuture<Collection<SlotOffer>> slotOfferFuture = jobMasterGateway.offerSlots(taskManagerLocation.getResourceID(), Collections.singleton(slotOffer), testingTimeout);

		assertThat(slotOfferFuture.get(), containsInAnyOrder(slotOffer));

		terminateHeartbeatVerification.set(true);

		// make sure that no assertion has been violated
		assertionFuture.get();
	} finally {
		RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
		jobManagerSharedServices.shutdown();
	}
}
 

/**
 * Tests that pending request is removed if task executor reports a slot with its allocation id.
 */
@Test
public void testSlotRequestRemovedIfTMReportAllocation() throws Exception {
	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(),
			new TestingResourceActionsBuilder().build())) {

		final JobID jobID = new JobID();
		final SlotRequest slotRequest1 = new SlotRequest(jobID, new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
		slotManager.registerSlotRequest(slotRequest1);

		final BlockingQueue<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestSlotQueue = new ArrayBlockingQueue<>(1);
		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(1);

		final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
				.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple6 -> {
					requestSlotQueue.offer(slotIDJobIDAllocationIDStringResourceManagerIdTuple6);
					try {
						return responseQueue.take();
					} catch (InterruptedException ignored) {
						return FutureUtils.completedExceptionally(new FlinkException("Response queue was interrupted."));
					}
				})
				.createTestingTaskExecutorGateway();

		final ResourceID taskExecutorResourceId = ResourceID.generate();
		final TaskExecutorConnection taskExecutionConnection = new TaskExecutorConnection(taskExecutorResourceId, testingTaskExecutorGateway);
		final SlotReport slotReport = new SlotReport(createEmptySlotStatus(new SlotID(taskExecutorResourceId, 0), ResourceProfile.ANY));

		final CompletableFuture<Acknowledge> firstManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(firstManualSlotRequestResponse);

		slotManager.registerTaskManager(taskExecutionConnection, slotReport);

		final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> firstRequest = requestSlotQueue.take();

		final CompletableFuture<Acknowledge> secondManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(secondManualSlotRequestResponse);

		final SlotRequest slotRequest2 = new SlotRequest(jobID, new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
		slotManager.registerSlotRequest(slotRequest2);

		// fail first request
		firstManualSlotRequestResponse.completeExceptionally(new TimeoutException("Test exception to fail first allocation"));

		final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> secondRequest = requestSlotQueue.take();

		// fail second request
		secondManualSlotRequestResponse.completeExceptionally(new SlotOccupiedException("Test exception", slotRequest1.getAllocationId(), jobID));

		assertThat(firstRequest.f2, equalTo(slotRequest1.getAllocationId()));
		assertThat(secondRequest.f2, equalTo(slotRequest2.getAllocationId()));
		assertThat(secondRequest.f0, equalTo(firstRequest.f0));

		secondManualSlotRequestResponse.complete(Acknowledge.get());

		final TaskManagerSlot slot = slotManager.getSlot(secondRequest.f0);
		assertThat(slot.getState(), equalTo(TaskManagerSlot.State.ALLOCATED));
		assertThat(slot.getAllocationId(), equalTo(firstRequest.f2));

		assertThat(slotManager.getNumberRegisteredSlots(), is(1));
	}
}
 

/**
 * Tests that unused offered slots are directly used to fulfill pending slot
 * requests.
 *
 * <p>Moreover it tests that the old slot request is canceled
 *
 * <p>See FLINK-8089, FLINK-8934
 */
@Test
public void testFulfillingSlotRequestsWithUnusedOfferedSlots() throws Exception {

	try (SlotPoolImpl slotPool = new SlotPoolImpl(jobId)) {
		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(2);
		resourceManagerGateway.setRequestSlotConsumer(
			(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));
		final ArrayBlockingQueue<AllocationID> canceledAllocations = new ArrayBlockingQueue<>(2);
		resourceManagerGateway.setCancelSlotConsumer(canceledAllocations::offer);
		final SlotRequestId slotRequestId1 = new SlotRequestId();
		final SlotRequestId slotRequestId2 = new SlotRequestId();
		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		final Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final ScheduledUnit scheduledUnit = new ScheduledUnit(
			new JobVertexID(),
			null,
			null);

		CompletableFuture<LogicalSlot> slotFuture1 = scheduler.allocateSlot(
			slotRequestId1,
			scheduledUnit,
			SlotProfile.noRequirements(),
			true,
			timeout);

		// wait for the first slot request
		final AllocationID allocationId1 = allocationIds.take();

		CompletableFuture<LogicalSlot> slotFuture2 = scheduler.allocateSlot(
			slotRequestId2,
			scheduledUnit,
			SlotProfile.noRequirements(),
			true,
			timeout);

		// wait for the second slot request
		final AllocationID allocationId2 = allocationIds.take();

		slotPool.releaseSlot(slotRequestId1, null);

		try {
			// this should fail with a CancellationException
			slotFuture1.get();
			fail("The first slot future should have failed because it was cancelled.");
		} catch (ExecutionException ee) {
			// expected
			assertTrue(ExceptionUtils.stripExecutionException(ee) instanceof FlinkException);
		}

		assertEquals(allocationId1, canceledAllocations.take());

		final SlotOffer slotOffer = new SlotOffer(allocationId1, 0, ResourceProfile.UNKNOWN);

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());

		assertTrue(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotOffer));

		// the slot offer should fulfill the second slot request
		assertEquals(allocationId1, slotFuture2.get().getAllocationId());

		// check that the second slot allocation has been canceled
		assertEquals(allocationId2, canceledAllocations.take());
	}
}
 

/**
 * Tests that the SlotManager retries allocating a slot if the TaskExecutor#requestSlot call
 * fails.
 */
@Test
public void testSlotRequestFailure() throws Exception {
	try (final SlotManager slotManager = createSlotManager(ResourceManagerId.generate(),
		new TestingResourceActionsBuilder().build())) {

		final SlotRequest slotRequest = new SlotRequest(new JobID(), new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
		slotManager.registerSlotRequest(slotRequest);

		final BlockingQueue<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>> requestSlotQueue = new ArrayBlockingQueue<>(1);
		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(1);

		final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
			.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple5 -> {
				requestSlotQueue.offer(slotIDJobIDAllocationIDStringResourceManagerIdTuple5);
				try {
					return responseQueue.take();
				} catch (InterruptedException ignored) {
					return FutureUtils.completedExceptionally(new FlinkException("Response queue was interrupted."));
				}
			})
			.createTestingTaskExecutorGateway();

		final ResourceID taskExecutorResourceId = ResourceID.generate();
		final TaskExecutorConnection taskExecutionConnection = new TaskExecutorConnection(taskExecutorResourceId, testingTaskExecutorGateway);
		final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(taskExecutorResourceId, 0), ResourceProfile.UNKNOWN));

		final CompletableFuture<Acknowledge> firstManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(firstManualSlotRequestResponse);

		slotManager.registerTaskManager(taskExecutionConnection, slotReport);

		final Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId> firstRequest = requestSlotQueue.take();

		final CompletableFuture<Acknowledge> secondManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(secondManualSlotRequestResponse);

		// fail first request
		firstManualSlotRequestResponse.completeExceptionally(new SlotAllocationException("Test exception"));

		final Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId> secondRequest = requestSlotQueue.take();

		assertThat(secondRequest.f2, equalTo(firstRequest.f2));
		assertThat(secondRequest.f0, equalTo(firstRequest.f0));

		secondManualSlotRequestResponse.complete(Acknowledge.get());

		final TaskManagerSlot slot = slotManager.getSlot(secondRequest.f0);
		assertThat(slot.getState(), equalTo(TaskManagerSlot.State.ALLOCATED));
		assertThat(slot.getAllocationId(), equalTo(secondRequest.f2));
	}
}
 

private static SlotProfile slotProfileForLocation(TaskManagerLocation... location) {
	return new SlotProfile(ResourceProfile.UNKNOWN, Arrays.asList(location), Collections.emptySet());
}
 

/**
 * Tests that unused offered slots are directly used to fulfill pending slot
 * requests.
 *
 * <p>Moreover it tests that the old slot request is canceled
 *
 * <p>See FLINK-8089, FLINK-8934
 */
@Test
public void testFulfillingSlotRequestsWithUnusedOfferedSlots() throws Exception {

	try (SlotPoolImpl slotPool = createSlotPoolImpl()) {
		final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(2);
		resourceManagerGateway.setRequestSlotConsumer(
			(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));
		final ArrayBlockingQueue<AllocationID> canceledAllocations = new ArrayBlockingQueue<>(2);
		resourceManagerGateway.setCancelSlotConsumer(canceledAllocations::offer);
		final SlotRequestId slotRequestId1 = new SlotRequestId();
		final SlotRequestId slotRequestId2 = new SlotRequestId();
		setupSlotPool(slotPool, resourceManagerGateway, mainThreadExecutor);
		final Scheduler scheduler = setupScheduler(slotPool, mainThreadExecutor);

		final ScheduledUnit scheduledUnit = new ScheduledUnit(
			new JobVertexID(),
			null,
			null);

		CompletableFuture<LogicalSlot> slotFuture1 = scheduler.allocateSlot(
			slotRequestId1,
			scheduledUnit,
			SlotProfile.noRequirements(),
			true,
			timeout);

		// wait for the first slot request
		final AllocationID allocationId1 = allocationIds.take();

		CompletableFuture<LogicalSlot> slotFuture2 = scheduler.allocateSlot(
			slotRequestId2,
			scheduledUnit,
			SlotProfile.noRequirements(),
			true,
			timeout);

		// wait for the second slot request
		final AllocationID allocationId2 = allocationIds.take();

		slotPool.releaseSlot(slotRequestId1, null);

		try {
			// this should fail with a CancellationException
			slotFuture1.get();
			fail("The first slot future should have failed because it was cancelled.");
		} catch (ExecutionException ee) {
			// expected
			assertTrue(ExceptionUtils.stripExecutionException(ee) instanceof FlinkException);
		}

		assertEquals(allocationId1, canceledAllocations.take());

		final SlotOffer slotOffer = new SlotOffer(allocationId1, 0, ResourceProfile.UNKNOWN);

		slotPool.registerTaskManager(taskManagerLocation.getResourceID());

		assertTrue(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotOffer));

		// the slot offer should fulfill the second slot request
		assertEquals(allocationId1, slotFuture2.get().getAllocationId());

		// check that the second slot allocation has been canceled
		assertEquals(allocationId2, canceledAllocations.take());
	}
}
 

private static SlotProfile slotProfileForLocation(TaskManagerLocation location) {
	return new SlotProfile(ResourceProfile.UNKNOWN, Collections.singletonList(location), Collections.emptySet());
}
 

/**
 * Tests that idle slots which cannot be released will be discarded. See FLINK-11059.
 */
@Test
public void testDiscardIdleSlotIfReleasingFailed() throws Exception {
	final ManualClock clock = new ManualClock();

	try (TestingSlotPoolImpl slotPool = createSlotPoolImpl(clock)) {

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

		final AllocationID expiredAllocationId = new AllocationID();
		final SlotOffer slotToExpire = new SlotOffer(expiredAllocationId, 0, ResourceProfile.UNKNOWN);

		OneShotLatch freeSlotLatch = new OneShotLatch();
		taskManagerGateway.setFreeSlotFunction((AllocationID allocationId, Throwable cause) -> {
			freeSlotLatch.trigger();
			return FutureUtils.completedExceptionally(new TimeoutException("Test failure"));
		});

		assertThat(slotPool.registerTaskManager(taskManagerLocation.getResourceID()), Matchers.is(true));

		assertThat(slotPool.offerSlot(taskManagerLocation, taskManagerGateway, slotToExpire), Matchers.is(true));

		clock.advanceTime(timeout.toMilliseconds() + 1, TimeUnit.MILLISECONDS);

		slotPool.triggerCheckIdleSlot();

		freeSlotLatch.await();

		CompletableFuture<LogicalSlot> allocatedSlotFuture = allocateSlot(scheduler, new SlotRequestId());

		try {
			// since the slot must have been discarded, we cannot fulfill the slot request
			allocatedSlotFuture.get(10L, TimeUnit.MILLISECONDS);
			fail("Expected to fail with a timeout.");
		} catch (TimeoutException ignored) {
			// expected
			assertEquals(0, slotPool.getAvailableSlots().size());
		}
	}
}
 

@Override
public ResourceProfile getResourceProfile() {
	return ResourceProfile.UNKNOWN;
}
 

/**
 * Tests that pending request is removed if task executor reports a slot with its allocation id.
 */
@Test
public void testSlotRequestRemovedIfTMReportAllocation() throws Exception {
	try (final SlotManagerImpl slotManager = createSlotManager(ResourceManagerId.generate(),
			new TestingResourceActionsBuilder().build())) {

		final JobID jobID = new JobID();
		final SlotRequest slotRequest1 = new SlotRequest(jobID, new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
		slotManager.registerSlotRequest(slotRequest1);

		final BlockingQueue<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>> requestSlotQueue = new ArrayBlockingQueue<>(1);
		final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(1);

		final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
				.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple5 -> {
					requestSlotQueue.offer(slotIDJobIDAllocationIDStringResourceManagerIdTuple5);
					try {
						return responseQueue.take();
					} catch (InterruptedException ignored) {
						return FutureUtils.completedExceptionally(new FlinkException("Response queue was interrupted."));
					}
				})
				.createTestingTaskExecutorGateway();

		final ResourceID taskExecutorResourceId = ResourceID.generate();
		final TaskExecutorConnection taskExecutionConnection = new TaskExecutorConnection(taskExecutorResourceId, testingTaskExecutorGateway);
		final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(taskExecutorResourceId, 0), ResourceProfile.UNKNOWN));

		final CompletableFuture<Acknowledge> firstManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(firstManualSlotRequestResponse);

		slotManager.registerTaskManager(taskExecutionConnection, slotReport);

		final Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId> firstRequest = requestSlotQueue.take();

		final CompletableFuture<Acknowledge> secondManualSlotRequestResponse = new CompletableFuture<>();
		responseQueue.offer(secondManualSlotRequestResponse);

		final SlotRequest slotRequest2 = new SlotRequest(jobID, new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
		slotManager.registerSlotRequest(slotRequest2);

		// fail first request
		firstManualSlotRequestResponse.completeExceptionally(new TimeoutException("Test exception to fail first allocation"));

		final Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId> secondRequest = requestSlotQueue.take();

		// fail second request
		secondManualSlotRequestResponse.completeExceptionally(new SlotOccupiedException("Test exception", slotRequest1.getAllocationId(), jobID));

		assertThat(firstRequest.f2, equalTo(slotRequest1.getAllocationId()));
		assertThat(secondRequest.f2, equalTo(slotRequest2.getAllocationId()));
		assertThat(secondRequest.f0, equalTo(firstRequest.f0));

		secondManualSlotRequestResponse.complete(Acknowledge.get());

		final TaskManagerSlot slot = slotManager.getSlot(secondRequest.f0);
		assertThat(slot.getState(), equalTo(TaskManagerSlot.State.ALLOCATED));
		assertThat(slot.getAllocationId(), equalTo(firstRequest.f2));

		assertThat(slotManager.getNumberRegisteredSlots(), is(1));
	}
}
 

/**
 * Tests that free slots which are reported as allocated won't be considered for fulfilling
 * other pending slot requests.
 *
 * <p>See: FLINK-8505
 */
@Test
public void testReportAllocatedSlot() throws Exception {
	final ResourceID taskManagerId = ResourceID.generate();
	final ResourceActions resourceActions = new TestingResourceActionsBuilder().build();
	final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway();
	final TaskExecutorConnection taskExecutorConnection = new TaskExecutorConnection(taskManagerId, taskExecutorGateway);

	try (final SlotManagerImpl slotManager = SlotManagerBuilder.newBuilder().build()) {

		slotManager.start(ResourceManagerId.generate(), Executors.directExecutor(), resourceActions);

		// initially report a single slot as free
		final SlotID slotId = new SlotID(taskManagerId, 0);
		final SlotStatus initialSlotStatus = new SlotStatus(
			slotId,
			ResourceProfile.UNKNOWN);
		final SlotReport initialSlotReport = new SlotReport(initialSlotStatus);

		slotManager.registerTaskManager(taskExecutorConnection, initialSlotReport);

		assertThat(slotManager.getNumberRegisteredSlots(), is(equalTo(1)));

		// Now report this slot as allocated
		final SlotStatus slotStatus = new SlotStatus(
			slotId,
			ResourceProfile.UNKNOWN,
			new JobID(),
			new AllocationID());
		final SlotReport slotReport = new SlotReport(
			slotStatus);

		slotManager.reportSlotStatus(
			taskExecutorConnection.getInstanceID(),
			slotReport);

		// this slot request should not be fulfilled
		final AllocationID allocationId = new AllocationID();
		final SlotRequest slotRequest = new SlotRequest(
			new JobID(),
			allocationId,
			ResourceProfile.UNKNOWN,
			"foobar");

		// This triggered an IllegalStateException before
		slotManager.registerSlotRequest(slotRequest);

		assertThat(slotManager.getSlotRequest(allocationId).isAssigned(), is(false));
	}
}