Java Code Examples for org.apache.hadoop.yarn.ipc.YarnRPC#create()

protected MRClientProtocol instantiateHistoryProxy(final Configuration conf,
    final InetSocketAddress hsAddress) throws IOException {

  if (LOG.isDebugEnabled()) {
    LOG.debug("Connecting to MRHistoryServer at: " + hsAddress);
  }
  final YarnRPC rpc = YarnRPC.create(conf);
  UserGroupInformation currentUser = UserGroupInformation.getCurrentUser();
  return currentUser.doAs(new PrivilegedAction<MRClientProtocol>() {
    @Override
    public MRClientProtocol run() {
      return (MRClientProtocol) rpc.getProxy(HSClientProtocol.class,
          hsAddress, conf);
    }
  });
}
 

@Override
public void init(Configuration conf) {
  systemConf = (TajoConf)conf;

  yarnRPC = YarnRPC.create(systemConf);

  connectYarnClient();

  taskRunnerLauncher = new YarnTaskRunnerLauncherImpl(queryTaskContext, yarnRPC);
  addService((Service) taskRunnerLauncher);
  queryTaskContext.getDispatcher().register(TaskRunnerGroupEvent.EventType.class, taskRunnerLauncher);

  rmAllocator = new YarnRMContainerAllocator(queryTaskContext);
  addService(rmAllocator);
  queryTaskContext.getDispatcher().register(ContainerAllocatorEventType.class, rmAllocator);
  super.init(conf);
}
 

@Override
protected void serviceStart() throws Exception {
  Configuration conf = getConfig();
  YarnRPC rpc = YarnRPC.create(conf);
  this.server =
      rpc.getServer(SCMAdminProtocol.class, this,
          clientBindAddress,
          conf, null, // Secret manager null for now (security not supported)
          conf.getInt(YarnConfiguration.SCM_ADMIN_CLIENT_THREAD_COUNT,
              YarnConfiguration.DEFAULT_SCM_ADMIN_CLIENT_THREAD_COUNT));

  // TODO: Enable service authorization (see YARN-2774)

  this.server.start();
  clientBindAddress =
      conf.updateConnectAddr(YarnConfiguration.SCM_ADMIN_ADDRESS,
          server.getListenerAddress());

  super.serviceStart();
}
 

private MRClientProtocol getMRClientProtocol(Token token,
    final InetSocketAddress hsAddress, String user, final Configuration conf) {
  UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);
  ugi.addToken(ConverterUtils.convertFromYarn(token, hsAddress));

  final YarnRPC rpc = YarnRPC.create(conf);
  MRClientProtocol hsWithDT = ugi
      .doAs(new PrivilegedAction<MRClientProtocol>() {

        @Override
        public MRClientProtocol run() {
          return (MRClientProtocol) rpc.getProxy(HSClientProtocol.class,
              hsAddress, conf);
        }
      });
  return hsWithDT;
}
 

@Before
public void startUp() {
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.SCM_STORE_CLASS,
      InMemorySCMStore.class.getName());
  conf.set(YarnConfiguration.SHARED_CACHE_ROOT, testDir.getPath());
  AppChecker appChecker = spy(new DummyAppChecker());
  store = new InMemorySCMStore(appChecker);
  store.init(conf);
  store.start();

  service = new ClientProtocolService(store);
  service.init(conf);
  service.start();

  YarnRPC rpc = YarnRPC.create(new Configuration());

  InetSocketAddress scmAddress =
      conf.getSocketAddr(YarnConfiguration.SCM_CLIENT_SERVER_ADDRESS,
          YarnConfiguration.DEFAULT_SCM_CLIENT_SERVER_ADDRESS,
          YarnConfiguration.DEFAULT_SCM_CLIENT_SERVER_PORT);

  clientSCMProxy =
      (ClientSCMProtocol) rpc.getProxy(ClientSCMProtocol.class, scmAddress,
          conf);
}
 

