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

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

private void initializeLocation(ExecutionVertex vertex, TaskManagerLocation location) throws Exception {
	// we need a bit of reflection magic to initialize the location without going through
	// scheduling paths. we choose to do that, rather than the alternatives:
	//  - mocking the scheduler created fragile tests that break whenever the scheduler is adjusted
	//  - exposing test methods in the ExecutionVertex leads to undesirable setters 

	SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));



	LogicalSlot slot = new SingleLogicalSlot(
		new SlotRequestId(),
		slotContext,
		null,
		Locality.LOCAL,
		mock(SlotOwner.class));

	if (!vertex.getCurrentExecutionAttempt().tryAssignResource(slot)) {
		throw new FlinkException("Could not assign resource.");
	}
}
 

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

public SimpleSlotProvider(JobID jobId, int numSlots, TaskManagerGateway taskManagerGateway) {
	checkNotNull(jobId, "jobId");
	checkArgument(numSlots >= 0, "numSlots must be >= 0");

	this.slots = new ArrayDeque<>(numSlots);

	for (int i = 0; i < numSlots; i++) {
		SimpleSlotContext as = new SimpleSlotContext(
			new AllocationID(),
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i),
			0,
			taskManagerGateway,
			ResourceProfile.UNKNOWN);
		slots.add(as);
	}

	allocatedSlots = new HashMap<>(slots.size());
}
 

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

private void initializeLocation(ExecutionVertex vertex, TaskManagerLocation location) throws Exception {
	// we need a bit of reflection magic to initialize the location without going through
	// scheduling paths. we choose to do that, rather than the alternatives:
	//  - mocking the scheduler created fragile tests that break whenever the scheduler is adjusted
	//  - exposing test methods in the ExecutionVertex leads to undesirable setters 

	SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));



	LogicalSlot slot = new SingleLogicalSlot(
		new SlotRequestId(),
		slotContext,
		null,
		Locality.LOCAL,
		mock(SlotOwner.class));

	if (!vertex.getCurrentExecutionAttempt().tryAssignResource(slot)) {
		throw new FlinkException("Could not assign resource.");
	}
}
 

public SimpleSlotProvider(JobID jobId, int numSlots, TaskManagerGateway taskManagerGateway) {
	checkNotNull(jobId, "jobId");
	checkArgument(numSlots >= 0, "numSlots must be >= 0");

	this.slots = new ArrayDeque<>(numSlots);

	for (int i = 0; i < numSlots; i++) {
		SimpleSlotContext as = new SimpleSlotContext(
			new AllocationID(),
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i),
			0,
			taskManagerGateway);
		slots.add(as);
	}

	allocatedSlots = new HashMap<>(slots.size());
}
 

private void initializeLocation(ExecutionVertex vertex, TaskManagerLocation location) throws Exception {
	// we need a bit of reflection magic to initialize the location without going through
	// scheduling paths. we choose to do that, rather than the alternatives:
	//  - mocking the scheduler created fragile tests that break whenever the scheduler is adjusted
	//  - exposing test methods in the ExecutionVertex leads to undesirable setters 

	SlotContext slot = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));

	SimpleSlot simpleSlot = new SimpleSlot(slot, mock(SlotOwner.class), 0);

	if (!vertex.getCurrentExecutionAttempt().tryAssignResource(simpleSlot)) {
		throw new FlinkException("Could not assign resource.");
	}
}
 

@Override
public Collection<SlotOffer> offerSlots(TaskManagerLocation taskManagerLocation, TaskManagerGateway taskManagerGateway, Collection<SlotOffer> offers) {
	hasReceivedSlotOffers.trigger();
	final Collection<SlotInfo> slotInfos = Optional.ofNullable(registeredSlots.get(taskManagerLocation.getResourceID()))
		.orElseThrow(() -> new FlinkRuntimeException("TaskManager not registered."));

	int slotIndex = slotInfos.size();

	for (SlotOffer offer : offers) {
		slotInfos.add(new SimpleSlotContext(
			offer.getAllocationId(),
			taskManagerLocation,
			slotIndex,
			taskManagerGateway));
		slotIndex++;
	}

	return offers;
}
 

