From c5c73182e5adab2299bfb27886bfd3ecc229009c Mon Sep 17 00:00:00 2001
From: Billie Rinaldi <billie@apache.org>
Date: Tue, 18 Dec 2018 08:43:09 -0800
Subject: [PATCH] YARN-9072. Send exit command to terminate docker exec on
 connection close. Contributed by Eric Yang

---
 .../server/nodemanager/webapp/ContainerShellWebSocket.java    | 4 ++++
 .../yarn/server/nodemanager/webapp/TestNMWebTerminal.java     | 2 +-
 2 files changed, 5 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/webapp/ContainerShellWebSocket.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/webapp/ContainerShellWebSocket.java
index 2fc1d16589..ade1211e7b 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/webapp/ContainerShellWebSocket.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/webapp/ContainerShellWebSocket.java
@@ -113,6 +113,7 @@ public void onConnect(Session session) {
           .fromString(cId));
       if (!checkAuthorization(session, container)) {
         session.close(1008, "Forbidden");
+        return;
       }
       LOG.info(session.getRemoteAddress().getHostString() + " connected!");
       LOG.info(
@@ -135,6 +136,9 @@ public void onConnect(Session session) {
   public void onClose(Session session, int status, String reason) {
     try {
       LOG.info(session.getRemoteAddress().getHostString() + " closed!");
+      String exit = "exit\r\n";
+      pair.out.write(exit.getBytes(Charset.forName("UTF-8")));
+      pair.out.flush();
       pair.in.close();
       pair.out.close();
     } catch (IOException e) {
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/webapp/TestNMWebTerminal.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/webapp/TestNMWebTerminal.java
index 8589b94ba1..ec7d62c803 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/webapp/TestNMWebTerminal.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/webapp/TestNMWebTerminal.java
@@ -48,7 +48,7 @@
 public class TestNMWebTerminal {
   private static final File TESTROOTDIR = new File("target",
       TestNMWebServer.class.getSimpleName());
-  private static File TESTLOGDIR = new File("target",
+  private static final File TESTLOGDIR = new File("target",
       TestNMWebServer.class.getSimpleName() + "LogDir");
   private NodeHealthCheckerService healthChecker;
   private WebServer server;