TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { @Override public void start() throws Exception { super.start(); try { haServices.getResourceManagerLeaderRetriever().start(new ResourceManagerLeaderListener()); } catch (Exception e) { onFatalError(e); } taskSlotTable.start(new SlotActionsImpl()); jobLeaderService.start(getAddress(), getRpcService(), haServices, new JobLeaderListenerImpl()); } TaskExecutor(
			RpcService rpcService,
			TaskManagerConfiguration taskManagerConfiguration,
			TaskManagerLocation taskManagerLocation,
			MemoryManager memoryManager,
			IOManager ioManager,
			NetworkEnvironment networkEnvironment,
			HighAvailabilityServices haServices,
			HeartbeatServices heartbeatServices,
			TaskManagerMetricGroup taskManagerMetricGroup,
			BroadcastVariableManager broadcastVariableManager,
			FileCache fileCache,
			TaskSlotTable taskSlotTable,
			JobManagerTable jobManagerTable,
			JobLeaderService jobLeaderService,
			FatalErrorHandler fatalErrorHandler); @Override void start(); @Override void postStop(); @Override CompletableFuture<Acknowledge> submitTask(
			TaskDeploymentDescriptor tdd,
			JobMasterId jobMasterId,
			Time timeout); @Override CompletableFuture<Acknowledge> cancelTask(ExecutionAttemptID executionAttemptID, Time timeout); @Override CompletableFuture<Acknowledge> stopTask(ExecutionAttemptID executionAttemptID, Time timeout); @Override CompletableFuture<Acknowledge> updatePartitions(
			final ExecutionAttemptID executionAttemptID,
			Iterable<PartitionInfo> partitionInfos,
			Time timeout); @Override void failPartition(ExecutionAttemptID executionAttemptID); @Override void heartbeatFromJobManager(ResourceID resourceID); @Override void heartbeatFromResourceManager(ResourceID resourceID); @Override CompletableFuture<Acknowledge> triggerCheckpoint(
			ExecutionAttemptID executionAttemptID,
			long checkpointId,
			long checkpointTimestamp,
			CheckpointOptions checkpointOptions); @Override CompletableFuture<Acknowledge> confirmCheckpoint(
			ExecutionAttemptID executionAttemptID,
			long checkpointId,
			long checkpointTimestamp); @Override CompletableFuture<Acknowledge> requestSlot(
		final SlotID slotId,
		final JobID jobId,
		final AllocationID allocationId,
		final String targetAddress,
		final ResourceManagerId resourceManagerId,
		final Time timeout); @Override void disconnectJobManager(JobID jobId, Exception cause); @Override void disconnectResourceManager(Exception cause); ResourceID getResourceID(); static final String TASK_MANAGER_NAME; static final int MAX_BLOB_PORT; }
