From 709f08813350b97862d1b09d992faa0024a3255e Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Mon, 9 Nov 2020 16:21:53 +0800
Subject: [PATCH] improve print agentId in welcome message. #1575

---
 .../arthas/core/server/ArthasBootstrap.java      | 13 ++-----------
 .../arthas/core/shell/impl/ShellServerImpl.java  | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java
index 087a3f035..8b25b8403 100644
--- a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java
+++ b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java
@@ -301,7 +301,6 @@ public class ArthasBootstrap {
                     System.getProperty(ArthasConstants.SPRING_APPLICATION_NAME, null)));
         }
 
-        String agentId = null;
         try {
             if (configure.getTunnelServer() != null) {
                 tunnelClient = new TunnelClient();
@@ -311,9 +310,6 @@ public class ArthasBootstrap {
                 tunnelClient.setVersion(ArthasBanner.version());
                 ChannelFuture channelFuture = tunnelClient.start();
                 channelFuture.await(10, TimeUnit.SECONDS);
-                if(channelFuture.isSuccess()) {
-                    agentId = tunnelClient.getId();
-                }
             }
         } catch (Throwable t) {
             logger().error("start tunnel client error", t);
@@ -322,17 +318,12 @@ public class ArthasBootstrap {
         try {
             ShellServerOptions options = new ShellServerOptions()
                             .setInstrumentation(instrumentation)
-                            .setPid(PidUtils.currentLongPid());
+                            .setPid(PidUtils.currentLongPid())
+                            .setWelcomeMessage(ArthasBanner.welcome());
             if (configure.getSessionTimeout() != null) {
                 options.setSessionTimeout(configure.getSessionTimeout() * 1000);
             }
 
-            if (agentId != null) {
-                Map<String, String> welcomeInfos = new HashMap<String, String>();
-                welcomeInfos.put("id", agentId);
-                options.setWelcomeMessage(ArthasBanner.welcome(welcomeInfos));
-            }
-
             shellServer = new ShellServerImpl(options);
             BuiltinCommandPack builtinCommands = new BuiltinCommandPack();
             List<CommandResolver> resolvers = new ArrayList<CommandResolver>();
diff --git a/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellServerImpl.java b/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellServerImpl.java
index 71e68d292..38934d79a 100644
--- a/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellServerImpl.java
+++ b/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellServerImpl.java
@@ -2,6 +2,7 @@ package com.taobao.arthas.core.shell.impl;
 
 import com.alibaba.arthas.deps.org.slf4j.Logger;
 import com.alibaba.arthas.deps.org.slf4j.LoggerFactory;
+import com.alibaba.arthas.tunnel.client.TunnelClient;
 import com.taobao.arthas.core.server.ArthasBootstrap;
 import com.taobao.arthas.core.shell.Shell;
 import com.taobao.arthas.core.shell.ShellServer;
@@ -19,10 +20,12 @@ import com.taobao.arthas.core.shell.system.impl.InternalCommandManager;
 import com.taobao.arthas.core.shell.system.impl.JobControllerImpl;
 import com.taobao.arthas.core.shell.term.Term;
 import com.taobao.arthas.core.shell.term.TermServer;
+import com.taobao.arthas.core.util.ArthasBanner;
 
 import java.lang.instrument.Instrumentation;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -93,6 +96,7 @@ public class ShellServerImpl extends ShellServer {
         }
 
         ShellImpl session = createShell(term);
+        tryUpdateWelcomeMessage();
         session.setWelcome(welcomeMessage);
         session.closedFuture.setHandler(new SessionClosedHandler(this, session));
         session.init();
@@ -100,6 +104,18 @@ public class ShellServerImpl extends ShellServer {
         session.readline(); // Now readline
     }
 
+    private void tryUpdateWelcomeMessage() {
+        TunnelClient tunnelClient = ArthasBootstrap.getInstance().getTunnelClient();
+        if (tunnelClient != null) {
+            String id = tunnelClient.getId();
+            if (id != null) {
+                Map<String, String> welcomeInfos = new HashMap<String, String>();
+                welcomeInfos.put("id", id);
+                this.welcomeMessage = ArthasBanner.welcome(welcomeInfos);
+            }
+        }
+    }
+
     @Override
     public ShellServer listen(final Handler<Future<Void>> listenHandler) {
         final List<TermServer> toStart;