@Test public void testDistributedLeaderStart() throws Exception { RMContextImpl rmContext = new RMContextImpl(); rmContext.setDistributedEnabled(true); rmContext.setHAEnabled(true); GroupMembershipService groupMembershipService = new GroupMembershipService(null, rmContext); groupMembershipService.init(conf); NMTokenSecretManagerInRM nmTokenSecretManager = new NMTokenSecretManagerInRM(conf, rmContext); try { groupMembershipService.start(); rmContext.setRMGroupMembershipService(groupMembershipService); rmContext.setStateStore(new NDBRMStateStore()); assertTrue("should be leader", rmContext.isLeader()); nmTokenSecretManager.start(); assertNotNull("did not roll master key", nmTokenSecretManager.getCurrentKey()); RMStateStore.RMState state = rmContext.getStateStore().loadState(rmContext); assertNotNull( "key not persisted to the database", state.getSecretTokenMamagerKey(RMStateStore.KeyType.CURRENTNMTOKENMASTERKEY)); NMTokenSecretManagerInRM nmTokenSecretManager2 = new NMTokenSecretManagerInRM(conf, rmContext); nmTokenSecretManager2.recover(state); assertEquals(nmTokenSecretManager.getCurrentKey(), nmTokenSecretManager2.getCurrentKey()); assertEquals(nmTokenSecretManager.getNextKey(), nmTokenSecretManager2.getNextKey()); } finally { groupMembershipService.stop(); nmTokenSecretManager.stop(); DefaultMetricsSystem.shutdown(); } }
@Test public void testNonLeaderKeyReception() throws InterruptedException, StorageInitializtionException, Exception { // create a groupMembershipService that will be leader RMContextImpl rmContext = new RMContextImpl(); rmContext.setDistributedEnabled(true); rmContext.setHAEnabled(true); GroupMembershipService groupMembershipService = new GroupMembershipService(null, rmContext); groupMembershipService.init(conf); NMTokenSecretManagerInRM nmTokenSecretManager = new NMTokenSecretManagerInRM(conf, rmContext); // create a resrouce tracker conf.setBoolean(YarnConfiguration.RM_HA_ENABLED, true); conf.setBoolean(YarnConfiguration.DISTRIBUTED_RM, true); MockRM mockRM = new MockRM(conf); mockRM.init(conf); try { groupMembershipService.start(); rmContext.setRMGroupMembershipService(groupMembershipService); rmContext.setStateStore(new NDBRMStateStore()); while (!rmContext.isLeader()) { Thread.sleep(1000); } mockRM.start(); if (mockRM.getRMContext().isDistributedEnabled() && !mockRM.getRMContext().isLeader()) { conf.set( YarnConfiguration.EVENT_RT_CONFIG_PATH, "target/test-classes/RT_EventAPIConfig.ini"); NdbRtStreamingProcessor rtStreamingProcessor = new NdbRtStreamingProcessor(mockRM.getRMContext()); RMStorageFactory.kickTheNdbEventStreamingAPI(false, conf); new Thread(rtStreamingProcessor).start(); } // this should be a resource tracker not a scheduler assertFalse(mockRM.getRMContext().isLeader()); // simulate creation of a token on the sheduler nmTokenSecretManager.start(); assertNotNull("did not roll master key", nmTokenSecretManager.getCurrentKey()); Thread.sleep(1000); dummyUpdate(); Thread.sleep(1000); RMStateStore.RMState state = rmContext.getStateStore().loadState(rmContext); assertEquals( "key not persisted to the database", state.getSecretTokenMamagerKey(RMStateStore.KeyType.CURRENTNMTOKENMASTERKEY), nmTokenSecretManager.getCurrentKey()); assertEquals( nmTokenSecretManager.getCurrentKey(), mockRM.getRMContext().getNMTokenSecretManager().getCurrentKey()); assertEquals( nmTokenSecretManager.getNextKey(), mockRM.getRMContext().getNMTokenSecretManager().getNextKey()); } finally { groupMembershipService.stop(); mockRM.stop(); nmTokenSecretManager.stop(); DefaultMetricsSystem.shutdown(); } }