From 1dd7f6c9bd31c4a507f46277b82ba10e655f19f8 Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Wed, 27 May 2020 00:04:41 +0800
Subject: [PATCH] profiler command support duration option. close #1197

---
 .../command/monitor200/ProfilerCommand.java   | 52 ++++++++++++++++---
 1 file changed, 44 insertions(+), 8 deletions(-)

diff --git a/core/src/main/java/com/taobao/arthas/core/command/monitor200/ProfilerCommand.java b/core/src/main/java/com/taobao/arthas/core/command/monitor200/ProfilerCommand.java
index 66fd8e995..96c6a3bb0 100644
--- a/core/src/main/java/com/taobao/arthas/core/command/monitor200/ProfilerCommand.java
+++ b/core/src/main/java/com/taobao/arthas/core/command/monitor200/ProfilerCommand.java
@@ -10,11 +10,13 @@ import java.util.Date;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 
 import com.alibaba.arthas.deps.org.slf4j.Logger;
 import com.alibaba.arthas.deps.org.slf4j.LoggerFactory;
 import com.taobao.arthas.common.OSUtils;
 import com.taobao.arthas.core.command.Constants;
+import com.taobao.arthas.core.server.ArthasBootstrap;
 import com.taobao.arthas.core.shell.cli.CliToken;
 import com.taobao.arthas.core.shell.cli.Completion;
 import com.taobao.arthas.core.shell.cli.CompletionUtils;
@@ -91,6 +93,11 @@ public class ProfilerCommand extends AnnotatedCommand {
      */
     private boolean alluser;
 
+    /**
+     * run profiling for <duration> seconds
+     */
+    private Long duration;
+
     private static String libPath;
     private static AsyncProfiler profiler = null;
 
@@ -180,6 +187,12 @@ public class ProfilerCommand extends AnnotatedCommand {
         this.alluser = alluser;
     }
 
+    @Option(shortName = "d", longName = "duration")
+    @Description("run profiling for <duration> seconds")
+    public void setDuration(long duration) {
+        this.duration = duration;
+    }
+
     private AsyncProfiler profilerInstance() {
         if (profiler != null) {
             return profiler;
@@ -256,7 +269,7 @@ public class ProfilerCommand extends AnnotatedCommand {
     }
 
     @Override
-    public void process(CommandProcess process) {
+    public void process(final CommandProcess process) {
         int status = 0;
         try {
             ProfilerAction profilerAction = ProfilerAction.valueOf(action);
@@ -266,7 +279,7 @@ public class ProfilerCommand extends AnnotatedCommand {
                 return;
             }
 
-            AsyncProfiler asyncProfiler = this.profilerInstance();
+            final AsyncProfiler asyncProfiler = this.profilerInstance();
 
             if (ProfilerAction.execute.equals(profilerAction)) {
                 if (actionArg == null) {
@@ -280,13 +293,28 @@ public class ProfilerCommand extends AnnotatedCommand {
                 String executeArgs = executeArgs(ProfilerAction.start);
                 String result = execute(asyncProfiler, executeArgs);
                 process.write(result);
-            } else if (ProfilerAction.stop.equals(profilerAction)) {
-                if (this.file == null) {
-                    this.file = new File("arthas-output",
-                            new SimpleDateFormat("yyyyMMdd-HHmmss").format(new Date()) + "." + this.format)
-                                    .getAbsolutePath();
+
+                if (this.duration != null) {
+                    final String outputFile = outputFile();
+                    final String stopExecuteArgs = executeArgs(ProfilerAction.stop);
+                    process.write(String.format("profiler will silent stop after %d seconds.\n", this.duration.longValue()));
+                    process.write("profiler output file will be: " + new File(outputFile).getAbsolutePath() + "\n");
+                    ArthasBootstrap.getInstance().getScheduledExecutorService().schedule(new Runnable() {
+                        @Override
+                        public void run() {
+                            try {
+                                logger.info("profiler output file: " + new File(outputFile).getAbsolutePath() + "\n");
+                                String result = execute(asyncProfiler, stopExecuteArgs);
+                                logger.info("profiler stop result: " + result);
+                            } catch (Throwable e) {
+                                logger.error("", e);
+                            }
+                        }
+                    }, this.duration, TimeUnit.SECONDS);
                 }
-                process.write("profiler output file: " + new File(this.file).getAbsolutePath() + "\n");
+            } else if (ProfilerAction.stop.equals(profilerAction)) {
+                String outputFile = outputFile();
+                process.write("profiler output file: " + new File(outputFile).getAbsolutePath() + "\n");
                 String executeArgs = executeArgs(ProfilerAction.stop);
                 String result = execute(asyncProfiler, executeArgs);
                 process.write(result);
@@ -342,6 +370,14 @@ public class ProfilerCommand extends AnnotatedCommand {
         }
     }
 
+    private String outputFile() {
+        if (this.file == null) {
+            this.file = new File("arthas-output",
+                    new SimpleDateFormat("yyyyMMdd-HHmmss").format(new Date()) + "." + this.format).getAbsolutePath();
+        }
+        return file;
+    }
+
     private List<String> events() {
         List<String> result = new ArrayList<String>();