content: package com.taobao.arthas.core.command.monitor200;
  
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.FileOutputStream;
  import java.lang.instrument.Instrumentation;
  import java.security.CodeSource;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  
  import com.alibaba.arthas.deps.org.slf4j.Logger;
  import com.alibaba.arthas.deps.org.slf4j.LoggerFactory;
  import com.taobao.arthas.common.IOUtils;
  import com.taobao.arthas.common.VmToolUtils;
  import com.taobao.arthas.core.command.Constants;
  import com.taobao.arthas.core.command.express.Express;
  import com.taobao.arthas.core.command.express.ExpressException;
  import com.taobao.arthas.core.command.express.ExpressFactory;
  import com.taobao.arthas.core.command.model.ClassLoaderVO;
  import com.taobao.arthas.core.command.model.ObjectVO;
  import com.taobao.arthas.core.command.model.VmToolModel;
  import com.taobao.arthas.core.shell.cli.Completion;
  import com.taobao.arthas.core.shell.cli.CompletionUtils;
  import com.taobao.arthas.core.shell.cli.OptionCompleteHandler;
  import com.taobao.arthas.core.shell.command.AnnotatedCommand;
  import com.taobao.arthas.core.shell.command.CommandProcess;
  import com.taobao.arthas.core.util.ClassLoaderUtils;
  import com.taobao.arthas.core.util.ClassUtils;
  import com.taobao.arthas.core.util.SearchUtils;
  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 arthas.VmTool;
  
  /**
   * 
   * @author hengyunabc 2021-04-27
   * @author ZhangZiCheng 2021-04-29
   *
   */
  //@formatter:off
  @Name("vmtool")
  @Summary("jvm tool")
  @Description(Constants.EXAMPLE
          + "  vmtool --action getInstances --className demo.MathGame\n"
          + "  vmtool --action getInstances --className demo.MathGame --express 'instances.length'\n"
          + "  vmtool --action getInstances --className demo.MathGame --express 'instances[0]'\n"
          + "  vmtool --action getInstances --className demo.MathGame -x 2\n"
          + "  vmtool --action getInstances --className java.lang.String --limit 10\n"
          + "  vmtool --action getInstances --classLoaderClass org.springframework.boot.loader.LaunchedURLClassLoader --className org.springframework.context.ApplicationContext\n"
          + "  vmtool --action forceGc\n"
          + "  vmtool --action interruptThread -t 1\n"
          + "  vmtool --action mallocTrim\n"
          + "  vmtool --action mallocStats\n"
          + Constants.WIKI + Constants.WIKI_HOME + "vmtool")
  //@formatter:on
  public class VmToolCommand extends AnnotatedCommand {
      private static final Logger logger = LoggerFactory.getLogger(VmToolCommand.class);
  
      private VmToolAction action;
      private String className;
      private String express;
      private int threadId;
      private String hashCode = null;
      private String classLoaderClass;
      /**
       * default value 1
       */
      private int expand;
  
      /**
       * default value 10
       */
      private int limit;
  
      private String libPath;
      private static String defaultLibPath;
      private static VmTool vmTool = null;
  
      static {
          String libName = VmToolUtils.detectLibName();
          if (libName != null) {
              CodeSource codeSource = VmToolCommand.class.getProtectionDomain().getCodeSource();
              if (codeSource != null) {
                  try {
                      File bootJarPath = new File(codeSource.getLocation().toURI().getSchemeSpecificPart());
                      File soFile = new File(bootJarPath.getParentFile(), "lib" + File.separator + libName);
                      if (soFile.exists()) {
                          defaultLibPath = soFile.getAbsolutePath();
                      }
                  } catch (Throwable e) {
                      logger.error("can not find VmTool so", e);
                  }
              }
          }
  
      }
  
      @Option(shortName = "a", longName = "action", required = true)
      @Description("Action to execute")
      public void setAction(VmToolAction action) {
          this.action = action;
      }
  
      @Option(longName = "className")
      @Description("The class name")
      public void setClassName(String className) {
          this.className = className;
      }
  
      @Option(shortName = "x", longName = "expand")
      @Description("Expand level of object (1 by default)")
      @DefaultValue("1")
      public void setExpand(int expand) {
          this.expand = expand;
      }
  
      @Option(shortName = "c", longName = "classloader")
      @Description("The hash code of the special class's classLoader")
      public void setHashCode(String hashCode) {
          this.hashCode = hashCode;
      }
  
      @Option(longName = "classLoaderClass")
      @Description("The class name of the special class's classLoader.")
      public void setClassLoaderClass(String classLoaderClass) {
          this.classLoaderClass = classLoaderClass;
      }
  
      @Option(shortName = "l", longName = "limit")
      @Description("Set the limit value of the getInstances action, default value is 10, set to -1 is unlimited")
      @DefaultValue("10")
      public void setLimit(int limit) {
          this.limit = limit;
      }
  
      @Option(longName = "libPath")
      @Description("The specify lib path.")
      public void setLibPath(String path) {
          libPath = path;
      }
  
      @Option(longName = "express", required = false)
      @Description("The ognl expression, default value is `instances`.")
      public void setExpress(String express) {
          this.express = express;
      }
  
      @Option(shortName = "t", longName = "threadId", required = false)
      @Description("The id of the thread to be interrupted")
      public void setThreadId(int threadId) {
          this.threadId = threadId;
      }
  
      public enum VmToolAction {
          getInstances, forceGc, interruptThread, mallocTrim, mallocStats
      }
  
      @Override
      public void process(final CommandProcess process) {
          try {
              Instrumentation inst = process.session().getInstrumentation();
  
              if (VmToolAction.getInstances.equals(action)) {
                  if (className == null) {
                      process.end(-1, "The className option cannot be empty!");
                      return;
                  }
                  ClassLoader classLoader = null;
                  if (hashCode != null) {
                      classLoader = ClassLoaderUtils.getClassLoader(inst, hashCode);
                      if (classLoader == null) {
                          process.end(-1, "Can not find classloader with hashCode: " + hashCode + ".");
                          return;
                      }
                  }else if ( classLoaderClass != null) {
                      List<ClassLoader> matchedClassLoaders = ClassLoaderUtils.getClassLoaderByClassName(inst,
                              classLoaderClass);
                      if (matchedClassLoaders.size() == 1) {
                          classLoader = matchedClassLoaders.get(0);
                          hashCode = Integer.toHexString(matchedClassLoaders.get(0).hashCode());
                      } else if (matchedClassLoaders.size() > 1) {
                          Collection<ClassLoaderVO> classLoaderVOList = ClassUtils
                                  .createClassLoaderVOList(matchedClassLoaders);
  
                          VmToolModel vmToolModel = new VmToolModel().setClassLoaderClass(classLoaderClass)
                                  .setMatchedClassLoaders(classLoaderVOList);
                          process.appendResult(vmToolModel);
                          process.end(-1,
                                  "Found more than one classloader by class name, please specify classloader with '-c <classloader hash>'");
                          return;
                      } else {
                          process.end(-1, "Can not find classloader by class name: " + classLoaderClass + ".");
                          return;
                      }
                  }else {
                      classLoader = ClassLoader.getSystemClassLoader();
                  }
  
                  List<Class<?>> matchedClasses = new ArrayList<Class<?>>(
                          SearchUtils.searchClassOnly(inst, className, false, hashCode));
                  int matchedClassSize = matchedClasses.size();
                  if (matchedClassSize == 0) {
                      process.end(-1, "Can not find class by class name: " + className + ".");
                      return;
                  } else if (matchedClassSize > 1) {
                      process.end(-1, "Found more than one class: " + matchedClasses + ", please specify classloader with '-c <classloader hash>'");
                      return;
                  } else {
                      Object[] instances = vmToolInstance().getInstances(matchedClasses.get(0), limit);
                      Object value = instances;
                      if (express != null) {
                          Express unpooledExpress = ExpressFactory.unpooledExpress(classLoader);
                          try {
                              value = unpooledExpress.bind(new InstancesWrapper(instances)).get(express);
                          } catch (ExpressException e) {
                              logger.warn("ognl: failed execute express: " + express, e);
                              process.end(-1, "Failed to execute ognl, exception message: " + e.getMessage()
                                      + ", please check $HOME/logs/arthas/arthas.log for more details. ");
                          }
                      }
  
                      VmToolModel vmToolModel = new VmToolModel().setValue(new ObjectVO(value, expand));
                      process.appendResult(vmToolModel);
                      process.end();
                  }
              } else if (VmToolAction.forceGc.equals(action)) {
                  vmToolInstance().forceGc();
                  process.write("\n");
                  process.end();
                  return;
              } else if (VmToolAction.interruptThread.equals(action)) {
                  vmToolInstance().interruptSpecialThread(threadId);
                  process.write("\n");
                  process.end();
  
                  return;
              } else if (VmToolAction.mallocTrim.equals(action)) {
                  int result = vmToolInstance().mallocTrim();
                  process.write("\n");
                  process.end(result == 1 ? 0 : -1, "mallocTrim result: " +
                      (result == 1 ? "true" : (result == 0 ? "false" : "not supported")));
                  return;
              } else if (VmToolAction.mallocStats.equals(action)) {
                  boolean result = vmToolInstance().mallocStats();
                  process.write("\n");
                  process.end(result ? 0 : -1, "mallocStats result: " +
                      (result ? "true" : "not supported"));
                  return;
              }
  
              process.end();
          } catch (Throwable e) {
              logger.error("vmtool error", e);
              process.end(1, "vmtool error: " + e.getMessage());
          }
      }
  
      static class InstancesWrapper {
          Object instances;
  
          public InstancesWrapper(Object instances) {
              this.instances = instances;
          }
  
          public Object getInstances() {
              return instances;
          }
  
          public void setInstances(Object instances) {
              this.instances = instances;
          }
      }
  
      private VmTool vmToolInstance() {
          if (vmTool != null) {
              return vmTool;
          } else {
              if (libPath == null) {
                  libPath = defaultLibPath;
              }
  
              // 尝试把lib文件复制到临时文件里，避免多次attach时出现 Native Library already loaded in another classloader
              FileOutputStream tmpLibOutputStream = null;
              FileInputStream libInputStream = null;
              try {
                  File tmpLibFile = File.createTempFile(VmTool.JNI_LIBRARY_NAME, null);
                  tmpLibOutputStream = new FileOutputStream(tmpLibFile);
                  libInputStream = new FileInputStream(libPath);
  
                  IOUtils.copy(libInputStream, tmpLibOutputStream);
                  libPath = tmpLibFile.getAbsolutePath();
                  logger.debug("copy {} to {}", libPath, tmpLibFile);
              } catch (Throwable e) {
                  logger.error("try to copy lib error! libPath: {}", libPath, e);
              } finally {
                  IOUtils.close(libInputStream);
                  IOUtils.close(tmpLibOutputStream);
              }
  
              vmTool = VmTool.getInstance(libPath);
          }
          return vmTool;
      }
  
      private Set<String> actions() {
          Set<String> values = new HashSet<String>();
          for (VmToolAction action : VmToolAction.values()) {
              values.add(action.toString());
          }
          return values;
      }
  
      @Override
      public void complete(Completion completion) {
          List<OptionCompleteHandler> handlers = new ArrayList<OptionCompleteHandler>();
  
          handlers.add(new OptionCompleteHandler() {
  
              @Override
              public boolean matchName(String token) {
                  return "-a".equals(token) || "--action".equals(token);
              }
  
              @Override
              public boolean complete(Completion completion) {
                  return CompletionUtils.complete(completion, actions());
              }
  
          });
  
          handlers.add(new OptionCompleteHandler() {
              @Override
              public boolean matchName(String token) {
                  return "--className".equals(token);
              }
  
              @Override
              public boolean complete(Completion completion) {
                  return CompletionUtils.completeClassName(completion);
              }
          });
  
          if (CompletionUtils.completeOptions(completion, handlers)) {
              return;
          }
  
          super.complete(completion);
      }
  
  }
  