public SimpleSlotProvider(int numSlots, TaskManagerGateway taskManagerGateway) {
	checkArgument(numSlots >= 0, "numSlots must be >= 0");

	this.slots = new ArrayDeque<>(numSlots);

	for (int i = 0; i < numSlots; i++) {
		SimpleSlotContext as = new SimpleSlotContext(
			new AllocationID(),
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i),
			0,
			taskManagerGateway,
			ResourceProfile.ANY);
		slots.add(as);
	}

	allocatedSlots = new HashMap<>(slots.size());
}
 

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

/**
 * Tests that the root slot are moved from unresolved to resolved once the
 * slot context future is successfully completed
 */
@Test
public void testRootSlotTransition() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	assertTrue(slotSharingManager.getUnresolvedRootSlots().contains(rootSlot));
	assertFalse(slotSharingManager.getResolvedRootSlots().contains(rootSlot));

	// now complete the slotContextFuture
	slotContextFuture.complete(
		new SimpleSlotContext(
			new AllocationID(),
			new LocalTaskManagerLocation(),
			0,
			new SimpleAckingTaskManagerGateway()));

	assertFalse(slotSharingManager.getUnresolvedRootSlots().contains(rootSlot));
	assertTrue(slotSharingManager.getResolvedRootSlots().contains(rootSlot));
}
 

private SlotSharingManager.MultiTaskSlot createRootSlot(TaskManagerLocation firstTaskExecutorLocation, SlotSharingManager slotSharingManager) {
	return slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				firstTaskExecutorLocation,
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());
}
 

@Test
public void testGetResolvedSlotWithResourceConfigured() {
	ResourceProfile rp1 = ResourceProfile.fromResources(1.0, 100);
	ResourceProfile rp2 = ResourceProfile.fromResources(2.0, 200);
	ResourceProfile allocatedSlotRp = ResourceProfile.fromResources(5.0, 500);

	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
			new SlotRequestId(),
			CompletableFuture.completedFuture(
					new SimpleSlotContext(
							new AllocationID(),
							new LocalTaskManagerLocation(),
							0,
							new SimpleAckingTaskManagerGateway(),
							allocatedSlotRp)),
			new SlotRequestId());

	rootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp1,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	Collection<SlotSelectionStrategy.SlotInfoAndResources> resolvedRoots =
		slotSharingManager.listResolvedRootSlotInfo(new AbstractID());
	assertEquals(1, resolvedRoots.size());
	assertEquals(allocatedSlotRp.subtract(rp1), resolvedRoots.iterator().next().getRemainingResources());

	rootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp2,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	resolvedRoots = slotSharingManager.listResolvedRootSlotInfo(new AbstractID());
	assertEquals(1, resolvedRoots.size());
	assertEquals(allocatedSlotRp.subtract(rp1).subtract(rp2), resolvedRoots.iterator().next().getRemainingResources());
}
 

private SimpleSlotContext createSimpleSlotContext() {
	return new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway());
}
 

private static SlotContext createSlotContext() {
	return new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway(),
		ResourceProfile.ANY);
}
 

@Override
public void returnLogicalSlot(LogicalSlot logicalSlot) {
	synchronized (lock) {
		SimpleSlotContext as = new SimpleSlotContext(
			logicalSlot.getAllocationId(),
			logicalSlot.getTaskManagerLocation(),
			logicalSlot.getPhysicalSlotNumber(),
			logicalSlot.getTaskManagerGateway(),
			ResourceProfile.UNKNOWN);

		slots.add(as);
		allocatedSlots.remove(logicalSlot.getSlotRequestId());
	}
}
 

@Override
public void returnLogicalSlot(LogicalSlot logicalSlot) {
	synchronized (lock) {
		SimpleSlotContext as = new SimpleSlotContext(
			logicalSlot.getAllocationId(),
			logicalSlot.getTaskManagerLocation(),
			logicalSlot.getPhysicalSlotNumber(),
			logicalSlot.getTaskManagerGateway(),
			ResourceProfile.ANY);

		slots.add(as);
		allocatedSlots.remove(logicalSlot.getSlotRequestId());
	}
}
 

