From 00905efab22edd9857e0a3828c201bf70f03cb96 Mon Sep 17 00:00:00 2001
From: Subru Krishnan <subru@apache.org>
Date: Fri, 6 Apr 2018 16:31:16 -0700
Subject: [PATCH] YARN-8110. AMRMProxy recover should catch for all throwable
 to avoid premature exit. (Botong Huang via Subru).

---
 .../amrmproxy/AMRMProxyService.java           |  2 +-
 .../amrmproxy/BaseAMRMProxyTest.java          |  5 +++
 .../amrmproxy/TestAMRMProxyService.java       | 42 +++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/AMRMProxyService.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/AMRMProxyService.java
index 815e39bfff..86fbb72ee2 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/AMRMProxyService.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/AMRMProxyService.java
@@ -261,7 +261,7 @@ public void recover() throws IOException {
         // Create the intercepter pipeline for the AM
         initializePipeline(attemptId, user, amrmToken, localToken,
             entry.getValue(), true, amCred);
-      } catch (IOException e) {
+      } catch (Throwable e) {
         LOG.error("Exception when recovering " + attemptId
             + ", removing it from NMStateStore and move on", e);
         this.nmContext.getNMStateStore().removeAMRMProxyAppContext(attemptId);
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/BaseAMRMProxyTest.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/BaseAMRMProxyTest.java
index 4b1a887ba0..677732d634 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/BaseAMRMProxyTest.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/BaseAMRMProxyTest.java
@@ -112,6 +112,11 @@ protected MockAMRMProxyService getAMRMProxyService() {
     return this.amrmProxyService;
   }
 
+  protected Context getNMContext() {
+    Assert.assertNotNull(this.nmContext);
+    return this.nmContext;
+  }
+
   @Before
   public void setUp() throws IOException {
     this.conf = createConfiguration();
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/TestAMRMProxyService.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/TestAMRMProxyService.java
index b955311bb9..1eefbd59a6 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/TestAMRMProxyService.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/amrmproxy/TestAMRMProxyService.java
@@ -44,6 +44,7 @@
 import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
 import org.apache.hadoop.yarn.server.MockResourceManagerFacade;
 import org.apache.hadoop.yarn.server.nodemanager.amrmproxy.AMRMProxyService.RequestInterceptorChainWrapper;
+import org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService.RecoveredAMRMProxyState;
 import org.apache.hadoop.yarn.util.Records;
 import org.junit.Assert;
 import org.junit.Test;
@@ -633,6 +634,35 @@ public void testRecovery() throws YarnException, Exception {
     mockRM = null;
   }
 
+  /**
+   * Test AMRMProxy restart with application recovery failure.
+   */
+  @Test
+  public void testAppRecoveryFailure() throws YarnException, Exception {
+    Configuration conf = createConfiguration();
+    // Use the MockRequestInterceptorAcrossRestart instead for the chain
+    conf.set(YarnConfiguration.AMRM_PROXY_INTERCEPTOR_CLASS_PIPELINE,
+        BadRequestInterceptorAcrossRestart.class.getName());
+
+    mockRM = new MockResourceManagerFacade(new YarnConfiguration(conf), 0);
+
+    createAndStartAMRMProxyService(conf);
+
+    // Create an app entry in NMSS
+    registerApplicationMaster(1);
+
+    RecoveredAMRMProxyState state =
+        getNMContext().getNMStateStore().loadAMRMProxyState();
+    Assert.assertEquals(1, state.getAppContexts().size());
+
+    // AMRMProxy restarts and recover
+    createAndStartAMRMProxyService(conf);
+
+    state = getNMContext().getNMStateStore().loadAMRMProxyState();
+    // The app that failed to recover should have been removed from NMSS
+    Assert.assertEquals(0, state.getAppContexts().size());
+  }
+
   /**
    * A mock intercepter implementation that uses the same mockRM instance across
    * restart.
@@ -672,4 +702,16 @@ public AllocateResponse allocate(AllocateRequest request)
     }
   }
 
+  /**
+   * A mock intercepter implementation that throws when recovering.
+   */
+  public static class BadRequestInterceptorAcrossRestart
+      extends MockRequestInterceptorAcrossRestart {
+
+    @Override
+    public void recover(Map<String, byte[]> recoveredDataMap) {
+      throw new RuntimeException("Kaboom");
+    }
+  }
+
 }