@Override
protected ContainerLauncher
    createContainerLauncher(final AppContext context) {
  return new ContainerLauncherImpl(context) {

    @Override
    public ContainerManagementProtocolProxyData getCMProxy(
        String containerMgrBindAddr, ContainerId containerId)
        throws IOException {
      InetSocketAddress addr = NetUtils.getConnectAddress(server);
      String containerManagerBindAddr =
          addr.getHostName() + ":" + addr.getPort();
      Token token =
          tokenSecretManager.createNMToken(
            containerId.getApplicationAttemptId(),
            NodeId.newInstance(addr.getHostName(), addr.getPort()), "user");
      ContainerManagementProtocolProxy cmProxy =
          new ContainerManagementProtocolProxy(conf);
      ContainerManagementProtocolProxyData proxy =
          cmProxy.new ContainerManagementProtocolProxyData(
            YarnRPC.create(conf), containerManagerBindAddr, containerId,
            token);
      return proxy;
    }
  };

}
 

@Override
protected void serviceStart() throws Exception {
  Configuration conf = getConfig();
  YarnRPC rpc = YarnRPC.create(conf);
  this.server =   
    rpc.getServer(ApplicationClientProtocol.class, this,
          clientBindAddress,
          conf, this.rmDTSecretManager,
          conf.getInt(YarnConfiguration.RM_CLIENT_THREAD_COUNT, 
              YarnConfiguration.DEFAULT_RM_CLIENT_THREAD_COUNT));
  
  // Enable service authorization?
  if (conf.getBoolean(
      CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHORIZATION, 
      false)) {
    InputStream inputStream =
        this.rmContext.getConfigurationProvider()
            .getConfigurationInputStream(conf,
                YarnConfiguration.HADOOP_POLICY_CONFIGURATION_FILE);
    if (inputStream != null) {
      conf.addResource(inputStream);
    }
    refreshServiceAcls(conf, RMPolicyProvider.getInstance());
  }
  
  this.server.start();
  clientBindAddress = conf.updateConnectAddr(YarnConfiguration.RM_BIND_HOST,
                                             YarnConfiguration.RM_ADDRESS,
                                             YarnConfiguration.DEFAULT_RM_ADDRESS,
                                             server.getListenerAddress());
  super.serviceStart();
}
 

private HSClientProtocol instantiateHistoryProxy() {
  final String serviceAddr =
      mrCluster.getConfig().get(JHAdminConfig.MR_HISTORY_ADDRESS);
  final YarnRPC rpc = YarnRPC.create(conf);
  HSClientProtocol historyClient =
      (HSClientProtocol) rpc.getProxy(HSClientProtocol.class,
          NetUtils.createSocketAddr(serviceAddr), mrCluster.getConfig());
  return historyClient;
}
 

@Override
protected GetUserMappingsProtocol getUgmProtocol() throws IOException {
  Configuration conf = getConf();
  
  final InetSocketAddress addr = conf.getSocketAddr(
      YarnConfiguration.RM_ADMIN_ADDRESS,
      YarnConfiguration.DEFAULT_RM_ADMIN_ADDRESS,
      YarnConfiguration.DEFAULT_RM_ADMIN_PORT);
  final YarnRPC rpc = YarnRPC.create(conf);
  
  ResourceManagerAdministrationProtocol adminProtocol = (ResourceManagerAdministrationProtocol) rpc.getProxy(
      ResourceManagerAdministrationProtocol.class, addr, getConf());

  return adminProtocol;
}
 