/**
 * Tests that the root slot are moved from unresolved to resolved once the
 * slot context future is successfully completed
 */
@Test
public void testRootSlotTransition() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	assertTrue(slotSharingManager.getUnresolvedRootSlots().contains(rootSlot));
	assertFalse(slotSharingManager.getResolvedRootSlots().contains(rootSlot));

	// now complete the slotContextFuture
	slotContextFuture.complete(
		new SimpleSlotContext(
			new AllocationID(),
			new LocalTaskManagerLocation(),
			0,
			new SimpleAckingTaskManagerGateway()));

	assertFalse(slotSharingManager.getUnresolvedRootSlots().contains(rootSlot));
	assertTrue(slotSharingManager.getResolvedRootSlots().contains(rootSlot));
}
 

private static SlotContext createSlotContext() {
	return new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway(),
		ResourceProfile.UNKNOWN);
}
 

public static SimpleSlot createMockSimpleSlot(TaskManagerGateway gateway) {
	final TaskManagerLocation location = new TaskManagerLocation(
			ResourceID.generate(), InetAddress.getLoopbackAddress(), 6572);

	final SimpleSlotContext allocatedSlot = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		gateway);

	return new SimpleSlot(
		allocatedSlot,
		mock(SlotOwner.class),
		0);
}
 

private SimpleSlot createSlot(TaskManagerGateway taskManager, JobID jobId, SlotOwner slotOwner) {
	TaskManagerLocation location = new TaskManagerLocation(
			ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slot = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManager);

	return new SimpleSlot(slot, slotOwner, 0);
}
 

@Test
public void testGetResolvedSlotWithResourceConfigured() {
	ResourceProfile rp1 = new ResourceProfile(1.0, 100);
	ResourceProfile rp2 = new ResourceProfile(2.0, 200);
	ResourceProfile allocatedSlotRp = new ResourceProfile(5.0, 500);

	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
			SLOT_SHARING_GROUP_ID,
			allocatedSlotActions,
			SLOT_OWNER);

	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
			new SlotRequestId(),
			CompletableFuture.completedFuture(
					new SimpleSlotContext(
							new AllocationID(),
							new LocalTaskManagerLocation(),
							0,
							new SimpleAckingTaskManagerGateway(),
							allocatedSlotRp)),
			new SlotRequestId());

	rootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp1,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	Collection<SlotSelectionStrategy.SlotInfoAndResources> resolvedRoots =
		slotSharingManager.listResolvedRootSlotInfo(new AbstractID());
	assertEquals(1, resolvedRoots.size());
	assertEquals(allocatedSlotRp.subtract(rp1), resolvedRoots.iterator().next().getRemainingResources());

	rootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp2,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	resolvedRoots = slotSharingManager.listResolvedRootSlotInfo(new AbstractID());
	assertEquals(1, resolvedRoots.size());
	assertEquals(allocatedSlotRp.subtract(rp1).subtract(rp2), resolvedRoots.iterator().next().getRemainingResources());
}
 

/**
 * Tests that we cannot retrieve a slot when it's releasing children.
 */
@Test
public void testResolvedSlotInReleasingIsNotAvailable() throws Exception {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	final SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	final SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				new LocalTaskManagerLocation(),
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());

	final AbstractID groupId1 = new AbstractID();
	final SlotSharingManager.SingleTaskSlot singleTaskSlot = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		groupId1,
		Locality.UNCONSTRAINED);

	final AtomicBoolean verified = new AtomicBoolean(false);

	final AbstractID groupId2 = new AbstractID();
	// register a verification in MultiTaskSlot's children releasing loop
	singleTaskSlot.getLogicalSlotFuture().get().tryAssignPayload(new LogicalSlot.Payload() {
		@Override
		public void fail(Throwable cause) {
			assertEquals(0, slotSharingManager.listResolvedRootSlotInfo(groupId2).size());

			verified.set(true);
		}

		@Override
		public CompletableFuture<?> getTerminalStateFuture() {
			return null;
		}
	});

	assertEquals(1, slotSharingManager.listResolvedRootSlotInfo(groupId2).size());

	rootSlot.release(new Exception("test exception"));

	// ensure the verification in Payload#fail is passed
	assertTrue(verified.get());
}
 

