YARN-11298. Improve Yarn Router Junit Test Close MockRM. (#4870)

This commit is contained in:
slfan1989 2022-09-09 02:42:36 +08:00 committed by GitHub
parent 0db3ee5b4b
commit 29c4d8d8f7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 62 additions and 2 deletions

View File

@ -31,8 +31,10 @@ import org.junit.Assert;
public class MockClientRequestInterceptor
extends DefaultClientRequestInterceptor {
MockRM mockRM = null;
public void init(String user) {
MockRM mockRM = new MockRM(super.getConf()) {
mockRM = new MockRM(super.getConf()) {
@Override
protected ClientRMService createClientRMService() {
return new ClientRMService(getRMContext(), getResourceScheduler(),
@ -68,4 +70,11 @@ public class MockClientRequestInterceptor
super.setRMClient(mockRM.getClientRMService());
}
@Override
public void shutdown() {
if (mockRM != null) {
mockRM.stop();
}
super.shutdown();
}
}

View File

@ -48,6 +48,8 @@ import org.apache.hadoop.yarn.server.resourcemanager.security.QueueACLsManager;
import org.apache.hadoop.yarn.server.resourcemanager.security.RMDelegationTokenSecretManager;
import org.apache.hadoop.yarn.server.security.ApplicationACLsManager;
import org.junit.Assert;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Extends the FederationClientInterceptor and overrides methods to provide a
@ -56,6 +58,9 @@ import org.junit.Assert;
public class TestableFederationClientInterceptor
extends FederationClientInterceptor {
private static final Logger LOG =
LoggerFactory.getLogger(TestableFederationClientInterceptor.class);
private ConcurrentHashMap<SubClusterId, MockRM> mockRMs =
new ConcurrentHashMap<>();
@ -161,4 +166,31 @@ public class TestableFederationClientInterceptor
throw new RuntimeException(e);
}
}
@Override
public void shutdown() {
if (mockRMs != null && !mockRMs.isEmpty()) {
for (Map.Entry<SubClusterId, MockRM> item : mockRMs.entrySet()) {
SubClusterId subClusterId = item.getKey();
// close mockNM
MockNM mockNM = mockNMs.getOrDefault(subClusterId, null);
try {
mockNM.unRegisterNode();
mockNM = null;
} catch (Exception e) {
LOG.error("mockNM unRegisterNode error.", e);
}
// close mockRM
MockRM mockRM = item.getValue();
if (mockRM != null) {
mockRM.stop();
}
}
}
mockNMs.clear();
mockRMs.clear();
super.shutdown();
}
}

View File

@ -32,8 +32,9 @@ import org.apache.hadoop.yarn.server.resourcemanager.MockRM;
public class MockRMAdminRequestInterceptor
extends DefaultRMAdminRequestInterceptor {
MockRM mockRM = null;
public void init(String user) {
MockRM mockRM = new MockRM(super.getConf()) {
mockRM = new MockRM(super.getConf()) {
@Override
protected AdminService createAdminService() {
return new AdminService(this) {
@ -60,4 +61,12 @@ public class MockRMAdminRequestInterceptor
mockRM.start();
super.setRMAdmin(mockRM.getAdminService());
}
@Override
public void shutdown() {
if (mockRM != null) {
mockRM.stop();
}
super.shutdown();
}
}

View File

@ -20,6 +20,7 @@ package org.apache.hadoop.yarn.server.router.secure;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
import org.apache.hadoop.metrics2.lib.DefaultMetricsSystem;
import org.apache.hadoop.minikdc.MiniKdc;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
@ -34,6 +35,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
@ -107,6 +109,8 @@ public abstract class AbstractSecureRouterTest {
// Router Kerberos KeyTab configuration
conf.set(YarnConfiguration.ROUTER_PRINCIPAL, ROUTER_LOCALHOST_REALM);
conf.set(YarnConfiguration.ROUTER_KEYTAB, routerKeytab.getAbsolutePath());
DefaultMetricsSystem.setMiniClusterMode(true);
}
/**
@ -214,6 +218,12 @@ public abstract class AbstractSecureRouterTest {
public static void afterSecureRouterTest() throws Exception {
LOG.info("teardown of kdc instance.");
teardownKDC();
if (mockRMs != null && mockRMs.isEmpty()) {
for (MockRM mockRM : mockRMs.values()) {
mockRM.close();
}
}
mockRMs.clear();
}
public static MiniKdc getKdc() {