@Override
protected void serviceStart() throws Exception {
  Configuration conf = getConfig();
  YarnRPC rpc = YarnRPC.create(conf);
  this.server =   
    rpc.getServer(ApplicationClientProtocol.class, this,
          clientBindAddress,
          conf, this.rmDTSecretManager,
          conf.getInt(YarnConfiguration.RM_CLIENT_THREAD_COUNT, 
              YarnConfiguration.DEFAULT_RM_CLIENT_THREAD_COUNT));
  
  // Enable service authorization?
  if (conf.getBoolean(
      CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHORIZATION, 
      false)) {
    InputStream inputStream =
        this.rmContext.getConfigurationProvider()
            .getConfigurationInputStream(conf,
                YarnConfiguration.HADOOP_POLICY_CONFIGURATION_FILE);
    if (inputStream != null) {
      conf.addResource(inputStream);
    }
    refreshServiceAcls(conf, RMPolicyProvider.getInstance());
  }
  
  this.server.start();
  clientBindAddress = conf.updateConnectAddr(YarnConfiguration.RM_BIND_HOST,
                                             YarnConfiguration.RM_ADDRESS,
                                             YarnConfiguration.DEFAULT_RM_ADDRESS,
                                             server.getListenerAddress());
  super.serviceStart();
}
 

private HSClientProtocol instantiateHistoryProxy() {
  final String serviceAddr =
      mrCluster.getConfig().get(JHAdminConfig.MR_HISTORY_ADDRESS);
  final YarnRPC rpc = YarnRPC.create(conf);
  HSClientProtocol historyClient =
      (HSClientProtocol) rpc.getProxy(HSClientProtocol.class,
          NetUtils.createSocketAddr(serviceAddr), mrCluster.getConfig());
  return historyClient;
}
 

@Override
protected void serviceStart() throws Exception {
  // All the clients to appsManager are supposed to be authenticated via
  // Kerberos if security is enabled, so no secretManager.
  YarnRPC rpc = YarnRPC.create(getConfig());
  Configuration clientServerConf = new Configuration(getConfig());
  this.server = rpc.getServer(ApplicationClientProtocol.class, this,
      clientBindAddress, clientServerConf, null, 1);
  this.server.start();
  super.serviceStart();
}
 

public ContainerManagementProtocolProxy(Configuration conf,
    NMTokenCache nmTokenCache) {
  this.conf = new Configuration(conf);
  this.nmTokenCache = nmTokenCache;

  maxConnectedNMs =
      conf.getInt(YarnConfiguration.NM_CLIENT_MAX_NM_PROXIES,
          YarnConfiguration.DEFAULT_NM_CLIENT_MAX_NM_PROXIES);
  if (maxConnectedNMs < 0) {
    throw new YarnRuntimeException(
        YarnConfiguration.NM_CLIENT_MAX_NM_PROXIES
            + " (" + maxConnectedNMs + ") can not be less than 0.");
  }
  LOG.info(YarnConfiguration.NM_CLIENT_MAX_NM_PROXIES + " : "
      + maxConnectedNMs);

  if (maxConnectedNMs > 0) {
    cmProxy =
        new LinkedHashMap<String, ContainerManagementProtocolProxyData>();
  } else {
    cmProxy = Collections.emptyMap();
    // Connections are not being cached so ensure connections close quickly
    // to avoid creating thousands of RPC client threads on large clusters.
    this.conf.setInt(
        CommonConfigurationKeysPublic.IPC_CLIENT_CONNECTION_MAXIDLETIME_KEY,
        0);
  }
  rpc = YarnRPC.create(conf);
}
 

