[OpAMP] Enabling and disabling call graphs via remote config

profiler/src/main/java/com/splunk/opentelemetry/profiler/ProfilingSupervisor.java

@@ -157,16 +157,19 @@ private void setJfrContextStorageEnabled(boolean enabled) {
    * request.
    */
   private void tryStart() {
-    updateJvmMemoryMetrics();
     if (isJfrRecordingActive()) {
       logger.fine("JFR is already running, not starting again.");
       return;
     }
+
+    updateJvmMemoryMetrics();
+
     if (!jfr.isAvailable()) {
       logger.warning(
           "JDK Flight Recorder (JFR) is not available in this JVM. Profiling will not start.");
       return;
     }
+
     configSupplier.get().log();
     setJfrContextStorageEnabled(true);
     activateJfrRecording(getResource(sdk));

profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ActiveSpanTracker.java

@@ -31,6 +31,8 @@ class ActiveSpanTracker implements ContextStorage, SpanTracker {
   private final ContextStorage delegate;
   private final TraceRegistry registry;
 
+  private boolean enabled;
+
   ActiveSpanTracker(ContextStorage delegate, TraceRegistry registry) {
     this.delegate = delegate;
     this.registry = registry;
@@ -39,6 +41,10 @@ class ActiveSpanTracker implements ContextStorage, SpanTracker {
   @Override
   public Scope attach(Context toAttach) {
     Scope scope = delegate.attach(toAttach);
+    if (!isEnabled()) {
+      return scope;
+    }
+
     SpanContext newSpanContext = Span.fromContext(toAttach).getSpanContext();
     if (doNotTrack(newSpanContext)) {
       return scope;
@@ -61,6 +67,16 @@ public Scope attach(Context toAttach) {
     };
   }
 
+  @Override
+  public void setEnabled(boolean enabled) {
+    this.enabled = enabled;
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return enabled;
+  }
+
   private boolean doNotTrack(SpanContext spanContext) {
     return !spanContext.isSampled() || !registry.isRegistered(spanContext);
   }

profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ContextStorageWrapper.java

@@ -49,6 +49,9 @@ private ContextStorage trackActiveSpans(ContextStorage storage, TraceRegistry re
   }
 
   private ContextStorage detectThreadChanges(ContextStorage storage, TraceRegistry registry) {
-    return new TraceThreadChangeDetector(storage, registry, StackTraceSampler.SUPPLIER);
+    TraceThreadChangeDetector detector =
+        new TraceThreadChangeDetector(storage, registry, StackTraceSampler.SUPPLIER);
+    TraceThreadChangeDetector.SUPPLIER.configure(detector);
+    return detector;
   }
 }

profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingAgentListener.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.profiler.snapshot;
+
+import com.google.auto.service.AutoService;
+import com.google.common.annotations.VisibleForTesting;
+import io.opentelemetry.javaagent.extension.AgentListener;
+import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
+import java.util.function.Function;
+
+@AutoService(AgentListener.class)
+public class SnapshotProfilingAgentListener implements AgentListener {
+  private final Function<AutoConfiguredOpenTelemetrySdk, SnapshotProfilingSupervisor>
+      snapshotProfilingSupervisorMaker;
+
+  public SnapshotProfilingAgentListener() {
+    this(SnapshotProfilingSupervisor::initialize);
+  }
+
+  @VisibleForTesting
+  SnapshotProfilingAgentListener(
+      Function<AutoConfiguredOpenTelemetrySdk, SnapshotProfilingSupervisor>
+          snapshotProfilingSupervisorMaker) {
+    this.snapshotProfilingSupervisorMaker = snapshotProfilingSupervisorMaker;
+  }
+
+  @Override
+  public void afterAgent(AutoConfiguredOpenTelemetrySdk sdk) {
+    // Must be always executed to initialize supervisor
+    SnapshotProfilingSupervisor supervisor = snapshotProfilingSupervisorMaker.apply(sdk);
+
+    SnapshotProfilingConfiguration configuration = SnapshotProfilingConfiguration.SUPPLIER.get();
+    if (!configuration.isEnabled()) {
+      return;
+    }
+
+    supervisor.startProfiling();
+  }
+}

profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingConfigurationCustomizerProvider.java

@@ -49,19 +49,12 @@ OpenTelemetryConfigurationModel customizeModel(OpenTelemetryConfigurationModel m
     SnapshotProfilingConfiguration snapshotProfiling = getSnapshotProfilingConfig(model);
     SnapshotProfilingConfiguration.SUPPLIER.configure(snapshotProfiling);
 
-    if (snapshotProfiling.isEnabled()) {
-      initActiveSpansTracking();
-      initStackTraceSampler(snapshotProfiling);
-      addSpanProcessors(model);
-    }
+    initActiveSpansTracking();
+    addSpanProcessors(model);
 
     return model;
   }
 
-  private void initStackTraceSampler(SnapshotProfilingConfiguration snapshotProfilingConfig) {
-    StackTraceSamplerInitializer.setupStackTraceSampler(snapshotProfilingConfig);
-  }
-
   private void addSpanProcessors(OpenTelemetryConfigurationModel model) {
     TracerProviderModel tracerProviderModel = model.getTracerProvider();
     if (tracerProviderModel == null) {

profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSdkCustomizer.java

@@ -22,56 +22,24 @@
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder;
-import java.time.Duration;
 import java.util.Collections;
 import java.util.Map;
-import java.util.Set;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
 @AutoService(AutoConfigurationCustomizerProvider.class)
 public class SnapshotProfilingSdkCustomizer implements AutoConfigurationCustomizerProvider {
   private final TraceRegistry registry;
-  private final Function<ConfigProperties, StackTraceSampler> samplerProvider;
   private final ContextStorageWrapper contextStorageWrapper;
 
   public SnapshotProfilingSdkCustomizer() {
-    this(
-        TraceRegistryHolder.getTraceRegistry(),
-        stackTraceSamplerProvider(),
-        new ContextStorageWrapper());
-  }
-
-  private static Function<ConfigProperties, StackTraceSampler> stackTraceSamplerProvider() {
-    return properties -> {
-      SnapshotProfilingConfiguration configuration = SnapshotProfilingConfiguration.SUPPLIER.get();
-      Duration samplingPeriod = configuration.getSamplingInterval();
-      StagingArea.SUPPLIER.configure(createStagingArea(configuration));
-      return new PeriodicStackTraceSampler(
-          StagingArea.SUPPLIER, SpanTracker.SUPPLIER, samplingPeriod);
-    };
-  }
-
-  private static StagingArea createStagingArea(SnapshotProfilingConfiguration configuration) {
-    Duration interval = configuration.getExportInterval();
-    int capacity = configuration.getStagingCapacity();
-    return new PeriodicallyExportingStagingArea(StackTraceExporter.SUPPLIER, interval, capacity);
+    this(TraceRegistryHolder.getTraceRegistry(), new ContextStorageWrapper());
   }
 
   @VisibleForTesting
   SnapshotProfilingSdkCustomizer(
-      TraceRegistry registry,
-      StackTraceSampler sampler,
-      ContextStorageWrapper contextStorageWrapper) {
-    this(registry, properties -> sampler, contextStorageWrapper);
-  }
-
-  private SnapshotProfilingSdkCustomizer(
-      TraceRegistry registry,
-      Function<ConfigProperties, StackTraceSampler> samplerProvider,
-      ContextStorageWrapper contextStorageWrapper) {
+      TraceRegistry registry, ContextStorageWrapper contextStorageWrapper) {
     this.registry = registry;
-    this.samplerProvider = samplerProvider;
     this.contextStorageWrapper = contextStorageWrapper;
   }
 
@@ -82,7 +50,6 @@ public void customize(AutoConfigurationCustomizer autoConfigurationCustomizer) {
 
     autoConfigurationCustomizer
         .addTracerProviderCustomizer(snapshotProfilingSpanProcessor(registry))
-        .addPropertiesCustomizer(setupStackTraceSampler())
         .addPropertiesCustomizer(startTrackingActiveSpans(registry))
         .addTracerProviderCustomizer(addShutdownHook());
   }
@@ -99,54 +66,28 @@ public void customize(AutoConfigurationCustomizer autoConfigurationCustomizer) {
   private BiFunction<SdkTracerProviderBuilder, ConfigProperties, SdkTracerProviderBuilder>
       addShutdownHook() {
     return (builder, properties) -> {
-      if (snapshotProfilingEnabled()) {
-        builder.addSpanProcessor(new SdkShutdownHook());
-      }
+      builder.addSpanProcessor(new SdkShutdownHook());
       return builder;
     };
   }
 
   private BiFunction<SdkTracerProviderBuilder, ConfigProperties, SdkTracerProviderBuilder>
       snapshotProfilingSpanProcessor(TraceRegistry registry) {
     return (builder, properties) -> {
-      if (snapshotProfilingEnabled()) {
-        double selectionProbability =
-            SnapshotProfilingConfiguration.SUPPLIER.get().getSnapshotSelectionProbability();
-
-        return builder.addSpanProcessor(
-            new SnapshotProfilingSpanProcessor(
-                registry, new TraceIdBasedSnapshotSelector(selectionProbability)));
-      }
-      return builder;
-    };
-  }
-
-  private boolean includeTraceContextPropagator(Set<String> configuredPropagators) {
-    return configuredPropagators.isEmpty();
-  }
+      double selectionProbability =
+          SnapshotProfilingConfiguration.SUPPLIER.get().getSnapshotSelectionProbability();
 
-  private Function<ConfigProperties, Map<String, String>> setupStackTraceSampler() {
-    return properties -> {
-      if (snapshotProfilingEnabled()) {
-        StackTraceSampler sampler = samplerProvider.apply(properties);
-        ConfigurableSupplier<StackTraceSampler> supplier = StackTraceSampler.SUPPLIER;
-        supplier.configure(sampler);
-      }
-      return Collections.emptyMap();
+      return builder.addSpanProcessor(
+          new SnapshotProfilingSpanProcessor(
+              registry, new TraceIdBasedSnapshotSelector(selectionProbability)));
     };
   }
 
   private Function<ConfigProperties, Map<String, String>> startTrackingActiveSpans(
       TraceRegistry registry) {
     return properties -> {
-      if (snapshotProfilingEnabled()) {
-        contextStorageWrapper.wrapContextStorage(registry);
-      }
+      contextStorageWrapper.wrapContextStorage(registry);
       return Collections.emptyMap();
     };
   }
-
-  private boolean snapshotProfilingEnabled() {
-    return SnapshotProfilingConfiguration.SUPPLIER.get().isEnabled();
-  }
 }

profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SnapshotProfilingSupervisor.java

@@ -0,0 +1,116 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.profiler.snapshot;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.splunk.opentelemetry.profiler.OtelLoggerFactory;
+import com.splunk.opentelemetry.profiler.util.OptionalConfigurableSupplier;
+import io.opentelemetry.sdk.autoconfigure.AutoConfigureUtil;
+import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
+import java.util.logging.Logger;
+
+public class SnapshotProfilingSupervisor {
+  public static final OptionalConfigurableSupplier<SnapshotProfilingSupervisor> SUPPLIER =
+      new OptionalConfigurableSupplier<>();
+  private static final Logger logger =
+      Logger.getLogger(SnapshotProfilingSupervisor.class.getName());
+
+  private final OptionalConfigurableSupplier<SnapshotProfilingConfiguration> configurationSupplier;
+  private final ConfigurableSupplier<SpanTracker> spanTrackerSupplier;
+  private final OptionalConfigurableSupplier<TraceThreadChangeDetector>
+      traceThreadChangeDetectorSupplier;
+  private final AutoConfiguredOpenTelemetrySdk sdk;
+  private final OtelLoggerFactory otelLoggerFactory;
+  private boolean running;
+
+  @VisibleForTesting
+  SnapshotProfilingSupervisor(
+      OptionalConfigurableSupplier<SnapshotProfilingConfiguration> configurationSupplier,
+      ConfigurableSupplier<SpanTracker> spanTrackerSupplier,
+      OptionalConfigurableSupplier<TraceThreadChangeDetector> traceThreadChangeDetectorSupplier,
+      AutoConfiguredOpenTelemetrySdk sdk,
+      OtelLoggerFactory otelLoggerFactory) {
+    this.configurationSupplier = configurationSupplier;
+    this.spanTrackerSupplier = spanTrackerSupplier;
+    this.traceThreadChangeDetectorSupplier = traceThreadChangeDetectorSupplier;
+    this.sdk = sdk;
+    this.otelLoggerFactory = otelLoggerFactory;
+  }
+
+  public static SnapshotProfilingSupervisor initialize(AutoConfiguredOpenTelemetrySdk sdk) {
+    if (SUPPLIER.isConfigured()) {
+      throw new IllegalStateException("Snapshot profiling already initialized");
+    }
+
+    SnapshotProfilingSupervisor supervisor =
+        new SnapshotProfilingSupervisor(
+            SnapshotProfilingConfiguration.SUPPLIER,
+            SpanTracker.SUPPLIER,
+            TraceThreadChangeDetector.SUPPLIER,
+            sdk,
+            new OtelLoggerFactory());
+    SUPPLIER.configure(supervisor);
+
+    return supervisor;
+  }
+
+  public synchronized void startProfiling() {
+    if (running) {
+      return;
+    }
+
+    configurationSupplier.get().log();
+
+    // StagingArea.SUPPLIER
+    // StackTraceSampler.SUPPLIER
+    // StackTraceExporter.SUPPLIER = AsyncStackTraceExporter
+    StackTraceSamplerInitializer.setupStackTraceSampler(configurationSupplier.get());
+    StackTraceSamplerInitializer.setupStackTraceExporter(
+        configurationSupplier.get(), AutoConfigureUtil.getResource(sdk), otelLoggerFactory);
+
+    // SpanTracker.SUPPLIER
+    spanTrackerSupplier.get().setEnabled(true);
+    traceThreadChangeDetectorSupplier.get().setEnabled(true);
+
+    // SnapshotProfilingSdkCustomizer/SnapshotProfilingSpanProcessorComponentProvider ->
+    // SnapshotProfilingSpanProcessor
+    // SdkShutdownHookComponentProvider -> SdkShutdownHook
+
+    running = true;
+    logger.info("Snapshot profiling is active.");
+  }
+
+  public synchronized void stopProfiling() {
+    if (!running) {
+      return;
+    }
+
+    StackTraceSampler.SUPPLIER.get().close();
+    StackTraceSampler.SUPPLIER.reset();
+
+    StagingArea.SUPPLIER.get().close();
+    StagingArea.SUPPLIER.reset();
+
+    StackTraceExporter.SUPPLIER.get().close();
+    StackTraceExporter.SUPPLIER.reset();
+
+    SpanTracker.SUPPLIER.get().setEnabled(false);
+    TraceThreadChangeDetector.SUPPLIER.get().setEnabled(false);
+
+    running = false;
+  }
+}