[OpAMP] Snapshot profiling preparation for remote config

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

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

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 volatile 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,51 @@
+/*
+ * 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()) {
+      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,31 @@ 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();
+      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();
-  }
+      SnapshotProfilingSpanProcessor spanProcessor =
+          new SnapshotProfilingSpanProcessor(
+              registry, new TraceIdBasedSnapshotSelector(selectionProbability));
+      SnapshotProfilingSpanProcessor.SUPPLIER.configure(spanProcessor);
 
-  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(spanProcessor);
     };
   }
 
   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/SnapshotProfilingSpanProcessor.java

@@ -18,6 +18,7 @@
 
 import static com.splunk.opentelemetry.profiler.ProfilingSemanticAttributes.SNAPSHOT_PROFILING;
 
+import com.splunk.opentelemetry.profiler.util.OptionalConfigurableSupplier;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.common.CompletableResultCode;
@@ -29,9 +30,12 @@
  * Custom {@link SpanProcessor} implementation that will register traces for snapshot profiling<br>
  */
 public class SnapshotProfilingSpanProcessor implements SpanProcessor {
+  public static final OptionalConfigurableSupplier<SnapshotProfilingSpanProcessor> SUPPLIER =
+      new OptionalConfigurableSupplier<>();
   private final TraceRegistry registry;
   private final SnapshotSelector selector;
   private final OrphanedTraceCleaner orphanedTraceCleaner;
+  private volatile boolean enabled;
 
   SnapshotProfilingSpanProcessor(TraceRegistry registry, SnapshotSelector selector) {
     this.registry = registry;
@@ -41,6 +45,9 @@ public class SnapshotProfilingSpanProcessor implements SpanProcessor {
 
   @Override
   public void onStart(Context context, ReadWriteSpan span) {
+    if (!isEnabled()) {
+      return;
+    }
     if (isEntry(span)) {
       SpanContext spanContext = span.getSpanContext();
       boolean selected = selector.select(spanContext);
@@ -70,6 +77,8 @@ public boolean isStartRequired() {
   @Override
   public void onEnd(ReadableSpan span) {
     if (isEntry(span)) {
+      // To prevent memory leaks do the cleanup even if this processor is disabled.
+      // Some spans may have been started when processor was enabled.
       registry.unregister(span.getSpanContext());
       orphanedTraceCleaner.unregister(span.getSpanContext());
     }
@@ -89,4 +98,12 @@ public CompletableResultCode shutdown() {
     orphanedTraceCleaner.close();
     return SpanProcessor.super.shutdown();
   }
+
+  public void setEnabled(boolean enabled) {
+    this.enabled = enabled;
+  }
+
+  public boolean isEnabled() {
+    return enabled;
+  }
 }

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

@@ -49,7 +49,11 @@ public SnapshotProfilingSpanProcessor create(
       DeclarativeConfigProperties declarativeConfigProperties) {
     double selectionProbability =
         SnapshotProfilingConfiguration.SUPPLIER.get().getSnapshotSelectionProbability();
-    return new SnapshotProfilingSpanProcessor(
-        traceRegistry, new TraceIdBasedSnapshotSelector(selectionProbability));
+
+    SnapshotProfilingSpanProcessor spanProcessor =
+        new SnapshotProfilingSpanProcessor(
+            traceRegistry, new TraceIdBasedSnapshotSelector(selectionProbability));
+    SnapshotProfilingSpanProcessor.SUPPLIER.configure(spanProcessor);
+    return spanProcessor;
   }
 }