@Test
public void testAuthorizedAccess() throws Exception {
  MyContainerManager containerManager = new MyContainerManager();
  rm =
      new MockRMWithAMS(conf, containerManager);
  rm.start();

  MockNM nm1 = rm.registerNode("localhost:1234", 5120);

  Map<ApplicationAccessType, String> acls =
      new HashMap<ApplicationAccessType, String>(2);
  acls.put(ApplicationAccessType.VIEW_APP, "*");
  RMApp app = rm.submitApp(1024, "appname", "appuser", acls);

  nm1.nodeHeartbeat(true);

  int waitCount = 0;
  while (containerManager.containerTokens == null && waitCount++ < 20) {
    LOG.info("Waiting for AM Launch to happen..");
    Thread.sleep(1000);
  }
  Assert.assertNotNull(containerManager.containerTokens);

  RMAppAttempt attempt = app.getCurrentAppAttempt();
  ApplicationAttemptId applicationAttemptId = attempt.getAppAttemptId();
  waitForLaunchedState(attempt);

  // Create a client to the RM.
  final Configuration conf = rm.getConfig();
  final YarnRPC rpc = YarnRPC.create(conf);

  UserGroupInformation currentUser = UserGroupInformation
      .createRemoteUser(applicationAttemptId.toString());
  Credentials credentials = containerManager.getContainerCredentials();
  final InetSocketAddress rmBindAddress =
      rm.getApplicationMasterService().getBindAddress();
  Token<? extends TokenIdentifier> amRMToken =
      MockRMWithAMS.setupAndReturnAMRMToken(rmBindAddress,
        credentials.getAllTokens());
  currentUser.addToken(amRMToken);
  ApplicationMasterProtocol client = currentUser
      .doAs(new PrivilegedAction<ApplicationMasterProtocol>() {
        @Override
        public ApplicationMasterProtocol run() {
          return (ApplicationMasterProtocol) rpc.getProxy(ApplicationMasterProtocol.class, rm
            .getApplicationMasterService().getBindAddress(), conf);
        }
      });

  RegisterApplicationMasterRequest request = Records
      .newRecord(RegisterApplicationMasterRequest.class);
  RegisterApplicationMasterResponse response =
      client.registerApplicationMaster(request);
  Assert.assertNotNull(response.getClientToAMTokenMasterKey());
  if (UserGroupInformation.isSecurityEnabled()) {
    Assert
      .assertTrue(response.getClientToAMTokenMasterKey().array().length > 0);
  }
  Assert.assertEquals("Register response has bad ACLs", "*",
      response.getApplicationACLs().get(ApplicationAccessType.VIEW_APP));
}
 

@Test
public void testValidateResourceBlacklistRequest() throws Exception {

  MyContainerManager containerManager = new MyContainerManager();
  final MockRMWithAMS rm =
      new MockRMWithAMS(new YarnConfiguration(), containerManager);
  rm.start();

  MockNM nm1 = rm.registerNode("localhost:1234", 5120);

  Map<ApplicationAccessType, String> acls =
      new HashMap<ApplicationAccessType, String>(2);
  acls.put(ApplicationAccessType.VIEW_APP, "*");
  RMApp app = rm.submitApp(1024, "appname", "appuser", acls);

  nm1.nodeHeartbeat(true);

  RMAppAttempt attempt = app.getCurrentAppAttempt();
  ApplicationAttemptId applicationAttemptId = attempt.getAppAttemptId();
  waitForLaunchedState(attempt);

  // Create a client to the RM.
  final Configuration conf = rm.getConfig();
  final YarnRPC rpc = YarnRPC.create(conf);

  UserGroupInformation currentUser = 
      UserGroupInformation.createRemoteUser(applicationAttemptId.toString());
  Credentials credentials = containerManager.getContainerCredentials();
  final InetSocketAddress rmBindAddress =
      rm.getApplicationMasterService().getBindAddress();
  Token<? extends TokenIdentifier> amRMToken =
      MockRMWithAMS.setupAndReturnAMRMToken(rmBindAddress,
        credentials.getAllTokens());
  currentUser.addToken(amRMToken);
  ApplicationMasterProtocol client =
      currentUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() {
        @Override
        public ApplicationMasterProtocol run() {
          return (ApplicationMasterProtocol) rpc.getProxy(
            ApplicationMasterProtocol.class, rmBindAddress, conf);
        }
      });

  RegisterApplicationMasterRequest request = Records
      .newRecord(RegisterApplicationMasterRequest.class);
  client.registerApplicationMaster(request);

  ResourceBlacklistRequest blacklistRequest =
      ResourceBlacklistRequest.newInstance(
          Collections.singletonList(ResourceRequest.ANY), null);

  AllocateRequest allocateRequest =
      AllocateRequest.newInstance(0, 0.0f, null, null, blacklistRequest);
  boolean error = false;
  try {
    client.allocate(allocateRequest);
  } catch (InvalidResourceBlacklistRequestException e) {
    error = true;
  }

  rm.stop();
  
  Assert.assertTrue(
      "Didn't not catch InvalidResourceBlacklistRequestException", error);
}
 