/**
 * Tests that the location preferences are honoured when looking for a resolved slot.
 */
@Test
public void testGetResolvedSlotWithLocationPreferences() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	SlotSharingManager.MultiTaskSlot rootSlot1 = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				new LocalTaskManagerLocation(),
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	SlotSharingManager.MultiTaskSlot rootSlot2 = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				taskManagerLocation,
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());

	AbstractID groupId = new AbstractID();

	SlotProfile slotProfile = SlotProfile.preferredLocality(ResourceProfile.UNKNOWN, Collections.singleton(taskManagerLocation));

	Collection<SlotSelectionStrategy.SlotInfoAndResources> slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	SlotSelectionStrategy.SlotInfoAndLocality slotInfoAndLocality =
		LocationPreferenceSlotSelectionStrategy.INSTANCE.selectBestSlotForProfile(slotInfos, slotProfile).get();
	SlotSharingManager.MultiTaskSlot resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());

	assertNotNull(resolvedRootSlot);
	assertEquals(Locality.LOCAL, slotInfoAndLocality.getLocality());
	assertEquals(rootSlot2.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());

	// occupy the slot
	resolvedRootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		groupId,
		slotInfoAndLocality.getLocality());

	slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	slotInfoAndLocality = LocationPreferenceSlotSelectionStrategy.INSTANCE.selectBestSlotForProfile(slotInfos, slotProfile).get();
	resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());
	assertNotNull(resolvedRootSlot);
	assertNotSame(Locality.LOCAL, (slotInfoAndLocality.getLocality()));
	assertEquals(rootSlot1.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());
}
 

/**
 * Tests that we can correctly retrieve resolved slots.
 */
@Test
public void testGetResolvedSlot() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				new LocalTaskManagerLocation(),
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());

	AbstractID groupId = new AbstractID();

	Collection<SlotSelectionStrategy.SlotInfoAndResources> slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	Assert.assertEquals(1, slotInfos.size());

	SlotSelectionStrategy.SlotInfoAndResources slotInfoAndRemainingResource = slotInfos.iterator().next();
	SlotSharingManager.MultiTaskSlot resolvedMultiTaskSlot =
		slotSharingManager.getResolvedRootSlot(slotInfoAndRemainingResource.getSlotInfo());

	SlotSelectionStrategy.SlotInfoAndLocality slotInfoAndLocality =
		LocationPreferenceSlotSelectionStrategy.INSTANCE.selectBestSlotForProfile(slotInfos, SlotProfile.noRequirements()).get();

	assertNotNull(resolvedMultiTaskSlot);
	assertEquals(Locality.UNCONSTRAINED, slotInfoAndLocality.getLocality());
	assertEquals(rootSlot.getSlotRequestId(), resolvedMultiTaskSlot.getSlotRequestId());

	// occupy the resolved root slot
	resolvedMultiTaskSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		groupId,
		Locality.UNCONSTRAINED);

	slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	assertTrue(slotInfos.isEmpty());
}
 

/**
 * Tests that the logical task slot futures are completed once the slot context
 * future is completed.
 */
