Exposed the AsyncProfiler stop method

agent/src/main/java/io/pyroscope/javaagent/PyroscopeAgent.java

@@ -7,10 +7,13 @@
 import io.pyroscope.javaagent.impl.*;
 
 import java.lang.instrument.Instrumentation;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class PyroscopeAgent {
-    private static final AtomicBoolean started = new AtomicBoolean(false);
+    // this is used to store the options passed to the agent
+    private static final AtomicReference<Options> startOptions = new AtomicReference<>(null);
+
+    private static final String stopLock = "";
 
     public static void premain(final String agentArgs,
                                final Instrumentation inst) {
@@ -35,14 +38,12 @@ public static void start(Config config) {
 
     public static void start(Options options) {
         Logger logger = options.logger;
-
-        if (!options.config.agentEnabled) {
-            logger.log(Logger.Level.INFO, "Pyroscope agent start disabled by configuration");
+        if (!startOptions.compareAndSet(null, options)) {
+            logger.log(Logger.Level.ERROR, "Failed to start profiling - already started");
             return;
         }
-
-        if (!started.compareAndSet(false, true)) {
-            logger.log(Logger.Level.ERROR, "Failed to start profiling - already started");
+        if (!options.config.agentEnabled) {
+            logger.log(Logger.Level.INFO, "Pyroscope agent start disabled by configuration");
             return;
         }
         logger.log(Logger.Level.DEBUG, "Config: %s", options.config);
@@ -54,6 +55,33 @@ public static void start(Options options) {
         }
     }
 
+    /**
+     * stop is used to stop profiling
+     */
+    public static void stop() {
+        synchronized (stopLock) {
+            if (startOptions.get() == null) {
+                return;
+            }
+            ProfilingScheduler scheduler = startOptions.get().scheduler;
+            Logger logger = startOptions.get().logger;
+            Profiler profiler = startOptions.get().profiler;
+
+            if (logger == null) {
+                return;
+            }
+            if (scheduler == null || profiler == null) {
+                logger.log(Logger.Level.ERROR, "Failed to stop profiling - already stopped");
+                return;
+            }
+
+            logger.log(Logger.Level.DEBUG, "Config: %s", startOptions.get().config);
+            scheduler.stop(profiler);
+            startOptions.set(null);
+            logger.log(Logger.Level.INFO, "Profiling stopped");
+        }
+    }
+
     /**
      * Options allow to swap pyroscope components:
      * - io.pyroscope.javaagent.api.ProfilingScheduler

agent/src/main/java/io/pyroscope/javaagent/api/ProfilingScheduler.java

@@ -36,4 +36,11 @@ public interface ProfilingScheduler {
      *
      **/
     void start(Profiler profiler);
+
+/**
+ * stop is ised tp stp[ the profiler
+ * {@link Profiler#stop()}
+ * {@link Profiler#dumpProfile(Instant, Instant)}
+ * **/
+    void stop(Profiler profiler);
 }

agent/src/main/java/io/pyroscope/javaagent/impl/ContinuousProfilingScheduler.java

@@ -67,6 +67,15 @@ public void start(Profiler profiler) {
 
     }
 
+    /**
+     * Stops the profiling scheduler which stops the AsyncProfiler.
+     * @param profiler
+     */
+    @Override
+    public void stop(Profiler profiler) {
+        profiler.stop();
+    }
+
     private void stop() {
         if (job != null) {
             job.cancel(true);

agent/src/main/java/io/pyroscope/javaagent/impl/SamplingProfilingScheduler.java

@@ -70,6 +70,15 @@ public void start(Profiler profiler) {
         );
     }
 
+    /**
+     * Stops the profiling scheduler which stops the AsyncProfiler.
+     * @param profiler
+     */
+    @Override
+    public void stop(Profiler profiler) {
+        profiler.stop();
+    }
+
     private void dumpProfile(final Profiler profiler, final long samplingDurationMillis, final Duration uploadInterval) {
         Instant profilingStartTime = Instant.now();
         try {

agent/src/test/java/io/pyroscope/javaagent/PyroscopeAgentTest.java

@@ -50,6 +50,10 @@ void startupTestWithEnabledAgent() {
         PyroscopeAgent.start(optionsAgentEnabled);
 
         verify(profilingScheduler, times(1)).start(any());
+
+        PyroscopeAgent.stop();
+
+        verify(profilingScheduler, times(1)).stop(any());
     }
 
     @Test
@@ -58,4 +62,5 @@ void startupTestWithDisabledAgent() {
 
         verify(profilingScheduler, never()).start(any());
     }
+
 }

demo/src/main/java/App.java

@@ -16,7 +16,7 @@
 public class App {
     public static final int N_THREADS = 8;
 
-    public static void main(String[] args) {
+    public static void main(String[] args) throws InterruptedException {
         PyroscopeAgent.start(
             new PyroscopeAgent.Options.Builder(
                 new Config.Builder()
@@ -32,6 +32,39 @@ public static void main(String[] args) {
         Pyroscope.setStaticLabels(mapOf("region", "us-east-1"));
 
         appLogic();
+
+        // Test this by creating a new thread, stop the agent,  sleeping for a second, and then restarting the agent.
+        Thread newThread = new Thread(new Runnable() {
+            public void run()
+            {
+                // Sleep for a second to make sure the agent has started
+                try {
+                    Thread.sleep(1000L);
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+                // Stop the agent after a second
+                PyroscopeAgent.stop();
+                try {
+                    Thread.sleep(1000L);
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+                // Start the agent back up again after a second
+                PyroscopeAgent.start(
+                    new PyroscopeAgent.Options.Builder(
+                        new Config.Builder()
+                            .setApplicationName("demo.app{qweqwe=asdasd}")
+                            .setServerAddress("http://localhost:4040")
+                            .setFormat(Format.JFR)
+                            .setLogLevel(Logger.Level.DEBUG)
+                            .setLabels(mapOf("user", "tolyan"))
+                            .build())
+//                .setExporter(new MyStdoutExporter())
+                        .build()
+                );
+            }});
+        newThread.start();
     }
 
     private static void appLogic() {