Examples with RecoveredNMTokensState org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService.RecoveredNMTokensState used on opensource projects

Example 1 with RecoveredNMTokensState

use of org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService.RecoveredNMTokensState in project hadoop by apache.

the class NMTokenSecretManagerInNM method recover.

public synchronized void recover() throws IOException {
    RecoveredNMTokensState state = stateStore.loadNMTokensState();
    MasterKey key = state.getCurrentMasterKey();
    if (key != null) {
        super.currentMasterKey = new MasterKeyData(key, createSecretKey(key.getBytes().array()));
    }
    key = state.getPreviousMasterKey();
    if (key != null) {
        previousMasterKey = new MasterKeyData(key, createSecretKey(key.getBytes().array()));
    }
    // restore the serial number from the current master key
    if (super.currentMasterKey != null) {
        super.serialNo = super.currentMasterKey.getMasterKey().getKeyId() + 1;
    }
    for (Map.Entry<ApplicationAttemptId, MasterKey> entry : state.getApplicationMasterKeys().entrySet()) {
        key = entry.getValue();
        oldMasterKeys.put(entry.getKey(), new MasterKeyData(key, createSecretKey(key.getBytes().array())));
    }
    // reconstruct app to app attempts map
    appToAppAttemptMap.clear();
    for (ApplicationAttemptId attempt : oldMasterKeys.keySet()) {
        ApplicationId app = attempt.getApplicationId();
        List<ApplicationAttemptId> attempts = appToAppAttemptMap.get(app);
        if (attempts == null) {
            attempts = new ArrayList<ApplicationAttemptId>();
            appToAppAttemptMap.put(app, attempts);
        }
        attempts.add(attempt);
    }
}

Example 2 with RecoveredNMTokensState

use of org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService.RecoveredNMTokensState in project hadoop by apache.

the class TestNMLeveldbStateStoreService method testNMTokenStorage.

@Test
public void testNMTokenStorage() throws IOException {
    // test empty when no state
    RecoveredNMTokensState state = stateStore.loadNMTokensState();
    assertNull(state.getCurrentMasterKey());
    assertNull(state.getPreviousMasterKey());
    assertTrue(state.getApplicationMasterKeys().isEmpty());
    // store a master key and verify recovered
    NMTokenSecretManagerForTest secretMgr = new NMTokenSecretManagerForTest();
    MasterKey currentKey = secretMgr.generateKey();
    stateStore.storeNMTokenCurrentMasterKey(currentKey);
    restartStateStore();
    state = stateStore.loadNMTokensState();
    assertEquals(currentKey, state.getCurrentMasterKey());
    assertNull(state.getPreviousMasterKey());
    assertTrue(state.getApplicationMasterKeys().isEmpty());
    // store a previous key and verify recovered
    MasterKey prevKey = secretMgr.generateKey();
    stateStore.storeNMTokenPreviousMasterKey(prevKey);
    restartStateStore();
    state = stateStore.loadNMTokensState();
    assertEquals(currentKey, state.getCurrentMasterKey());
    assertEquals(prevKey, state.getPreviousMasterKey());
    assertTrue(state.getApplicationMasterKeys().isEmpty());
    // store a few application keys and verify recovered
    ApplicationAttemptId attempt1 = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    MasterKey attemptKey1 = secretMgr.generateKey();
    stateStore.storeNMTokenApplicationMasterKey(attempt1, attemptKey1);
    ApplicationAttemptId attempt2 = ApplicationAttemptId.newInstance(ApplicationId.newInstance(2, 3), 4);
    MasterKey attemptKey2 = secretMgr.generateKey();
    stateStore.storeNMTokenApplicationMasterKey(attempt2, attemptKey2);
    restartStateStore();
    state = stateStore.loadNMTokensState();
    assertEquals(currentKey, state.getCurrentMasterKey());
    assertEquals(prevKey, state.getPreviousMasterKey());
    Map<ApplicationAttemptId, MasterKey> loadedAppKeys = state.getApplicationMasterKeys();
    assertEquals(2, loadedAppKeys.size());
    assertEquals(attemptKey1, loadedAppKeys.get(attempt1));
    assertEquals(attemptKey2, loadedAppKeys.get(attempt2));
    // add/update/remove keys and verify recovered
    ApplicationAttemptId attempt3 = ApplicationAttemptId.newInstance(ApplicationId.newInstance(5, 6), 7);
    MasterKey attemptKey3 = secretMgr.generateKey();
    stateStore.storeNMTokenApplicationMasterKey(attempt3, attemptKey3);
    stateStore.removeNMTokenApplicationMasterKey(attempt1);
    attemptKey2 = prevKey;
    stateStore.storeNMTokenApplicationMasterKey(attempt2, attemptKey2);
    prevKey = currentKey;
    stateStore.storeNMTokenPreviousMasterKey(prevKey);
    currentKey = secretMgr.generateKey();
    stateStore.storeNMTokenCurrentMasterKey(currentKey);
    restartStateStore();
    state = stateStore.loadNMTokensState();
    assertEquals(currentKey, state.getCurrentMasterKey());
    assertEquals(prevKey, state.getPreviousMasterKey());
    loadedAppKeys = state.getApplicationMasterKeys();
    assertEquals(2, loadedAppKeys.size());
    assertNull(loadedAppKeys.get(attempt1));
    assertEquals(attemptKey2, loadedAppKeys.get(attempt2));
    assertEquals(attemptKey3, loadedAppKeys.get(attempt3));
}