/**
 * This tests whether a containerId is serialized/deserialized with epoch.
 *
 * @throws IOException
 * @throws InterruptedException
 * @throws YarnException
 */
private void testContainerTokenWithEpoch(Configuration conf)
    throws IOException, InterruptedException, YarnException {

  LOG.info("Running test for serializing/deserializing containerIds");

  NMTokenSecretManagerInRM nmTokenSecretManagerInRM =
      yarnCluster.getResourceManager().getRMContext()
          .getNMTokenSecretManager();
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  ApplicationAttemptId appAttemptId =
      ApplicationAttemptId.newInstance(appId, 0);
  ContainerId cId = ContainerId.newContainerId(appAttemptId, (5L << 40) | 3L);
  NodeManager nm = yarnCluster.getNodeManager(0);
  NMTokenSecretManagerInNM nmTokenSecretManagerInNM =
      nm.getNMContext().getNMTokenSecretManager();
  String user = "test";

  waitForNMToReceiveNMTokenKey(nmTokenSecretManagerInNM, nm);

  NodeId nodeId = nm.getNMContext().getNodeId();

  // Both id should be equal.
  Assert.assertEquals(nmTokenSecretManagerInNM.getCurrentKey().getKeyId(),
      nmTokenSecretManagerInRM.getCurrentKey().getKeyId());

  // Creating a normal Container Token
  RMContainerTokenSecretManager containerTokenSecretManager =
      yarnCluster.getResourceManager().getRMContext().
          getContainerTokenSecretManager();
  Resource r = Resource.newInstance(1230, 2);
  Token containerToken =
      containerTokenSecretManager.createContainerToken(cId, nodeId, user, r,
          Priority.newInstance(0), 0);
  
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier();
  byte[] tokenIdentifierContent = containerToken.getIdentifier().array();
  DataInputBuffer dib = new DataInputBuffer();
  dib.reset(tokenIdentifierContent, tokenIdentifierContent.length);
  containerTokenIdentifier.readFields(dib);
  
  
  Assert.assertEquals(cId, containerTokenIdentifier.getContainerID());
  Assert.assertEquals(
      cId.toString(), containerTokenIdentifier.getContainerID().toString());

  Token nmToken =
      nmTokenSecretManagerInRM.createNMToken(appAttemptId, nodeId, user);

  YarnRPC rpc = YarnRPC.create(conf);
  testStartContainer(rpc, appAttemptId, nodeId, containerToken, nmToken,
      false);

  List<ContainerId> containerIds = new LinkedList<ContainerId>();
  containerIds.add(cId);
  ContainerManagementProtocol proxy
      = getContainerManagementProtocolProxy(rpc, nmToken, nodeId, user);
  GetContainerStatusesResponse res = proxy.getContainerStatuses(
      GetContainerStatusesRequest.newInstance(containerIds));
  Assert.assertNotNull(res.getContainerStatuses().get(0));
  Assert.assertEquals(
      cId, res.getContainerStatuses().get(0).getContainerId());
  Assert.assertEquals(cId.toString(),
      res.getContainerStatuses().get(0).getContainerId().toString());
}
 

