private void updateNodeLabels(NodeId nodeId, String... nodeLabelsStr) { nodeLabelsMap.put(nodeId, toNodeLabelSet(nodeLabelsStr)); }
@Test public void testWithNodeLabelUpdateEnabled() throws Exception { conf.setLong(YarnConfiguration.RM_NODE_LABELS_PROVIDER_FETCH_INTERVAL_MS, 1000); MockRM rm = new MockRM(conf); rm.init(conf); rm.getRMContext().getRMDelegatedNodeLabelsUpdater().nodeLabelsUpdateInterval = 3 * 1000; rm.start(); RMNodeLabelsManager mgr = rm.getRMContext().getNodeLabelManager(); mgr.addToCluserNodeLabelsWithDefaultExclusivity(ImmutableSet.of("x", "y")); NodeId nodeId = toNodeId("h1:1234"); assertEquals(0, mgr.getLabelsOnNode(nodeId).size()); updateNodeLabels(nodeId, "x"); registerNode(rm, nodeId); Thread.sleep(4000); assertCollectionEquals(ImmutableSet.of("x"), mgr.getLabelsOnNode(nodeId)); // Ensure that node labels are updated if NodeLabelsProvider // gives different labels updateNodeLabels(nodeId, "y"); Thread.sleep(4000); assertCollectionEquals(ImmutableSet.of("y"), mgr.getLabelsOnNode(nodeId)); rm.stop(); }
@Test public void testWithNodeLabelUpdateDisabled() throws Exception { conf.setLong(YarnConfiguration.RM_NODE_LABELS_PROVIDER_FETCH_INTERVAL_MS, RMDelegatedNodeLabelsUpdater.DISABLE_DELEGATED_NODE_LABELS_UPDATE); MockRM rm = new MockRM(conf); rm.init(conf); rm.getRMContext().getRMDelegatedNodeLabelsUpdater().nodeLabelsUpdateInterval = 3 * 1000; rm.start(); RMNodeLabelsManager mgr = rm.getRMContext().getNodeLabelManager(); mgr.addToCluserNodeLabelsWithDefaultExclusivity(ImmutableSet.of("x")); NodeId nodeId = toNodeId("h1:1234"); updateNodeLabels(nodeId, "x"); registerNode(rm, nodeId); Thread.sleep(4000); // Ensure that even though timer is not run, node labels are fetched // when node is registered assertCollectionEquals(ImmutableSet.of("x"), mgr.getLabelsOnNode(nodeId)); rm.stop(); }