@Test public void testSlotAcceptance() throws Exception { final JobID jobId = new JobID(); final Configuration configuration = new Configuration(); final TaskManagerConfiguration taskManagerConfiguration = TaskManagerConfiguration.fromConfiguration(configuration); final ResourceID resourceId = new ResourceID("foobar"); final TaskManagerLocation taskManagerLocation = new TaskManagerLocation(resourceId, InetAddress.getLoopbackAddress(), 1234); final TestingHighAvailabilityServices haServices = new TestingHighAvailabilityServices(); final TimerService<AllocationID> timerService = mock(TimerService.class); final TaskSlotTable taskSlotTable = new TaskSlotTable(Arrays.asList(mock(ResourceProfile.class), mock(ResourceProfile.class)), timerService); final JobManagerTable jobManagerTable = new JobManagerTable(); final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation); final TestingFatalErrorHandler testingFatalErrorHandler = new TestingFatalErrorHandler(); final String resourceManagerAddress = "rm"; final UUID resourceManagerLeaderId = UUID.randomUUID(); final ResourceID resourceManagerResourceId = new ResourceID(resourceManagerAddress); final String jobManagerAddress = "jm"; final UUID jobManagerLeaderId = UUID.randomUUID(); final LeaderRetrievalService resourceManagerLeaderRetrievalService = new TestingLeaderRetrievalService(resourceManagerAddress, resourceManagerLeaderId); final LeaderRetrievalService jobManagerLeaderRetrievalService = new TestingLeaderRetrievalService(jobManagerAddress, jobManagerLeaderId); haServices.setResourceManagerLeaderRetriever(resourceManagerLeaderRetrievalService); haServices.setJobMasterLeaderRetriever(jobId, jobManagerLeaderRetrievalService); final ResourceManagerGateway resourceManagerGateway = mock(ResourceManagerGateway.class); final InstanceID registrationId = new InstanceID(); when(resourceManagerGateway.registerTaskExecutor( any(String.class), eq(resourceId), any(SlotReport.class), any(Time.class))).thenReturn(CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(registrationId, resourceManagerResourceId, 1000L))); final ResourceID jmResourceId = new ResourceID(jobManagerAddress); final int blobPort = 42; final AllocationID allocationId1 = new AllocationID(); final AllocationID allocationId2 = new AllocationID(); final SlotOffer offer1 = new SlotOffer(allocationId1, 0, ResourceProfile.UNKNOWN); final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class); when(jobMasterGateway.registerTaskManager( any(String.class), eq(taskManagerLocation), any(Time.class) )).thenReturn(CompletableFuture.completedFuture(new JMTMRegistrationSuccess(jmResourceId, blobPort))); when(jobMasterGateway.getHostname()).thenReturn(jobManagerAddress); when(jobMasterGateway.offerSlots( any(ResourceID.class), any(Iterable.class), any(Time.class))) .thenReturn(CompletableFuture.completedFuture((Collection<SlotOffer>)Collections.singleton(offer1))); rpc.registerGateway(resourceManagerAddress, resourceManagerGateway); rpc.registerGateway(jobManagerAddress, jobMasterGateway); TaskExecutor taskManager = new TaskExecutor( rpc, taskManagerConfiguration, taskManagerLocation, mock(MemoryManager.class), mock(IOManager.class), mock(NetworkEnvironment.class), haServices, mock(HeartbeatServices.class, RETURNS_MOCKS), mock(TaskManagerMetricGroup.class), mock(BroadcastVariableManager.class), mock(FileCache.class), taskSlotTable, jobManagerTable, jobLeaderService, testingFatalErrorHandler); try { taskManager.start(); verify(resourceManagerGateway, Mockito.timeout(timeout.toMilliseconds())).registerTaskExecutor( eq(taskManager.getAddress()), eq(resourceId), any(SlotReport.class), any(Time.class)); taskSlotTable.allocateSlot(0, jobId, allocationId1, Time.milliseconds(10000L)); taskSlotTable.allocateSlot(1, jobId, allocationId2, Time.milliseconds(10000L)); jobLeaderService.addJob(jobId, jobManagerAddress); verify(resourceManagerGateway, Mockito.timeout(timeout.toMilliseconds())).notifySlotAvailable( eq(registrationId), eq(new SlotID(resourceId, 1)), eq(allocationId2)); assertTrue(taskSlotTable.existsActiveSlot(jobId, allocationId1)); assertFalse(taskSlotTable.existsActiveSlot(jobId, allocationId2)); assertTrue(taskSlotTable.isSlotFree(1)); testingFatalErrorHandler.rethrowError(); } finally { taskManager.shutDown(); taskManager.getTerminationFuture().get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS); } }
@Test public void testFilterOutDuplicateJobMasterRegistrations() throws Exception { final long verificationTimeout = 500L; final Configuration configuration = new Configuration(); final TestingFatalErrorHandler testingFatalErrorHandler = new TestingFatalErrorHandler(); final JobLeaderService jobLeaderService = mock(JobLeaderService.class); final TaskManagerConfiguration taskManagerConfiguration = TaskManagerConfiguration.fromConfiguration(configuration); final TaskManagerLocation taskManagerLocation = new TaskManagerLocation(ResourceID.generate(), InetAddress.getLocalHost(), 1234); final HighAvailabilityServices haServicesMock = mock(HighAvailabilityServices.class, Mockito.RETURNS_MOCKS); final HeartbeatServices heartbeatServicesMock = mock(HeartbeatServices.class, Mockito.RETURNS_MOCKS); final JobID jobId = new JobID(); final JobMasterGateway jobMasterGateway = mock(JobMasterGateway.class); when(jobMasterGateway.getHostname()).thenReturn("localhost"); final JMTMRegistrationSuccess registrationMessage = new JMTMRegistrationSuccess(ResourceID.generate(), 1); final JobManagerTable jobManagerTableMock = spy(new JobManagerTable()); final TaskExecutor taskExecutor = new TaskExecutor( rpc, taskManagerConfiguration, taskManagerLocation, mock(MemoryManager.class), mock(IOManager.class), mock(NetworkEnvironment.class), haServicesMock, heartbeatServicesMock, mock(TaskManagerMetricGroup.class), mock(BroadcastVariableManager.class), mock(FileCache.class), mock(TaskSlotTable.class), jobManagerTableMock, jobLeaderService, testingFatalErrorHandler); try { taskExecutor.start(); ArgumentCaptor<JobLeaderListener> jobLeaderListenerArgumentCaptor = ArgumentCaptor.forClass(JobLeaderListener.class); verify(jobLeaderService).start(anyString(), any(RpcService.class), any(HighAvailabilityServices.class), jobLeaderListenerArgumentCaptor.capture()); JobLeaderListener taskExecutorListener = jobLeaderListenerArgumentCaptor.getValue(); taskExecutorListener.jobManagerGainedLeadership(jobId, jobMasterGateway, registrationMessage); taskExecutorListener.jobManagerGainedLeadership(jobId, jobMasterGateway, registrationMessage); ArgumentCaptor<JobManagerConnection> jobManagerConnectionArgumentCaptor = ArgumentCaptor.forClass(JobManagerConnection.class); verify(jobManagerTableMock, Mockito.timeout(verificationTimeout).times(1)).put(eq(jobId), jobManagerConnectionArgumentCaptor.capture()); JobManagerConnection jobManagerConnection = jobManagerConnectionArgumentCaptor.getValue(); assertEquals(jobMasterGateway, jobManagerConnection.getJobManagerGateway()); testingFatalErrorHandler.rethrowError(); } finally { taskExecutor.shutDown(); taskExecutor.getTerminationFuture().get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS); } }