@Before
public void before() throws IOException {
mgr = new MockNodeLabelManager();
conf = new Configuration();
conf.setBoolean(YarnConfiguration.NODE_LABELS_ENABLED, true);
File tempDir = File.createTempFile("nlb", ".tmp");
tempDir.delete();
tempDir.mkdirs();
tempDir.deleteOnExit();
conf.set(YarnConfiguration.FS_NODE_LABELS_STORE_ROOT_DIR, tempDir.getAbsolutePath());
mgr.init(conf);
mgr.start();
}
@SuppressWarnings({"unchecked", "rawtypes"})
@Test(timeout = 10000)
public void testRecoverWithMirror() throws Exception {
mgr.addToCluserNodeLabels(toSet("p1", "p2", "p3"));
mgr.addToCluserNodeLabels(toSet("p4"));
mgr.addToCluserNodeLabels(toSet("p5", "p6"));
mgr.replaceLabelsOnNode(
(Map) ImmutableMap.of(toNodeId("n1"), toSet("p1"), toNodeId("n2"), toSet("p2")));
mgr.replaceLabelsOnNode(
(Map)
ImmutableMap.of(
toNodeId("n3"),
toSet("p3"),
toNodeId("n4"),
toSet("p4"),
toNodeId("n5"),
toSet("p5"),
toNodeId("n6"),
toSet("p6"),
toNodeId("n7"),
toSet("p6")));
/*
* node -> partition p1: n1 p2: n2 p3: n3 p4: n4 p5: n5 p6: n6, n7
*/
mgr.removeFromClusterNodeLabels(toSet("p1"));
mgr.removeFromClusterNodeLabels(Arrays.asList("p3", "p5"));
/*
* After removed p2: n2 p4: n4 p6: n6, n7
*/
// shutdown mgr and start a new mgr
mgr.stop();
mgr = new MockNodeLabelManager();
mgr.init(conf);
mgr.start();
// check variables
Assert.assertEquals(3, mgr.getClusterNodeLabels().size());
Assert.assertTrue(mgr.getClusterNodeLabels().containsAll(Arrays.asList("p2", "p4", "p6")));
assertMapContains(
mgr.getNodeLabels(),
ImmutableMap.of(
toNodeId("n2"),
toSet("p2"),
toNodeId("n4"),
toSet("p4"),
toNodeId("n6"),
toSet("p6"),
toNodeId("n7"),
toSet("p6")));
assertLabelsToNodesEquals(
mgr.getLabelsToNodes(),
ImmutableMap.of(
"p6", toSet(toNodeId("n6"), toNodeId("n7")),
"p4", toSet(toNodeId("n4")),
"p2", toSet(toNodeId("n2"))));
// stutdown mgr and start a new mgr
mgr.stop();
mgr = new MockNodeLabelManager();
mgr.init(conf);
mgr.start();
// check variables
Assert.assertEquals(3, mgr.getClusterNodeLabels().size());
Assert.assertTrue(mgr.getClusterNodeLabels().containsAll(Arrays.asList("p2", "p4", "p6")));
assertMapContains(
mgr.getNodeLabels(),
ImmutableMap.of(
toNodeId("n2"),
toSet("p2"),
toNodeId("n4"),
toSet("p4"),
toNodeId("n6"),
toSet("p6"),
toNodeId("n7"),
toSet("p6")));
assertLabelsToNodesEquals(
mgr.getLabelsToNodes(),
ImmutableMap.of(
"p6", toSet(toNodeId("n6"), toNodeId("n7")),
"p4", toSet(toNodeId("n4")),
"p2", toSet(toNodeId("n2"))));
mgr.stop();
}
@After
public void after() throws IOException {
getStore().fs.delete(getStore().fsWorkingPath, true);
mgr.stop();
}
@SuppressWarnings({"unchecked", "rawtypes"})
@Test // (timeout = 10000)
public void testSerilizationAfterRecovery() throws Exception {
mgr.addToCluserNodeLabels(toSet("p1", "p2", "p3"));
mgr.addToCluserNodeLabels(toSet("p4"));
mgr.addToCluserNodeLabels(toSet("p5", "p6"));
mgr.replaceLabelsOnNode(
ImmutableMap.of(toNodeId("n1"), toSet("p1"), toNodeId("n2"), toSet("p2")));
mgr.replaceLabelsOnNode(
(Map)
ImmutableMap.of(
toNodeId("n3"),
toSet("p3"),
toNodeId("n4"),
toSet("p4"),
toNodeId("n5"),
toSet("p5"),
toNodeId("n6"),
toSet("p6"),
toNodeId("n7"),
toSet("p6")));
/*
* node -> labels
* p1: n1
* p2: n2
* p3: n3
* p4: n4
* p5: n5
* p6: n6, n7
*/
mgr.removeFromClusterNodeLabels(toSet("p1"));
mgr.removeFromClusterNodeLabels(Arrays.asList("p3", "p5"));
/*
* After removed
* p2: n2
* p4: n4
* p6: n6, n7
*/
// shutdown mgr and start a new mgr
mgr.stop();
mgr = new MockNodeLabelManager();
mgr.init(conf);
mgr.start();
// check variables
Assert.assertEquals(3, mgr.getClusterNodeLabels().size());
Assert.assertTrue(mgr.getClusterNodeLabels().containsAll(Arrays.asList("p2", "p4", "p6")));
assertMapContains(
mgr.getNodeLabels(),
ImmutableMap.of(
toNodeId("n2"),
toSet("p2"),
toNodeId("n4"),
toSet("p4"),
toNodeId("n6"),
toSet("p6"),
toNodeId("n7"),
toSet("p6")));
assertLabelsToNodesEquals(
mgr.getLabelsToNodes(),
ImmutableMap.of(
"p6", toSet(toNodeId("n6"), toNodeId("n7")),
"p4", toSet(toNodeId("n4")),
"p2", toSet(toNodeId("n2"))));
/*
* Add label p7,p8 then shutdown
*/
mgr = new MockNodeLabelManager();
mgr.init(conf);
mgr.start();
mgr.addToCluserNodeLabels(toSet("p7", "p8"));
mgr.stop();
/*
* Restart, add label p9 and shutdown
*/
mgr = new MockNodeLabelManager();
mgr.init(conf);
mgr.start();
mgr.addToCluserNodeLabels(toSet("p9"));
mgr.stop();
/*
* Recovery, and see if p9 added
*/
mgr = new MockNodeLabelManager();
mgr.init(conf);
mgr.start();
// check variables
Assert.assertEquals(6, mgr.getClusterNodeLabels().size());
Assert.assertTrue(
mgr.getClusterNodeLabels().containsAll(Arrays.asList("p2", "p4", "p6", "p7", "p8", "p9")));
mgr.stop();
}