@Test
public void testValidateResourceBlacklistRequest() throws Exception {

  MyContainerManager containerManager = new MyContainerManager();
  final MockRMWithAMS rm =
      new MockRMWithAMS(new YarnConfiguration(), containerManager);
  rm.start();

  MockNM nm1 = rm.registerNode("localhost:1234", 5120);

  Map<ApplicationAccessType, String> acls =
      new HashMap<ApplicationAccessType, String>(2);
  acls.put(ApplicationAccessType.VIEW_APP, "*");
  RMApp app = rm.submitApp(1024, "appname", "appuser", acls);

  nm1.nodeHeartbeat(true);

  RMAppAttempt attempt = app.getCurrentAppAttempt();
  ApplicationAttemptId applicationAttemptId = attempt.getAppAttemptId();
  waitForLaunchedState(attempt);

  // Create a client to the RM.
  final Configuration conf = rm.getConfig();
  final YarnRPC rpc = YarnRPC.create(conf);

  UserGroupInformation currentUser = 
      UserGroupInformation.createRemoteUser(applicationAttemptId.toString());
  Credentials credentials = containerManager.getContainerCredentials();
  final InetSocketAddress rmBindAddress =
      rm.getApplicationMasterService().getBindAddress();
  Token<? extends TokenIdentifier> amRMToken =
      MockRMWithAMS.setupAndReturnAMRMToken(rmBindAddress,
        credentials.getAllTokens());
  currentUser.addToken(amRMToken);
  ApplicationMasterProtocol client =
      currentUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() {
        @Override
        public ApplicationMasterProtocol run() {
          return (ApplicationMasterProtocol) rpc.getProxy(
            ApplicationMasterProtocol.class, rmBindAddress, conf);
        }
      });

  RegisterApplicationMasterRequest request = Records
      .newRecord(RegisterApplicationMasterRequest.class);
  client.registerApplicationMaster(request);

  ResourceBlacklistRequest blacklistRequest =
      ResourceBlacklistRequest.newInstance(
          Collections.singletonList(ResourceRequest.ANY), null);

  AllocateRequest allocateRequest =
      AllocateRequest.newInstance(0, 0.0f, null, null, blacklistRequest);
  boolean error = false;
  try {
    client.allocate(allocateRequest);
  } catch (InvalidResourceBlacklistRequestException e) {
    error = true;
  }

  rm.stop();
  
  Assert.assertTrue(
      "Didn't not catch InvalidResourceBlacklistRequestException", error);
}
 

@Test
public void testGetClusterNodes() throws Exception {
  MockRM rm = new MockRM() {
    protected ClientRMService createClientRMService() {
      return new ClientRMService(this.rmContext, scheduler,
        this.rmAppManager, this.applicationACLsManager, this.queueACLsManager,
        this.getRMContext().getRMDelegationTokenSecretManager());
    };
  };
  rm.start();
  RMNodeLabelsManager labelsMgr = rm.getRMContext().getNodeLabelManager();
  labelsMgr.addToCluserNodeLabels(ImmutableSet.of("x", "y"));

  // Add a healthy node with label = x
  MockNM node = rm.registerNode("host1:1234", 1024);
  Map<NodeId, Set<String>> map = new HashMap<NodeId, Set<String>>();
  map.put(node.getNodeId(), ImmutableSet.of("x"));
  labelsMgr.replaceLabelsOnNode(map);
  rm.sendNodeStarted(node);
  node.nodeHeartbeat(true);
  
  // Add and lose a node with label = y
  MockNM lostNode = rm.registerNode("host2:1235", 1024);
  rm.sendNodeStarted(lostNode);
  lostNode.nodeHeartbeat(true);
  rm.NMwaitForState(lostNode.getNodeId(), NodeState.RUNNING);
  rm.sendNodeLost(lostNode);

  // Create a client.
  Configuration conf = new Configuration();
  YarnRPC rpc = YarnRPC.create(conf);
  InetSocketAddress rmAddress = rm.getClientRMService().getBindAddress();
  LOG.info("Connecting to ResourceManager at " + rmAddress);
  ApplicationClientProtocol client =
      (ApplicationClientProtocol) rpc
        .getProxy(ApplicationClientProtocol.class, rmAddress, conf);

  // Make call
  GetClusterNodesRequest request =
      GetClusterNodesRequest.newInstance(EnumSet.of(NodeState.RUNNING));
  List<NodeReport> nodeReports =
      client.getClusterNodes(request).getNodeReports();
  Assert.assertEquals(1, nodeReports.size());
  Assert.assertNotSame("Node is expected to be healthy!", NodeState.UNHEALTHY,
      nodeReports.get(0).getNodeState());
  
  // Check node's label = x
  Assert.assertTrue(nodeReports.get(0).getNodeLabels().contains("x"));

  // Now make the node unhealthy.
  node.nodeHeartbeat(false);

  // Call again
  nodeReports = client.getClusterNodes(request).getNodeReports();
  Assert.assertEquals("Unhealthy nodes should not show up by default", 0,
      nodeReports.size());
  
  // Change label of host1 to y
  map = new HashMap<NodeId, Set<String>>();
  map.put(node.getNodeId(), ImmutableSet.of("y"));
  labelsMgr.replaceLabelsOnNode(map);
  
  // Now query for UNHEALTHY nodes
  request = GetClusterNodesRequest.newInstance(EnumSet.of(NodeState.UNHEALTHY));
  nodeReports = client.getClusterNodes(request).getNodeReports();
  Assert.assertEquals(1, nodeReports.size());
  Assert.assertEquals("Node is expected to be unhealthy!", NodeState.UNHEALTHY,
      nodeReports.get(0).getNodeState());
  
  Assert.assertTrue(nodeReports.get(0).getNodeLabels().contains("y"));
  
  // Remove labels of host1
  map = new HashMap<NodeId, Set<String>>();
  map.put(node.getNodeId(), ImmutableSet.of("y"));
  labelsMgr.removeLabelsFromNode(map);
  
  // Query all states should return all nodes
  rm.registerNode("host3:1236", 1024);
  request = GetClusterNodesRequest.newInstance(EnumSet.allOf(NodeState.class));
  nodeReports = client.getClusterNodes(request).getNodeReports();
  Assert.assertEquals(3, nodeReports.size());
  
  // All host1-3's label should be empty (instead of null)
  for (NodeReport report : nodeReports) {
    Assert.assertTrue(report.getNodeLabels() != null
        && report.getNodeLabels().isEmpty());
  }
  
  rpc.stopProxy(client, conf);
  rm.close();
}
 

