@Test public void testStopWorker() throws Exception { new Context() {{ runTest(() -> { CompletableFuture<?> registerSlotRequestFuture = resourceManager.runInMainThread(() -> { rmServices.slotManager.registerSlotRequest( new SlotRequest(new JobID(), new AllocationID(), resourceProfile1, taskHost)); return null; }); registerSlotRequestFuture.get(); Container testingContainer = mockContainer("container", 1234, 1, resourceManager.getContainerResource()); doReturn(Collections.singletonList(Collections.singletonList(resourceManager.getContainerRequest()))) .when(mockResourceManagerClient).getMatchingRequests(any(Priority.class), anyString(), any(Resource.class)); resourceManager.onContainersAllocated(ImmutableList.of(testingContainer)); verify(mockResourceManagerClient).addContainerRequest(any(AMRMClient.ContainerRequest.class)); verify(mockNMClient).startContainer(eq(testingContainer), any(ContainerLaunchContext.class)); TaskExecutorGateway mockTaskExecutorGateway = mock(TaskExecutorGateway.class); rpcService.registerGateway(taskHost, mockTaskExecutorGateway); final ResourceManagerGateway rmGateway = resourceManager.getSelfGateway(ResourceManagerGateway.class); final ResourceID taskManagerResourceId = new ResourceID(testingContainer.getId().toString()); final SlotReport slotReport = new SlotReport( new SlotStatus( new SlotID(taskManagerResourceId, 1), new ResourceProfile(10, 1, 1, 1, 0, Collections.emptyMap()))); CompletableFuture<Integer> numberRegisteredSlotsFuture = rmGateway .registerTaskExecutor( taskHost, taskManagerResourceId, dataPort, hardwareDescription, Time.seconds(10L)) .thenCompose( (RegistrationResponse response) -> { assertThat(response, instanceOf(TaskExecutorRegistrationSuccess.class)); final TaskExecutorRegistrationSuccess success = (TaskExecutorRegistrationSuccess) response; return rmGateway.sendSlotReport( taskManagerResourceId, success.getRegistrationId(), slotReport, Time.seconds(10L)); }) .handleAsync( (Acknowledge ignored, Throwable throwable) -> rmServices.slotManager.getNumberRegisteredSlots(), resourceManager.getMainThreadExecutorForTesting()); final int numberRegisteredSlots = numberRegisteredSlotsFuture.get(); assertEquals(1, numberRegisteredSlots); CompletableFuture<?> unregisterAndReleaseFuture = resourceManager.runInMainThread(() -> { rmServices.slotManager.unregisterTaskManagersAndReleaseResources(); return null; }); unregisterAndReleaseFuture.get(); verify(mockNMClient).stopContainer(any(ContainerId.class), any(NodeId.class)); verify(mockResourceManagerClient).releaseAssignedContainer(any(ContainerId.class)); }); assertThat(rmServices.slotManager.getNumberRegisteredSlots(), Matchers.equalTo(0)); assertThat(resourceManager.getNumberOfRegisteredTaskManagers().get(), Matchers.equalTo(0)); }}; }
@Override public boolean stopWorker(final YarnWorkerNode workerNode) { final Container container = workerNode.getContainer(); log.info("Stopping container {}.", container.getId()); try { nodeManagerClient.stopContainer(container.getId(), container.getNodeId()); } catch (final Exception e) { log.warn("Error while calling YARN Node Manager to stop container", e); } resourceManagerClient.releaseAssignedContainer(container.getId()); workerNodeMap.remove(workerNode.getResourceID()); return true; }
YarnResourceManager extends ResourceManager<YarnWorkerNode> implements AMRMClientAsync.CallbackHandler { @Override public boolean stopWorker(final YarnWorkerNode workerNode) { final Container container = workerNode.getContainer(); log.info("Stopping container {}.", container.getId()); try { nodeManagerClient.stopContainer(container.getId(), container.getNodeId()); } catch (final Exception e) { log.warn("Error while calling YARN Node Manager to stop container", e); } resourceManagerClient.releaseAssignedContainer(container.getId()); workerNodeMap.remove(workerNode.getResourceID()); return true; } }
YarnResourceManager extends ResourceManager<YarnWorkerNode> implements AMRMClientAsync.CallbackHandler { @Override public boolean stopWorker(final YarnWorkerNode workerNode) { final Container container = workerNode.getContainer(); log.info("Stopping container {}.", container.getId()); try { nodeManagerClient.stopContainer(container.getId(), container.getNodeId()); } catch (final Exception e) { log.warn("Error while calling YARN Node Manager to stop container", e); } resourceManagerClient.releaseAssignedContainer(container.getId()); workerNodeMap.remove(workerNode.getResourceID()); return true; } YarnResourceManager(
			RpcService rpcService,
			String resourceManagerEndpointId,
			ResourceID resourceId,
			Configuration flinkConfig,
			Map<String, String> env,
			HighAvailabilityServices highAvailabilityServices,
			HeartbeatServices heartbeatServices,
			SlotManager slotManager,
			MetricRegistry metricRegistry,
			JobLeaderIdService jobLeaderIdService,
			ClusterInformation clusterInformation,
			FatalErrorHandler fatalErrorHandler,
			@Nullable String webInterfaceUrl,
			JobManagerMetricGroup jobManagerMetricGroup); }
