Stop Snapshot Profiler On OpenTelemetry SDK Shutdown (#2262)

* Define StackTraceSampler as Closeable.

* Stop all running samplers and prevent new trace samplers from being started when ScheduledExecutorStackTraceSampler is closed.

* Define StagingArea as Closeable and have AccumulatingStagingArea prevent new StackTraces from being staged when closed as well as export every staged StackTrace.

* Define StackTraceExporter as Closeable and have AsyncStackTraceExporter prevent new StackTraces from being exported when closed.

* Add the StackTraceSamplerProvider.

* Referece the StackTraceSampler via the StackTraceSamplerProvider in SnapshotProfilingSpanProcessor.

* Add the StagingAreaProvider.

* Referece the StagingArea via the StagingAreaProvider in AccumulatingStagingArea.

* Add the Closer class.

* Add the CloserProvider.

* Add the SpanProcessor to hook into the OpenTelemetry SDK shutdown process.

* Add test to verify that snapshot profiling is stopped when the OpenTelemetry SDK is shutdown.

* Remove unused classes.

* Add the generic ConfigurableSupplier class.

* Add constant ConfigurableSupplier instances available on SpanTracker, StackTraceSampler, StagingArea, and StackTraceExporter.

* Rename tests.

* Apply spotless code formatting.

* Initialize a ConfigurableSupplier value to the defaultValue.

* Mark the 'value' in ConfigurableSupplier volatile.

* Maintain a 'closed' field in AsyncStackTraceExporter and avoid submitting new stack trace export jobs when closed is true.

* Wait for schedulers to fully shutdown in ScheduledExecutorStackTraceSampler and AsyncStackTraceExporter before returning from the 'close' method.

* Wait for shutdown prior to emptying staging area.

* Apply spotless code formatting.

* Apply spotless code formatting.

* fix merge

* Update tests to compare staged stack traces after ScheduledExecutorStackTraceSampler has been stopped.

---------

Co-authored-by: Lauri Tulmin <[email protected]>

Author: tduncan

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

@@ -25,13 +25,18 @@
 class AccumulatingStagingArea implements StagingArea {
   private final ConcurrentMap<String, Queue<StackTrace>> stackTraces = new ConcurrentHashMap<>();
   private final Supplier<StackTraceExporter> exporter;
+  private volatile boolean closed = false;
 
   AccumulatingStagingArea(Supplier<StackTraceExporter> exporter) {
     this.exporter = exporter;
   }
 
   @Override
   public void stage(String traceId, StackTrace stackTrace) {
+    if (closed) {
+      return;
+    }
+
     stackTraces.compute(
         traceId,
         (id, stackTraces) -> {
@@ -50,4 +55,11 @@ public void empty(String traceId) {
       exporter.get().export(stackTraces);
     }
   }
+
+  @Override
+  public void close() {
+    closed = true;
+    stackTraces.values().forEach(exporter.get()::export);
+    stackTraces.clear();
+  }
 }

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

@@ -23,6 +23,7 @@
 import io.opentelemetry.api.logs.Logger;
 import java.util.Collection;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 
 class AsyncStackTraceExporter implements StackTraceExporter {
@@ -33,6 +34,7 @@ class AsyncStackTraceExporter implements StackTraceExporter {
       HelpfulExecutors.newSingleThreadExecutor("async-stack-trace-exporter");
   private final Logger otelLogger;
   private final int maxDepth;
+  private volatile boolean closed = false;
 
   AsyncStackTraceExporter(Logger logger, int maxDepth) {
     this.otelLogger = logger;
@@ -41,9 +43,27 @@ class AsyncStackTraceExporter implements StackTraceExporter {
 
   @Override
   public void export(Collection<StackTrace> stackTraces) {
+    if (closed) {
+      return;
+    }
     executor.submit(pprofExporter(otelLogger, stackTraces));
   }
 
+  @Override
+  public void close() {
+    closed = true;
+
+    try {
+      executor.shutdown();
+      if (!executor.awaitTermination(1, TimeUnit.SECONDS)) {
+        executor.shutdownNow();
+      }
+    } catch (InterruptedException e) {
+      executor.shutdownNow();
+      Thread.currentThread().interrupt();
+    }
+  }
+
   private Runnable pprofExporter(Logger otelLogger, Collection<StackTrace> stackTraces) {
     return () -> {
       try {

profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/SpanTrackerProvider.java renamed to profiler/src/main/java/com/splunk/opentelemetry/profiler/snapshot/ConfigurableSupplier.java

@@ -19,19 +19,25 @@
 import java.util.Objects;
 import java.util.function.Supplier;
 
-class SpanTrackerProvider implements Supplier<SpanTracker> {
-  public static final SpanTrackerProvider INSTANCE = new SpanTrackerProvider();
+class ConfigurableSupplier<T> implements Supplier<T> {
+  private final T defaultValue;
+  private volatile T value;
 
-  private SpanTracker tracker = SpanTracker.NOOP;
+  ConfigurableSupplier(T defaultValue) {
+    this.defaultValue = Objects.requireNonNull(defaultValue);
+    this.value = defaultValue;
+  }
 
   @Override
-  public SpanTracker get() {
-    return tracker;
+  public T get() {
+    return value;
   }
 
-  void configure(SpanTracker tracker) {
-    this.tracker = Objects.requireNonNull(tracker);
+  void configure(T value) {
+    this.value = Objects.requireNonNull(value);
   }
 
-  private SpanTrackerProvider() {}
+  void reset() {
+    this.value = defaultValue;
+  }
 }

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

@@ -39,7 +39,7 @@ public void activate(TraceRegistry registry) {
     contextStorageWrappingFunction.accept(
         contextStorage -> {
           ActiveSpanTracker tracker = new ActiveSpanTracker(contextStorage, registry);
-          SpanTrackerProvider.INSTANCE.configure(tracker);
+          SpanTracker.SUPPLIER.configure(tracker);
           return tracker;
         });
   }

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

@@ -38,19 +38,26 @@ class ScheduledExecutorStackTraceSampler implements StackTraceSampler {
 
   private final ConcurrentMap<String, ThreadSampler> samplers = new ConcurrentHashMap<>();
   private final ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
-  private final StagingArea stagingArea;
+  private final Supplier<StagingArea> stagingArea;
   private final Supplier<SpanTracker> spanTracker;
   private final Duration samplingPeriod;
+  private volatile boolean closed = false;
 
   ScheduledExecutorStackTraceSampler(
-      StagingArea stagingArea, Supplier<SpanTracker> spanTracker, Duration samplingPeriod) {
+      Supplier<StagingArea> stagingArea,
+      Supplier<SpanTracker> spanTracker,
+      Duration samplingPeriod) {
     this.stagingArea = stagingArea;
     this.spanTracker = spanTracker;
     this.samplingPeriod = samplingPeriod;
   }
 
   @Override
   public void start(SpanContext spanContext) {
+    if (closed) {
+      return;
+    }
+
     samplers.computeIfAbsent(
         spanContext.getTraceId(), id -> new ThreadSampler(spanContext, samplingPeriod));
   }
@@ -62,13 +69,34 @@ public void stop(SpanContext spanContext) {
         (traceId, sampler) -> {
           if (spanContext.equals(sampler.getSpanContext())) {
             sampler.shutdown();
-            stagingArea.empty(spanContext.getTraceId());
+            waitForShutdown(sampler);
+            stagingArea.get().empty(spanContext.getTraceId());
             return null;
           }
           return sampler;
         });
   }
 
+  @Override
+  public void close() {
+    closed = true;
+
+    samplers.values().forEach(ThreadSampler::shutdown);
+    samplers.values().forEach(this::waitForShutdown);
+    samplers.clear();
+  }
+
+  private void waitForShutdown(ThreadSampler sampler) {
+    try {
+      if (!sampler.awaitTermination(samplingPeriod.multipliedBy(2))) {
+        sampler.shutdownNow();
+      }
+    } catch (InterruptedException e) {
+      sampler.shutdownNow();
+      Thread.currentThread().interrupt();
+    }
+  }
+
   private class ThreadSampler {
     private final ScheduledExecutorService scheduler;
     private final SpanContext spanContext;
@@ -91,6 +119,14 @@ void shutdown() {
       gatherer.run();
     }
 
+    void shutdownNow() {
+      scheduler.shutdownNow();
+    }
+
+    boolean awaitTermination(Duration timeout) throws InterruptedException {
+      return scheduler.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS);
+    }
+
     SpanContext getSpanContext() {
       return spanContext;
     }
@@ -117,7 +153,7 @@ public void run() {
         String spanId = retrieveActiveSpan(thread).getSpanId();
         StackTrace stackTrace =
             StackTrace.from(Instant.now(), samplingPeriod, threadInfo, traceId, spanId);
-        stagingArea.stage(traceId, stackTrace);
+        stagingArea.get().stage(traceId, stackTrace);
       } catch (Exception e) {
         logger.log(Level.SEVERE, e, samplerErrorMessage(traceId, thread.getId()));
       } finally {

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

@@ -0,0 +1,62 @@
+/*
+ * 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 io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+import java.io.Closeable;
+import java.util.ArrayList;
+import java.util.List;
+
+class SdkShutdownHook implements SpanProcessor {
+  @Override
+  public CompletableResultCode shutdown() {
+    List<CompletableResultCode> results = new ArrayList<>();
+    results.add(close(StackTraceSampler.SUPPLIER.get()));
+    results.add(close(StagingArea.SUPPLIER.get()));
+    results.add(close(StackTraceExporter.SUPPLIER.get()));
+    return CompletableResultCode.ofAll(results);
+  }
+
+  private CompletableResultCode close(Closeable closeable) {
+    try {
+      closeable.close();
+      return CompletableResultCode.ofSuccess();
+    } catch (Exception e) {
+      return CompletableResultCode.ofExceptionalFailure(e);
+    }
+  }
+
+  @Override
+  public void onStart(Context parentContext, ReadWriteSpan span) {}
+
+  @Override
+  public boolean isStartRequired() {
+    return false;
+  }
+
+  @Override
+  public void onEnd(ReadableSpan span) {}
+
+  @Override
+  public boolean isEndRequired() {
+    return false;
+  }
+}

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

@@ -47,10 +47,9 @@ public SnapshotProfilingSdkCustomizer() {
   private static Function<ConfigProperties, StackTraceSampler> stackTraceSamplerProvider() {
     return properties -> {
       Duration samplingPeriod = Configuration.getSnapshotProfilerSamplingInterval(properties);
-      return new ScheduledExecutorStackTraceSampler(
-          new AccumulatingStagingArea(StackTraceExporterProvider.INSTANCE),
-          SpanTrackerProvider.INSTANCE,
-          samplingPeriod);
+      ConfigurableSupplier<StagingArea> supplier = StagingArea.SUPPLIER;
+      supplier.configure(new AccumulatingStagingArea(StackTraceExporter.SUPPLIER));
+      return new ScheduledExecutorStackTraceSampler(supplier, SpanTracker.SUPPLIER, samplingPeriod);
     };
   }
 
@@ -74,15 +73,28 @@ public void customize(AutoConfigurationCustomizer autoConfigurationCustomizer) {
     autoConfigurationCustomizer
         .addPropertiesCustomizer(autoConfigureSnapshotVolumePropagator())
         .addTracerProviderCustomizer(snapshotProfilingSpanProcessor(registry))
-        .addPropertiesCustomizer(startTrackingActiveSpans(registry));
+        .addPropertiesCustomizer(startTrackingActiveSpans(registry))
+        .addTracerProviderCustomizer(addShutdownHook());
+  }
+
+  private BiFunction<SdkTracerProviderBuilder, ConfigProperties, SdkTracerProviderBuilder>
+      addShutdownHook() {
+    return (builder, properties) -> {
+      if (snapshotProfilingEnabled(properties)) {
+        builder.addSpanProcessor(new SdkShutdownHook());
+      }
+      return builder;
+    };
   }
 
   private BiFunction<SdkTracerProviderBuilder, ConfigProperties, SdkTracerProviderBuilder>
       snapshotProfilingSpanProcessor(TraceRegistry registry) {
     return (builder, properties) -> {
       if (snapshotProfilingEnabled(properties)) {
         StackTraceSampler sampler = samplerProvider.apply(properties);
-        return builder.addSpanProcessor(new SnapshotProfilingSpanProcessor(registry, sampler));
+        ConfigurableSupplier<StackTraceSampler> supplier = StackTraceSampler.SUPPLIER;
+        supplier.configure(sampler);
+        return builder.addSpanProcessor(new SnapshotProfilingSpanProcessor(registry, supplier));
       }
       return builder;
     };

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

@@ -23,6 +23,7 @@
 import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SpanProcessor;
+import java.util.function.Supplier;
 
 /**
  * Custom {@link SpanProcessor} implementation that will 1. register traces for snapshot profiling
@@ -37,9 +38,9 @@
  */
 public class SnapshotProfilingSpanProcessor implements SpanProcessor {
   private final TraceRegistry registry;
-  private final StackTraceSampler sampler;
+  private final Supplier<StackTraceSampler> sampler;
 
-  SnapshotProfilingSpanProcessor(TraceRegistry registry, StackTraceSampler sampler) {
+  SnapshotProfilingSpanProcessor(TraceRegistry registry, Supplier<StackTraceSampler> sampler) {
     this.registry = registry;
     this.sampler = sampler;
   }
@@ -54,7 +55,7 @@ public void onStart(Context context, ReadWriteSpan span) {
     }
 
     if (isEntry(span) && registry.isRegistered(span.getSpanContext())) {
-      sampler.start(span.getSpanContext());
+      sampler.get().start(span.getSpanContext());
       span.setAttribute(SNAPSHOT_PROFILING, true);
     }
   }
@@ -75,7 +76,7 @@ public boolean isStartRequired() {
   public void onEnd(ReadableSpan span) {
     if (isEntry(span)) {
       registry.unregister(span.getSpanContext());
-      sampler.stop(span.getSpanContext());
+      sampler.get().stop(span.getSpanContext());
     }
   }
 

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

@@ -21,6 +21,7 @@
 
 interface SpanTracker {
   SpanTracker NOOP = thread -> Optional.empty();
+  ConfigurableSupplier<SpanTracker> SUPPLIER = new ConfigurableSupplier<>(SpanTracker.NOOP);
 
   Optional<SpanContext> getActiveSpan(Thread thread);
 }

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

@@ -16,11 +16,15 @@
 
 package com.splunk.opentelemetry.profiler.snapshot;
 
+import java.io.Closeable;
 import java.util.Collection;
 
 /** Works in concert with the {@link StagingArea} to export a batch of {@link StackTrace}s */
-interface StackTraceExporter {
+interface StackTraceExporter extends Closeable {
   StackTraceExporter NOOP = stackTraces -> {};
+  ConfigurableSupplier<StackTraceExporter> SUPPLIER = new ConfigurableSupplier<>(NOOP);
 
   void export(Collection<StackTrace> stackTraces);
+
+  default void close() {}
 }