@Test
public void testGetNodeLabels() throws Exception {
  MockRM rm = new MockRM() {
    protected ClientRMService createClientRMService() {
      return new ClientRMService(this.rmContext, scheduler,
          this.rmAppManager, this.applicationACLsManager,
          this.queueACLsManager, this.getRMContext()
              .getRMDelegationTokenSecretManager());
    };
  };
  rm.start();
  RMNodeLabelsManager labelsMgr = rm.getRMContext().getNodeLabelManager();
  labelsMgr.addToCluserNodeLabels(ImmutableSet.of("x", "y"));

  Map<NodeId, Set<String>> map = new HashMap<NodeId, Set<String>>();
  map.put(NodeId.newInstance("host1", 0), ImmutableSet.of("x"));
  map.put(NodeId.newInstance("host2", 0), ImmutableSet.of("y"));
  labelsMgr.replaceLabelsOnNode(map);

  // Create a client.
  Configuration conf = new Configuration();
  YarnRPC rpc = YarnRPC.create(conf);
  InetSocketAddress rmAddress = rm.getClientRMService().getBindAddress();
  LOG.info("Connecting to ResourceManager at " + rmAddress);
  ApplicationClientProtocol client =
      (ApplicationClientProtocol) rpc.getProxy(
          ApplicationClientProtocol.class, rmAddress, conf);

  // Get node labels collection
  GetClusterNodeLabelsResponse response =
      client.getClusterNodeLabels(GetClusterNodeLabelsRequest.newInstance());
  Assert.assertTrue(response.getNodeLabels().containsAll(
      Arrays.asList("x", "y")));

  // Get node labels mapping
  GetNodesToLabelsResponse response1 =
      client.getNodeToLabels(GetNodesToLabelsRequest.newInstance());
  Map<NodeId, Set<String>> nodeToLabels = response1.getNodeToLabels();
  Assert.assertTrue(nodeToLabels.keySet().containsAll(
      Arrays.asList(NodeId.newInstance("host1", 0),
          NodeId.newInstance("host2", 0))));
  Assert.assertTrue(nodeToLabels.get(NodeId.newInstance("host1", 0))
      .containsAll(Arrays.asList("x")));
  Assert.assertTrue(nodeToLabels.get(NodeId.newInstance("host2", 0))
      .containsAll(Arrays.asList("y")));
  
  rpc.stopProxy(client, conf);
  rm.close();
}
 