YarnResourceManager extends ResourceManager<YarnWorkerNode> implements AMRMClientAsync.CallbackHandler { @Override public boolean stopWorker(final YarnWorkerNode workerNode) { final Container container = workerNode.getContainer(); log.info("Stopping container {}.", container.getId()); try { nodeManagerClient.stopContainer(container.getId(), container.getNodeId()); } catch (final Exception e) { log.warn("Error while calling YARN Node Manager to stop container", e); } resourceManagerClient.releaseAssignedContainer(container.getId()); workerNodeMap.remove(workerNode.getResourceID()); return true; } YarnResourceManager(
			RpcService rpcService,
			String resourceManagerEndpointId,
			ResourceID resourceId,
			Configuration flinkConfig,
			Map<String, String> env,
			HighAvailabilityServices highAvailabilityServices,
			HeartbeatServices heartbeatServices,
			SlotManager slotManager,
			MetricRegistry metricRegistry,
			JobLeaderIdService jobLeaderIdService,
			ClusterInformation clusterInformation,
			FatalErrorHandler fatalErrorHandler,
			@Nullable String webInterfaceUrl,
			JobManagerMetricGroup jobManagerMetricGroup); @Override CompletableFuture<Void> onStop(); @Override Collection<ResourceProfile> startNewWorker(ResourceProfile resourceProfile); @Override boolean stopWorker(final YarnWorkerNode workerNode); @Override float getProgress(); @Override void onContainersCompleted(final List<ContainerStatus> statuses); @Override void onContainersAllocated(List<Container> containers); @Override void onShutdownRequest(); @Override void onNodesUpdated(List<NodeReport> list); @Override void onError(Throwable error); }
YarnResourceManager extends ResourceManager<YarnWorkerNode> implements AMRMClientAsync.CallbackHandler { @Override public boolean stopWorker(final YarnWorkerNode workerNode) { final Container container = workerNode.getContainer(); log.info("Stopping container {}.", container.getId()); try { nodeManagerClient.stopContainer(container.getId(), container.getNodeId()); } catch (final Exception e) { log.warn("Error while calling YARN Node Manager to stop container", e); } resourceManagerClient.releaseAssignedContainer(container.getId()); workerNodeMap.remove(workerNode.getResourceID()); return true; } YarnResourceManager(
			RpcService rpcService,
			String resourceManagerEndpointId,
			ResourceID resourceId,
			Configuration flinkConfig,
			Map<String, String> env,
			HighAvailabilityServices highAvailabilityServices,
			HeartbeatServices heartbeatServices,
			SlotManager slotManager,
			MetricRegistry metricRegistry,
			JobLeaderIdService jobLeaderIdService,
			ClusterInformation clusterInformation,
			FatalErrorHandler fatalErrorHandler,
			@Nullable String webInterfaceUrl,
			JobManagerMetricGroup jobManagerMetricGroup); @Override CompletableFuture<Void> onStop(); @Override Collection<ResourceProfile> startNewWorker(ResourceProfile resourceProfile); @Override boolean stopWorker(final YarnWorkerNode workerNode); @Override float getProgress(); @Override void onContainersCompleted(final List<ContainerStatus> statuses); @Override void onContainersAllocated(List<Container> containers); @Override void onShutdownRequest(); @Override void onNodesUpdated(List<NodeReport> list); @Override void onError(Throwable error);  }