feat: PyroscopeAgent.stop (#149)

.github/workflows/test.yml

@@ -13,7 +13,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu, macos]
-        java: ['8', '11', '17']
+        java: ['8', '11', '17', '21']
     steps:
       - uses: actions/checkout@v2
       - name: Set up JDK
@@ -26,4 +26,4 @@ jobs:
       - name: Test with Gradle
         uses: gradle/gradle-build-action@937999e9cc2425eddc7fd62d1053baf041147db7
         with:
-          arguments: test
+          arguments: test --stacktrace

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

@@ -7,10 +7,10 @@
 import io.pyroscope.javaagent.impl.*;
 
 import java.lang.instrument.Instrumentation;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 public class PyroscopeAgent {
-    private static final AtomicBoolean started = new AtomicBoolean(false);
+    private static final Object sLock = new Object();
+    private static Options sOptions = null;
 
     public static void premain(final String agentArgs,
                                final Instrumentation inst) {
@@ -34,23 +34,50 @@ public static void start(Config config) {
     }
 
     public static void start(Options options) {
-        Logger logger = options.logger;
+        synchronized (sLock) {
+            Logger logger = options.logger;
 
-        if (!options.config.agentEnabled) {
-            logger.log(Logger.Level.INFO, "Pyroscope agent start disabled by configuration");
-            return;
+            if (!options.config.agentEnabled) {
+                logger.log(Logger.Level.INFO, "Pyroscope agent start disabled by configuration");
+                return;
+            }
+
+            if (sOptions != null) {
+                logger.log(Logger.Level.ERROR, "Failed to start profiling - already started");
+                return;
+            }
+            sOptions = options;
+            logger.log(Logger.Level.DEBUG, "Config: %s", options.config);
+            try {
+                options.scheduler.start(options.profiler);
+                logger.log(Logger.Level.INFO, "Profiling started");
+            } catch (final Throwable e) {
+                logger.log(Logger.Level.ERROR, "Error starting profiler %s", e);
+                sOptions = null;
+            }
         }
+    }
 
-        if (!started.compareAndSet(false, true)) {
-            logger.log(Logger.Level.ERROR, "Failed to start profiling - already started");
-            return;
+    public static void stop() {
+        synchronized (sLock) {
+            if (sOptions == null) {
+                DefaultLogger.PRECONFIG_LOGGER.log(Logger.Level.ERROR, "Error stopping profiler: not started");
+                return;
+            }
+            try {
+                sOptions.scheduler.stop();
+                sOptions.logger.log(Logger.Level.INFO, "Profiling stopped");
+            } catch (Throwable e) {
+                sOptions.logger.log(Logger.Level.ERROR, "Error stopping profiler %s", e);
+            }
+
+            sOptions = null;
         }
-        logger.log(Logger.Level.DEBUG, "Config: %s", options.config);
-        try {
-            options.scheduler.start(options.profiler);
-            logger.log(Logger.Level.INFO, "Profiling started");
-        } catch (final Throwable e) {
-            logger.log(Logger.Level.ERROR, "Error starting profiler %s", e);
+    }
+
+    public static boolean isStarted() {
+        synchronized (sLock) {
+            return sOptions != null;
         }
     }
 
@@ -65,12 +92,14 @@ public static class Options {
         final ProfilingScheduler scheduler;
         final Logger logger;
         final Profiler profiler;
+        final Exporter exporter;
 
         private Options(Builder b) {
             this.config = b.config;
             this.profiler = b.profiler;
             this.scheduler = b.scheduler;
             this.logger = b.logger;
+            this.exporter = b.exporter;
         }
 
         public static class Builder {

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

@@ -36,4 +36,6 @@ public interface ProfilingScheduler {
      *
      **/
     void start(Profiler profiler);
+
+    void stop();
 }

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

@@ -10,25 +10,26 @@
 
 import java.time.Duration;
 import java.time.Instant;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledFuture;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 
 
 public class ContinuousProfilingScheduler implements ProfilingScheduler {
-    final Config config;
-
-    final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(r -> {
+    public static final ThreadFactory THREAD_FACTORY = r -> {
         Thread t = Executors.defaultThreadFactory().newThread(r);
         t.setName("PyroscopeProfilingScheduler");
         t.setDaemon(true);
         return t;
-    });
+    };
+    private final Config config;
+
+    private ScheduledExecutorService executor;
     private final Exporter exporter;
-    private Logger logger;
+    private final Logger logger;
+    private final Object lock = new Object();
     private Instant profilingIntervalStartTime;
     private ScheduledFuture<?> job;
+    private boolean started;
+    private Profiler profiler;
 
     public ContinuousProfilingScheduler(Config config, Exporter exporter, Logger logger) {
         this.config = config;
@@ -38,14 +39,85 @@ public ContinuousProfilingScheduler(Config config, Exporter exporter, Logger log
 
     @Override
     public void start(Profiler profiler) {
-        Duration firstProfilingDuration;
+        this.logger.log(Logger.Level.DEBUG, "ContinuousProfilingScheduler starting");
+        synchronized (lock) {
+            if (started) {
+                throw new IllegalStateException("already started");
+            }
+            Duration firstProfilingDuration;
+            try {
+                firstProfilingDuration = startFirst(profiler);
+            } catch (Throwable throwable) {
+                stopSchedulerLocked();
+                throw new IllegalStateException(throwable);
+            }
+            this.profiler = profiler;
+            this.executor = Executors.newSingleThreadScheduledExecutor(THREAD_FACTORY);
+            this.job = executor.scheduleAtFixedRate(this::schedulerTick,
+                firstProfilingDuration.toMillis(), config.uploadInterval.toMillis(), TimeUnit.MILLISECONDS);
+            this.started = true;
+            logger.log(Logger.Level.DEBUG, "ContinuousProfilingScheduler started");
+        }
+    }
+
+    @Override
+    public void stop() {
+        ScheduledExecutorService svc = null;
+        try {
+            synchronized (lock) {
+                try {
+                    stopSchedulerLocked();
+                } finally {
+                    svc = this.executor;
+                    this.executor = null;
+                }
+            }
+            this.logger.log(Logger.Level.DEBUG, "ContinuousProfilingScheduler stopped");
+        } finally {
+            // shutdown here not under lock to avoid deadlock ( the task may block to wait for lock and
+            // we are holding the lock and waiting for task to finish)
+            // There is still synchronization happens from the PyroscopeAgent class,
+            // so there are no concurrent calls to start/stop. So there is no lock here
+            awaitTermination(svc);
+        }
+    }
+
+    private static void awaitTermination(ScheduledExecutorService svc) {
+        try {
+            boolean terminated = svc.awaitTermination(10, TimeUnit.SECONDS);
+            if (!terminated) {
+                throw new IllegalStateException("failed to terminate scheduler's executor");
+            }
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new IllegalStateException("failed to terminate scheduler's executor", e);
+        }
+    }
+
+    private void stopSchedulerLocked() {
+        if (!this.started) {
+            return;
+        }
+        this.logger.log(Logger.Level.DEBUG, "ContinuousProfilingScheduler stopping");
         try {
-            firstProfilingDuration = startFirst(profiler);
+            this.profiler.stop();
         } catch (Throwable throwable) {
-            stop();
             throw new IllegalStateException(throwable);
+        } finally {
+            job.cancel(true);
+            executor.shutdown();
+            this.started = false;
         }
-        final Runnable dumpProfile = () -> {
+    }
+
+
+    private void schedulerTick() {
+
+        synchronized (lock) {
+            if (!started) {
+                return;
+            }
+            logger.log(Logger.Level.DEBUG, "ContinuousProfilingScheduler#schedulerTick");
             Snapshot snapshot;
             Instant now;
             try {
@@ -55,25 +127,15 @@ public void start(Profiler profiler) {
                 profiler.start();
             } catch (Throwable throwable) {
                 logger.log(Logger.Level.ERROR, "Error dumping profiler %s", throwable);
-                stop();
+                stopSchedulerLocked();
                 return;
             }
             profilingIntervalStartTime = now;
             exporter.export(snapshot);
-        };
-
-        job = executor.scheduleAtFixedRate(dumpProfile,
-            firstProfilingDuration.toMillis(), config.uploadInterval.toMillis(), TimeUnit.MILLISECONDS);
-
-    }
-
-    private void stop() {
-        if (job != null) {
-            job.cancel(true);
         }
-        executor.shutdown();
     }
 
+
     /**
      * Starts the first profiling interval.
      * profilingIntervalStartTime is set to now
@@ -86,7 +148,7 @@ private Duration startFirst(Profiler profiler) {
 
         long uploadIntervalMillis = config.uploadInterval.toMillis();
         float randomOffset = Random.Default.nextFloat();
-        uploadIntervalMillis = (long)((float)uploadIntervalMillis * randomOffset);
+        uploadIntervalMillis = (long) ((float) uploadIntervalMillis * randomOffset);
         if (uploadIntervalMillis < 2000) {
             uploadIntervalMillis = 2000;
         }

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

@@ -70,13 +70,18 @@ public void start(Profiler profiler) {
         );
     }
 
+    @Override
+    public void stop() {
+        throw new RuntimeException("not implemented");
+    }
+
     private void dumpProfile(final Profiler profiler, final long samplingDurationMillis, final Duration uploadInterval) {
         Instant profilingStartTime = Instant.now();
         try {
             profiler.start();
         } catch (Throwable e) {
             logger.log(Logger.Level.ERROR, "Error starting profiler %s", e);
-            stop();
+            stopProfiling();
             return;
         }
         try {
@@ -90,7 +95,7 @@ private void dumpProfile(final Profiler profiler, final long samplingDurationMil
         exporter.export(snapshot);
     }
 
-    private void stop() {
+    private void stopProfiling() {
         if (job != null) {
             job.cancel(true);
         }

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

@@ -3,6 +3,7 @@
 import io.pyroscope.javaagent.api.Logger;
 import io.pyroscope.javaagent.api.ProfilingScheduler;
 import io.pyroscope.javaagent.config.Config;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -45,6 +46,11 @@ void setUp() {
             .build();
     }
 
+    @AfterEach
+    void tearDown() {
+        PyroscopeAgent.stop();
+    }
+
     @Test
     void startupTestWithEnabledAgent() {
         PyroscopeAgent.start(optionsAgentEnabled);

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

@@ -0,0 +1,52 @@
+package io.pyroscope.javaagent;
+
+import io.pyroscope.http.Format;
+import io.pyroscope.javaagent.api.Logger;
+import io.pyroscope.javaagent.config.Config;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class StartStopTest {
+
+    public static final Config INVALID = new Config.Builder()
+        .setApplicationName("demo.app{qweqwe=asdasd}")
+        .setFormat(Format.JFR)
+        .setProfilingAlloc("512k")
+        .setAPExtraArguments("event=qwe") // java.lang.IllegalArgumentException: Duplicate event argument
+        .setProfilingEvent(EventType.ITIMER)
+        .setLogLevel(Logger.Level.DEBUG)
+        .build();
+
+    public static final Config VALID = new Config.Builder()
+        .setApplicationName("demo.app{qweqwe=asdasd}")
+        .setFormat(Format.JFR)
+        .setProfilingEvent(EventType.ITIMER)
+        .setLogLevel(Logger.Level.DEBUG)
+        .build();
+
+
+    @Test
+    void testStartFail() {
+        assertFalse(PyroscopeAgent.isStarted());
+
+        PyroscopeAgent.start(INVALID);
+        assertFalse(PyroscopeAgent.isStarted());
+
+        PyroscopeAgent.start(INVALID);
+        assertFalse(PyroscopeAgent.isStarted());
+
+        PyroscopeAgent.stop();
+        assertFalse(PyroscopeAgent.isStarted());
+        PyroscopeAgent.stop();
+        assertFalse(PyroscopeAgent.isStarted());
+
+        PyroscopeAgent.start(VALID);
+        assertTrue(PyroscopeAgent.isStarted());
+
+        PyroscopeAgent.stop();
+        assertFalse(PyroscopeAgent.isStarted());
+    }
+
+}