@Test
public void testGetLabelsToNodes() throws Exception {
  MockRM rm = new MockRM() {
    protected ClientRMService createClientRMService() {
      return new ClientRMService(this.rmContext, scheduler,
          this.rmAppManager, this.applicationACLsManager,
          this.queueACLsManager, this.getRMContext()
              .getRMDelegationTokenSecretManager());
    };
  };
  rm.start();
  RMNodeLabelsManager labelsMgr = rm.getRMContext().getNodeLabelManager();
  labelsMgr.addToCluserNodeLabels(ImmutableSet.of("x", "y", "z"));

  Map<NodeId, Set<String>> map = new HashMap<NodeId, Set<String>>();
  map.put(NodeId.newInstance("host1", 0), ImmutableSet.of("x"));
  map.put(NodeId.newInstance("host1", 1), ImmutableSet.of("z"));
  map.put(NodeId.newInstance("host2", 0), ImmutableSet.of("y"));
  map.put(NodeId.newInstance("host3", 0), ImmutableSet.of("y"));
  map.put(NodeId.newInstance("host3", 1), ImmutableSet.of("z"));
  labelsMgr.replaceLabelsOnNode(map);

  // Create a client.
  Configuration conf = new Configuration();
  YarnRPC rpc = YarnRPC.create(conf);
  InetSocketAddress rmAddress = rm.getClientRMService().getBindAddress();
  LOG.info("Connecting to ResourceManager at " + rmAddress);
  ApplicationClientProtocol client =
      (ApplicationClientProtocol) rpc.getProxy(
          ApplicationClientProtocol.class, rmAddress, conf);

  // Get node labels collection
  GetClusterNodeLabelsResponse response =
      client.getClusterNodeLabels(GetClusterNodeLabelsRequest.newInstance());
  Assert.assertTrue(response.getNodeLabels().containsAll(
      Arrays.asList("x", "y", "z")));

  // Get labels to nodes mapping
  GetLabelsToNodesResponse response1 =
      client.getLabelsToNodes(GetLabelsToNodesRequest.newInstance());
  Map<String, Set<NodeId>> labelsToNodes = response1.getLabelsToNodes();
  Assert.assertTrue(
      labelsToNodes.keySet().containsAll(Arrays.asList("x", "y", "z")));
  Assert.assertTrue(
      labelsToNodes.get("x").containsAll(Arrays.asList(
      NodeId.newInstance("host1", 0))));
  Assert.assertTrue(
      labelsToNodes.get("y").containsAll(Arrays.asList(
      NodeId.newInstance("host2", 0), NodeId.newInstance("host3", 0))));
  Assert.assertTrue(
      labelsToNodes.get("z").containsAll(Arrays.asList(
      NodeId.newInstance("host1", 1), NodeId.newInstance("host3", 1))));

  // Get labels to nodes mapping for specific labels
  Set<String> setlabels =
      new HashSet<String>(Arrays.asList(new String[]{"x", "z"}));
  GetLabelsToNodesResponse response2 =
      client.getLabelsToNodes(GetLabelsToNodesRequest.newInstance(setlabels));
  labelsToNodes = response2.getLabelsToNodes();
  Assert.assertTrue(
      labelsToNodes.keySet().containsAll(Arrays.asList("x", "z")));
  Assert.assertTrue(
      labelsToNodes.get("x").containsAll(Arrays.asList(
      NodeId.newInstance("host1", 0))));
  Assert.assertTrue(
      labelsToNodes.get("z").containsAll(Arrays.asList(
      NodeId.newInstance("host1", 1), NodeId.newInstance("host3", 1))));
  Assert.assertEquals(labelsToNodes.get("y"), null);

  rpc.stopProxy(client, conf);
  rm.close();
}