feat: session命令支持显示和设置session data

以便可以设置上下文参供其他命令提供默认值
alibaba · Oct 28, 2023 · 19895b4 · 19895b4
1 parent e638b97
commit 19895b4
Show file tree

Hide file tree

Showing 5 changed files with 113 additions and 1 deletion.
diff --git a/core/src/main/java/com/taobao/arthas/core/command/basic1000/SessionCommand.java b/core/src/main/java/com/taobao/arthas/core/command/basic1000/SessionCommand.java
@@ -1,32 +1,118 @@
 package com.taobao.arthas.core.command.basic1000;
 
+import com.taobao.arthas.core.command.Constants;
 import com.taobao.arthas.core.command.model.SessionModel;
 import com.taobao.arthas.core.server.ArthasBootstrap;
 import com.taobao.arthas.core.shell.command.AnnotatedCommand;
 import com.taobao.arthas.core.shell.command.CommandProcess;
 import com.taobao.arthas.core.shell.session.Session;
 import com.taobao.arthas.core.util.UserStatUtil;
+import com.taobao.middleware.cli.annotations.DefaultValue;
+import com.taobao.middleware.cli.annotations.Description;
 import com.taobao.middleware.cli.annotations.Name;
+import com.taobao.middleware.cli.annotations.Option;
 import com.taobao.middleware.cli.annotations.Summary;
 
 import com.alibaba.arthas.tunnel.client.TunnelClient;
 
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
 /**
  * 查看会话状态命令
  *
  * @author vlinux on 15/5/3.
  */
 @Name("session")
 @Summary("Display current session information")
+@Description(value =Constants.EXAMPLE +
+        "  session -s\n" +
+        "  session -I *\n" +
+        "  session -I '${arthas-cmd}-args=10'\n" +
+        "  session -I 'history-args=10'\n" +
+        "  session -I 'trace-args=-n 3 -v'\n" +
+        Constants.WIKI + Constants.WIKI_HOME + "session")
 public class SessionCommand extends AnnotatedCommand {
 
+    private static List<String> RESERVED_WORDS = Arrays.asList( "id", "tty","pid","server","lastAccessedTime", 
+        "createTime","subject", "instrumentation", "arthas-command-manager");
+
+    private boolean showSessionData = true;
+
+    private boolean showReservedNames = false;
+
+    private List<String> customSesionData;
+
+    @Option(shortName = "H", longName = "hiddenSessionData", flag = true)
+    @Description("Hidden artahs session data")
+    public void setHiddenSessionData(final boolean hiddenSessionData) {
+        this.showSessionData = !hiddenSessionData;
+    }
+
+    @Option(shortName = "A", longName = "showAll", flag = true)
+    @Description("Show all session data")
+    public void setShowAll(final boolean showReservedNames) {
+        this.showReservedNames = showReservedNames;
+    }
+
+    @Option(shortName = "I", longName = "input")
+    @Description("Input cutom session data, ie: 'trace-args=-n 3' or * for cleanup")
+    public void setCustomSessionData(final List<String> customSesionData) {
+        this.customSesionData = customSesionData;
+    }
+
     @Override
     public void process(CommandProcess process) {
         SessionModel result = new SessionModel();
         Session session = process.session();
         result.setJavaPid(session.getPid());
         result.setSessionId(session.getSessionId());
-
+        if (customSesionData != null) {
+            process.echoTips("customSesionData: " + customSesionData + "\n");
+            if (customSesionData.size() == 1 && "*".equals(customSesionData.get(0))) {
+                final Map<String, Object> sessionData = session.cloneSessionData();
+                for (final Map.Entry<String, Object> entry : sessionData.entrySet()) {
+                     if (!RESERVED_WORDS.contains(entry.getKey())) {
+                        session.remove(entry.getKey());
+                     }
+                }
+            } else {
+                final Pattern pattern = Pattern.compile("\\s*([^:=]+)\\s*[:=]\\s*(.+)");
+                for (final String pv : customSesionData) {
+                    final Matcher matcher = pattern.matcher(pv);
+                    if (matcher.find()) {
+                        final int groupCount = matcher.groupCount();
+                        if (groupCount >= 2) {
+                            final String name = matcher.group(1);
+                            if (RESERVED_WORDS.contains(name)){
+                                process.echoTips("ignore: " + name);
+                                return;
+                            }
+                            session.put(name, matcher.group(2));
+                        }
+                    }
+                }
+            }
+            showSessionData = true;
+        }
+        if (showSessionData) {
+            final Map<String, Object> sessionData = session.cloneSessionData();
+            if (!this.showReservedNames) {
+                final Iterator<Entry<String, Object>> iter = sessionData.entrySet().iterator();
+                while(iter.hasNext()) {
+                    if (RESERVED_WORDS.contains(iter.next().getKey())) {
+                        iter.remove();
+                        continue;
+                    }
+                }
+            }
+            result.setSessionData(sessionData);
+        }
         //tunnel
         TunnelClient tunnelClient = ArthasBootstrap.getInstance().getTunnelClient();
         if (tunnelClient != null) {

diff --git a/core/src/main/java/com/taobao/arthas/core/command/model/SessionModel.java b/core/src/main/java/com/taobao/arthas/core/command/model/SessionModel.java
@@ -1,5 +1,7 @@
 package com.taobao.arthas.core.command.model;
 
+import java.util.Map;
+
 /**
  * Session command result model
  *
@@ -15,6 +17,8 @@ public class SessionModel extends ResultModel {
 
     private boolean tunnelConnected;
 
+    private Map<String, Object> sessionData;
+
     @Override
     public String getType() {
         return "session";
@@ -67,4 +71,12 @@ public boolean isTunnelConnected() {
     public void setTunnelConnected(boolean tunnelConnected) {
         this.tunnelConnected = tunnelConnected;
     }
+
+    public Map<String, Object> getSessionData() {
+        return this.sessionData;
+    }
+
+    public void setSessionData(Map<String, Object> env) {
+        this.sessionData = env;
+    }
 }
diff --git a/core/src/main/java/com/taobao/arthas/core/command/view/SessionView.java b/core/src/main/java/com/taobao/arthas/core/command/view/SessionView.java
@@ -31,6 +31,9 @@ public void draw(CommandProcess process, SessionModel result) {
         if (result.getStatUrl() != null) {
             table.row("STAT_URL", result.getStatUrl());
         }
+        if (result.getSessionData() != null) {
+             table.row("SESION_ENV", result.getSessionData().toString());
+        }
         process.write(RenderUtil.render(table, process.width()));
     }
 

diff --git a/core/src/main/java/com/taobao/arthas/core/shell/session/Session.java b/core/src/main/java/com/taobao/arthas/core/shell/session/Session.java
@@ -6,6 +6,7 @@
 
 import java.lang.instrument.Instrumentation;
 import java.util.List;
+import java.util.Map;
 
 /**
  * A shell session.
@@ -168,4 +169,9 @@ public interface Session {
      * Whether the session is tty term
      */
     boolean isTty();
+
+    /**
+     * @return Clone the session data.
+     */
+    Map<String, Object> cloneSessionData();
 }
diff --git a/core/src/main/java/com/taobao/arthas/core/shell/session/impl/SessionImpl.java b/core/src/main/java/com/taobao/arthas/core/shell/session/impl/SessionImpl.java
@@ -28,6 +28,11 @@ public SessionImpl() {
         this.setLastAccessTime(now);
     }
 
+    @Override
+    public Map<String, Object> cloneSessionData() {
+        return new HashMap<String, Object>(data);
+    }
+
     @Override
     public Session put(String key, Object obj) {
         if (obj == null) {