SlotPoolTest.java example

Explorer
flink-master
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.runtime.instance;

import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.runtime.clusterframework.types.AllocationID;
import org.apache.flink.runtime.clusterframework.types.ResourceID;
import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
import org.apache.flink.runtime.concurrent.Future;
import org.apache.flink.runtime.jobmanager.scheduler.ScheduledUnit;
import org.apache.flink.runtime.jobmanager.slots.AllocatedSlot;
import org.apache.flink.runtime.jobmanager.slots.TaskManagerGateway;
import org.apache.flink.runtime.resourcemanager.ResourceManagerGateway;
import org.apache.flink.runtime.resourcemanager.SlotRequest;
import org.apache.flink.runtime.rpc.MainThreadValidatorUtil;
import org.apache.flink.runtime.rpc.RpcService;
import org.apache.flink.runtime.rpc.TestingSerialRpcService;
import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
import org.apache.flink.util.TestLogger;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.mockito.ArgumentCaptor;

import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

import static org.apache.flink.runtime.instance.AvailableSlotsTest.DEFAULT_TESTING_PROFILE;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.RETURNS_MOCKS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public class SlotPoolTest extends TestLogger {

	private RpcService rpcService;

	private JobID jobId;

	private MainThreadValidatorUtil mainThreadValidatorUtil;

	private SlotPool slotPool;

	private ResourceManagerGateway resourceManagerGateway;

	@Before
	public void setUp() throws Exception {

		this.rpcService = new TestingSerialRpcService();
		this.jobId = new JobID();
		this.slotPool = new SlotPool(rpcService, jobId);

		this.mainThreadValidatorUtil = new MainThreadValidatorUtil(slotPool);

		mainThreadValidatorUtil.enterMainThread();

		final String jobManagerAddress = "foobar";

		slotPool.start(UUID.randomUUID(), jobManagerAddress);

		this.resourceManagerGateway = mock(ResourceManagerGateway.class);
		when(resourceManagerGateway
			.requestSlot(any(UUID.class), any(UUID.class), any(SlotRequest.class), any(Time.class)))
			.thenReturn(mock(Future.class, RETURNS_MOCKS));

		slotPool.connectToResourceManager(UUID.randomUUID(), resourceManagerGateway);
	}

	@After
	public void tearDown() throws Exception {
		mainThreadValidatorUtil.exitMainThread();
	}

	@Test
	public void testAllocateSimpleSlot() throws Exception {
		ResourceID resourceID = new ResourceID("resource");
		slotPool.registerTaskManager(resourceID);

		ScheduledUnit task = mock(ScheduledUnit.class);
		Future<SimpleSlot> future = slotPool.allocateSlot(task, DEFAULT_TESTING_PROFILE, null);
		assertFalse(future.isDone());

		ArgumentCaptor<SlotRequest> slotRequestArgumentCaptor = ArgumentCaptor.forClass(SlotRequest.class);
		verify(resourceManagerGateway).requestSlot(any(UUID.class), any(UUID.class), slotRequestArgumentCaptor.capture(), any(Time.class));

		final SlotRequest slotRequest = slotRequestArgumentCaptor.getValue();

		AllocatedSlot allocatedSlot = createAllocatedSlot(resourceID, slotRequest.getAllocationId(), jobId, DEFAULT_TESTING_PROFILE);
		assertTrue(slotPool.offerSlot(allocatedSlot));

		SimpleSlot slot = future.get(1, TimeUnit.SECONDS);
		assertTrue(future.isDone());
		assertTrue(slot.isAlive());
		assertEquals(resourceID, slot.getTaskManagerID());
		assertEquals(jobId, slot.getJobID());
		assertEquals(slotPool.getSlotOwner(), slot.getOwner());
		assertEquals(slotPool.getAllocatedSlots().get(slot.getAllocatedSlot().getSlotAllocationId()), slot);
	}

	@Test
	public void testAllocationFulfilledByReturnedSlot() throws Exception {
		ResourceID resourceID = new ResourceID("resource");
		slotPool.registerTaskManager(resourceID);

		Future<SimpleSlot> future1 = slotPool.allocateSlot(mock(ScheduledUnit.class),DEFAULT_TESTING_PROFILE, null);
		Future<SimpleSlot> future2 = slotPool.allocateSlot(mock(ScheduledUnit.class),DEFAULT_TESTING_PROFILE, null);

		assertFalse(future1.isDone());
		assertFalse(future2.isDone());

		ArgumentCaptor<SlotRequest> slotRequestArgumentCaptor = ArgumentCaptor.forClass(SlotRequest.class);
		verify(resourceManagerGateway, times(2))
			.requestSlot(any(UUID.class), any(UUID.class), slotRequestArgumentCaptor.capture(), any(Time.class));

		final List<SlotRequest> slotRequests = slotRequestArgumentCaptor.getAllValues();

		AllocatedSlot allocatedSlot = createAllocatedSlot(resourceID, slotRequests.get(0).getAllocationId(), jobId, DEFAULT_TESTING_PROFILE);
		assertTrue(slotPool.offerSlot(allocatedSlot));

		SimpleSlot slot1 = future1.get(1, TimeUnit.SECONDS);
		assertTrue(future1.isDone());
		assertFalse(future2.isDone());

		// return this slot to pool
		slot1.releaseSlot();

		// second allocation fulfilled by previous slot returning
		SimpleSlot slot2 = future2.get(1, TimeUnit.SECONDS);
		assertTrue(future2.isDone());

		assertNotEquals(slot1, slot2);
		assertTrue(slot1.isReleased());
		assertTrue(slot2.isAlive());
		assertEquals(slot1.getTaskManagerID(), slot2.getTaskManagerID());
		assertEquals(slot1.getSlotNumber(), slot2.getSlotNumber());
		assertEquals(slotPool.getAllocatedSlots().get(slot1.getAllocatedSlot().getSlotAllocationId()), slot2);
	}

	@Test
	public void testAllocateWithFreeSlot() throws Exception {
		ResourceID resourceID = new ResourceID("resource");
		slotPool.registerTaskManager(resourceID);

		Future<SimpleSlot> future1 = slotPool.allocateSlot(mock(ScheduledUnit.class),DEFAULT_TESTING_PROFILE, null);
		assertFalse(future1.isDone());

		ArgumentCaptor<SlotRequest> slotRequestArgumentCaptor = ArgumentCaptor.forClass(SlotRequest.class);
		verify(resourceManagerGateway).requestSlot(any(UUID.class), any(UUID.class), slotRequestArgumentCaptor.capture(), any(Time.class));

		final SlotRequest slotRequest = slotRequestArgumentCaptor.getValue();

		AllocatedSlot allocatedSlot = createAllocatedSlot(resourceID, slotRequest.getAllocationId(), jobId, DEFAULT_TESTING_PROFILE);
		assertTrue(slotPool.offerSlot(allocatedSlot));

		SimpleSlot slot1 = future1.get(1, TimeUnit.SECONDS);
		assertTrue(future1.isDone());

		// return this slot to pool
		slot1.releaseSlot();

		Future<SimpleSlot> future2 = slotPool.allocateSlot(mock(ScheduledUnit.class),DEFAULT_TESTING_PROFILE, null);

		// second allocation fulfilled by previous slot returning
		SimpleSlot slot2 = future2.get(1, TimeUnit.SECONDS);
		assertTrue(future2.isDone());

		assertNotEquals(slot1, slot2);
		assertTrue(slot1.isReleased());
		assertTrue(slot2.isAlive());
		assertEquals(slot1.getTaskManagerID(), slot2.getTaskManagerID());
		assertEquals(slot1.getSlotNumber(), slot2.getSlotNumber());
	}

	@Test
	public void testOfferSlot() throws Exception {
		ResourceID resourceID = new ResourceID("resource");
		slotPool.registerTaskManager(resourceID);

		Future<SimpleSlot> future = slotPool.allocateSlot(mock(ScheduledUnit.class),DEFAULT_TESTING_PROFILE, null);
		assertFalse(future.isDone());

		ArgumentCaptor<SlotRequest> slotRequestArgumentCaptor = ArgumentCaptor.forClass(SlotRequest.class);
		verify(resourceManagerGateway).requestSlot(any(UUID.class), any(UUID.class), slotRequestArgumentCaptor.capture(), any(Time.class));

		final SlotRequest slotRequest = slotRequestArgumentCaptor.getValue();

		// slot from unregistered resource
		AllocatedSlot invalid = createAllocatedSlot(new ResourceID("unregistered"), slotRequest.getAllocationId(), jobId, DEFAULT_TESTING_PROFILE);
		assertFalse(slotPool.offerSlot(invalid));

		AllocatedSlot notRequested = createAllocatedSlot(resourceID, new AllocationID(), jobId, DEFAULT_TESTING_PROFILE);

		// we'll also accept non requested slots
		assertTrue(slotPool.offerSlot(notRequested));

		AllocatedSlot allocatedSlot = createAllocatedSlot(resourceID, slotRequest.getAllocationId(), jobId, DEFAULT_TESTING_PROFILE);

		// accepted slot
		assertTrue(slotPool.offerSlot(allocatedSlot));
		SimpleSlot slot = future.get(1, TimeUnit.SECONDS);
		assertTrue(future.isDone());
		assertTrue(slot.isAlive());

		// duplicated offer with using slot
		assertTrue(slotPool.offerSlot(allocatedSlot));
		assertTrue(future.isDone());
		assertTrue(slot.isAlive());

		// duplicated offer with free slot
		slot.releaseSlot();
		assertTrue(slot.isReleased());
		assertTrue(slotPool.offerSlot(allocatedSlot));
	}

	@Test
	public void testReleaseResource() throws Exception {
		ResourceID resourceID = new ResourceID("resource");
		slotPool.registerTaskManager(resourceID);

		Future<SimpleSlot> future1 = slotPool.allocateSlot(mock(ScheduledUnit.class),DEFAULT_TESTING_PROFILE, null);

		ArgumentCaptor<SlotRequest> slotRequestArgumentCaptor = ArgumentCaptor.forClass(SlotRequest.class);
		verify(resourceManagerGateway).requestSlot(any(UUID.class), any(UUID.class), slotRequestArgumentCaptor.capture(), any(Time.class));

		final SlotRequest slotRequest = slotRequestArgumentCaptor.getValue();

		Future<SimpleSlot> future2 = slotPool.allocateSlot(mock(ScheduledUnit.class),DEFAULT_TESTING_PROFILE, null);

		AllocatedSlot allocatedSlot = createAllocatedSlot(resourceID, slotRequest.getAllocationId(), jobId, DEFAULT_TESTING_PROFILE);
		assertTrue(slotPool.offerSlot(allocatedSlot));

		SimpleSlot slot1 = future1.get(1, TimeUnit.SECONDS);
		assertTrue(future1.isDone());
		assertFalse(future2.isDone());

		slotPool.releaseTaskManager(resourceID);
		assertTrue(slot1.isReleased());

		// slot released and not usable, second allocation still not fulfilled
		Thread.sleep(10);
		assertFalse(future2.isDone());
	}

	static AllocatedSlot createAllocatedSlot(
			final ResourceID resourceId,
			final AllocationID allocationId,
			final JobID jobId,
			final ResourceProfile resourceProfile) {
		TaskManagerLocation mockTaskManagerLocation = mock(TaskManagerLocation.class);
		when(mockTaskManagerLocation.getResourceID()).thenReturn(resourceId);

		TaskManagerGateway mockTaskManagerGateway = mock(TaskManagerGateway.class);

		return new AllocatedSlot(
			allocationId,
			jobId,
			mockTaskManagerLocation,
			0,
			resourceProfile,
			mockTaskManagerGateway);
	}

}