@Test
public void testSlotContextFutureCompletion() throws Exception {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	final SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	final SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway());

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	Locality locality1 = Locality.LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot1 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality1);

	Locality locality2 = Locality.HOST_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot2 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality2);

	CompletableFuture<LogicalSlot> logicalSlotFuture1 = singleTaskSlot1.getLogicalSlotFuture();
	CompletableFuture<LogicalSlot> logicalSlotFuture2 = singleTaskSlot2.getLogicalSlotFuture();
	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());

	slotContextFuture.complete(slotContext);

	assertTrue(logicalSlotFuture1.isDone());
	assertTrue(logicalSlotFuture2.isDone());

	final LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	final LogicalSlot logicalSlot2 = logicalSlotFuture2.get();

	assertEquals(logicalSlot1.getAllocationId(), slotContext.getAllocationId());
	assertEquals(logicalSlot2.getAllocationId(), slotContext.getAllocationId());
	assertEquals(locality1, logicalSlot1.getLocality());
	assertEquals(locality2, logicalSlot2.getLocality());

	Locality locality3 = Locality.NON_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot3 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality3);

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = singleTaskSlot3.getLogicalSlotFuture();

	assertTrue(logicalSlotFuture3.isDone());
	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();

	assertEquals(locality3, logicalSlot3.getLocality());
	assertEquals(slotContext.getAllocationId(), logicalSlot3.getAllocationId());
}
 

/**
 * Tests that the location preferences are honoured when looking for a resolved slot.
 */
@Test
public void testGetResolvedSlotWithLocationPreferences() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	SlotSharingManager.MultiTaskSlot rootSlot1 = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				new LocalTaskManagerLocation(),
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());

	LocalTaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
	SlotSharingManager.MultiTaskSlot rootSlot2 = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				taskManagerLocation,
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());

	AbstractID groupId = new AbstractID();

	SlotProfile slotProfile = SlotProfile.preferredLocality(ResourceProfile.UNKNOWN, Collections.singleton(taskManagerLocation));

	Collection<SlotInfo> slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	SlotSelectionStrategy.SlotInfoAndLocality slotInfoAndLocality =
		LocationPreferenceSlotSelectionStrategy.INSTANCE.selectBestSlotForProfile(slotInfos, slotProfile).get();
	SlotSharingManager.MultiTaskSlot resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());

	assertNotNull(resolvedRootSlot);
	assertEquals(Locality.LOCAL, slotInfoAndLocality.getLocality());
	assertEquals(rootSlot2.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());

	// occupy the slot
	resolvedRootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		groupId,
		slotInfoAndLocality.getLocality());

	slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	slotInfoAndLocality = LocationPreferenceSlotSelectionStrategy.INSTANCE.selectBestSlotForProfile(slotInfos, slotProfile).get();
	resolvedRootSlot = slotSharingManager.getResolvedRootSlot(slotInfoAndLocality.getSlotInfo());
	assertNotNull(resolvedRootSlot);
	assertNotSame(Locality.LOCAL, (slotInfoAndLocality.getLocality()));
	assertEquals(rootSlot1.getSlotRequestId(), resolvedRootSlot.getSlotRequestId());
}
 

/**
 * Tests that we can correctly retrieve resolved slots.
 */
@Test
public void testGetResolvedSlot() {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		CompletableFuture.completedFuture(
			new SimpleSlotContext(
				new AllocationID(),
				new LocalTaskManagerLocation(),
				0,
				new SimpleAckingTaskManagerGateway())),
		new SlotRequestId());

	AbstractID groupId = new AbstractID();

	Collection<SlotInfo> slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	Assert.assertEquals(1, slotInfos.size());

	SlotInfo slotInfo = slotInfos.iterator().next();
	SlotSharingManager.MultiTaskSlot resolvedMultiTaskSlot =
		slotSharingManager.getResolvedRootSlot(slotInfo);

	SlotSelectionStrategy.SlotInfoAndLocality slotInfoAndLocality =
		LocationPreferenceSlotSelectionStrategy.INSTANCE.selectBestSlotForProfile(slotInfos, SlotProfile.noRequirements()).get();

	assertNotNull(resolvedMultiTaskSlot);
	assertEquals(Locality.UNCONSTRAINED, slotInfoAndLocality.getLocality());
	assertEquals(rootSlot.getSlotRequestId(), resolvedMultiTaskSlot.getSlotRequestId());

	// occupy the resolved root slot
	resolvedMultiTaskSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		groupId,
		Locality.UNCONSTRAINED);

	slotInfos = slotSharingManager.listResolvedRootSlotInfo(groupId);
	assertTrue(slotInfos.isEmpty());
}