Collect Callstacks For Threads Processing Traces (#2207)

* Add the StackTraceSampler interface.

* Add a ScheduledExecutor implementation of the StackTraceSampler interface.

* Start and stop profiling when entry spans start and end respectively.

* Add a noop StagingArea implementation and test builder for the SDK customizer to insulate individual tests from additional fields being added to the customizer.

* Use a concurrent map in ScheduledExecutorStackTraceSampler.

* Apply spotless code formatting.

* Convert SAMPLING_PERIOD constant to Duration.

* Remove debug logging from ScheduledExecutorStackTraceSampler.

* InMemoryStagingArea doesn't need to be a JUnit extension.

* Remove unused JUnit extension method from ObservableStackTraceSampler.

* Rename LOGGER to lowercase.

* Log explanatory message and stacktrace if something goes wrong collecting callstacks from profilied threads.

* Add explanatory comment to SnapshotProfilingSpanProcessor.

* Expand the StackTraceSampler interface to accept a SpanContext to eventually track traces, not only thread IDs.

* Associate thread profilers with trace IDs to account for different threads managing a spans lifecycle (e.g. async web frameworks) and avoid a resource leak and runaway profiling thread.

* Add test verifying that ending a span from a different thread will stop trace profiling.

* Only allow a single thread per trace to be profiled at a time.

* Rename the 'Profilng' class to 'Snapshotting'.

* Attempt to make the explanation about the multithreaded trace limitation more clear.

Author: tduncan

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

@@ -0,0 +1,25 @@
+/*
+ * 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;
+
+class NoopStagingArea implements StagingArea {
+  @Override
+  public void stage(String traceId, StackTrace stackTrace) {}
+
+  @Override
+  public void empty(String traceId) {}
+}

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

@@ -0,0 +1,111 @@
+/*
+ * 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 io.opentelemetry.api.trace.SpanContext;
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadInfo;
+import java.lang.management.ThreadMXBean;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+class ScheduledExecutorStackTraceSampler implements StackTraceSampler {
+  private static final Logger logger =
+      Logger.getLogger(ScheduledExecutorStackTraceSampler.class.getName());
+  private static final int SCHEDULER_INITIAL_DELAY = 0;
+  private static final Duration SCHEDULER_PERIOD = Duration.ofMillis(20);
+  private static final int MAX_ENTRY_DEPTH = 200;
+
+  private final ConcurrentMap<String, ScheduledExecutorService> samplers =
+      new ConcurrentHashMap<>();
+  private final ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
+  private final StagingArea stagingArea;
+  private final Duration samplingPeriod;
+
+  ScheduledExecutorStackTraceSampler(StagingArea stagingArea) {
+    this(stagingArea, SCHEDULER_PERIOD);
+  }
+
+  @VisibleForTesting
+  ScheduledExecutorStackTraceSampler(StagingArea stagingArea, Duration samplingPeriod) {
+    this.stagingArea = stagingArea;
+    this.samplingPeriod = samplingPeriod;
+  }
+
+  @Override
+  public void start(SpanContext spanContext) {
+    if (samplers.containsKey(spanContext.getTraceId())) {
+      return;
+    }
+
+    ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
+    samplers.put(spanContext.getTraceId(), scheduler);
+    scheduler.scheduleAtFixedRate(
+        new StackTraceGatherer(spanContext.getTraceId(), Thread.currentThread().getId()),
+        SCHEDULER_INITIAL_DELAY,
+        samplingPeriod.toMillis(),
+        TimeUnit.MILLISECONDS);
+  }
+
+  @Override
+  public void stop(SpanContext spanContext) {
+    ScheduledExecutorService scheduler = samplers.remove(spanContext.getTraceId());
+    if (scheduler != null) {
+      scheduler.shutdown();
+    }
+    stagingArea.empty(spanContext.getTraceId());
+  }
+
+  class StackTraceGatherer implements Runnable {
+    private final String traceId;
+    private final long threadId;
+
+    StackTraceGatherer(String traceId, long threadId) {
+      this.traceId = traceId;
+      this.threadId = threadId;
+    }
+
+    @Override
+    public void run() {
+      Instant now = Instant.now();
+      try {
+        ThreadInfo threadInfo = threadMXBean.getThreadInfo(threadId, MAX_ENTRY_DEPTH);
+        StackTrace stackTrace = StackTrace.from(now, threadInfo);
+        stagingArea.stage(traceId, stackTrace);
+      } catch (Exception e) {
+        logger.log(Level.SEVERE, e, samplerErrorMessage(traceId, threadId));
+      }
+    }
+
+    private Supplier<String> samplerErrorMessage(String traceId, long threadId) {
+      return () ->
+          "Exception thrown attempting to stage callstacks for trace ID ' "
+              + traceId
+              + "' on profiled thread "
+              + threadId;
+    }
+  }
+}

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

@@ -29,14 +29,16 @@
 @AutoService(AutoConfigurationCustomizerProvider.class)
 public class SnapshotProfilingSdkCustomizer implements AutoConfigurationCustomizerProvider {
   private final TraceRegistry registry;
+  private final StackTraceSampler sampler;
 
   public SnapshotProfilingSdkCustomizer() {
-    this(new TraceRegistry());
+    this(new TraceRegistry(), new ScheduledExecutorStackTraceSampler(new NoopStagingArea()));
   }
 
   @VisibleForTesting
-  SnapshotProfilingSdkCustomizer(TraceRegistry registry) {
+  SnapshotProfilingSdkCustomizer(TraceRegistry registry, StackTraceSampler sampler) {
     this.registry = registry;
+    this.sampler = sampler;
   }
 
   @Override
@@ -49,7 +51,7 @@ public void customize(AutoConfigurationCustomizer autoConfigurationCustomizer) {
       snapshotProfilingSpanProcessor(TraceRegistry registry) {
     return (builder, properties) -> {
       if (snapshotProfilingEnabled(properties)) {
-        return builder.addSpanProcessor(new SnapshotProfilingSpanProcessor(registry));
+        return builder.addSpanProcessor(new SnapshotProfilingSpanProcessor(registry, sampler));
       }
       return builder;
     };

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

@@ -24,11 +24,24 @@
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SpanProcessor;
 
+/**
+ * Custom {@link SpanProcessor} implementation that will 1. register traces for snapshot profiling
+ * and 2. activate profiling for the thread processing the trace. <br>
+ * <br>
+ * <b>Implementation Note</b><br>
+ * The current snapshot profiling extension supports profiling one thread per trace at a time. If
+ * the trace spans multiple threads -- whether because a service is invoked multiple time
+ * concurrently by an upstream caller or work is delegated to background threads -- only a single
+ * thread associated with that trace (specifically the thread associated with the "Entry" span) will
+ * be profiled at a time.
+ */
 public class SnapshotProfilingSpanProcessor implements SpanProcessor {
   private final TraceRegistry registry;
+  private final StackTraceSampler sampler;
 
-  SnapshotProfilingSpanProcessor(TraceRegistry registry) {
+  SnapshotProfilingSpanProcessor(TraceRegistry registry, StackTraceSampler sampler) {
     this.registry = registry;
+    this.sampler = sampler;
   }
 
   @Override
@@ -41,6 +54,7 @@ public void onStart(Context context, ReadWriteSpan span) {
     }
 
     if (isEntry(span) && registry.isRegistered(span.getSpanContext())) {
+      sampler.start(span.getSpanContext());
       span.setAttribute(SNAPSHOT_PROFILING, true);
     }
   }
@@ -61,6 +75,7 @@ public boolean isStartRequired() {
   public void onEnd(ReadableSpan span) {
     if (isEntry(span)) {
       registry.unregister(span.getSpanContext());
+      sampler.stop(span.getSpanContext());
     }
   }
 

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

@@ -0,0 +1,56 @@
+/*
+ * 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 java.lang.management.ThreadInfo;
+import java.time.Instant;
+
+class StackTrace {
+  static StackTrace from(Instant timestamp, ThreadInfo thread) {
+    return new StackTrace(
+        timestamp, thread.getThreadId(), thread.getThreadName(), thread.getStackTrace());
+  }
+
+  private final Instant timestamp;
+  private final long threadId;
+  private final String threadName;
+  private final StackTraceElement[] stackFrames;
+
+  private StackTrace(
+      Instant timestamp, long threadId, String threadName, StackTraceElement[] stackFrames) {
+    this.timestamp = timestamp;
+    this.threadId = threadId;
+    this.threadName = threadName;
+    this.stackFrames = stackFrames;
+  }
+
+  Instant getTimestamp() {
+    return timestamp;
+  }
+
+  long getThreadId() {
+    return threadId;
+  }
+
+  String getThreadName() {
+    return threadName;
+  }
+
+  StackTraceElement[] getStackFrames() {
+    return stackFrames;
+  }
+}

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

@@ -0,0 +1,25 @@
+/*
+ * 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.api.trace.SpanContext;
+
+interface StackTraceSampler {
+  void start(SpanContext spanContext);
+
+  void stop(SpanContext spanContext);
+}

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

@@ -0,0 +1,27 @@
+/*
+ * 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;
+
+/**
+ * Acts as a location to stockpile gathered {@link StackTrace}s segmented by thread ID for bulk
+ * exportation at some later point in time.
+ */
+interface StagingArea {
+  void stage(String traceId, StackTrace stackTrace);
+
+  void empty(String traceId);
+}

profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/DistributedProfilingSignalTest.java

@@ -31,7 +31,7 @@
 class DistributedProfilingSignalTest {
   private final RecordingTraceRegistry downstreamRegistry = new RecordingTraceRegistry();
   private final SnapshotProfilingSdkCustomizer downstreamCustomizer =
-      new SnapshotProfilingSdkCustomizer(downstreamRegistry);
+      Snapshotting.customizer().with(downstreamRegistry).build();
 
   @RegisterExtension
   public final OpenTelemetrySdkExtension downstreamSdk =
@@ -57,7 +57,7 @@ class DistributedProfilingSignalTest {
 
   private final RecordingTraceRegistry upstreamRegistry = new RecordingTraceRegistry();
   private final SnapshotProfilingSdkCustomizer upstreamCustomizer =
-      new SnapshotProfilingSdkCustomizer(upstreamRegistry);
+      Snapshotting.customizer().with(upstreamRegistry).build();
 
   @RegisterExtension
   public final OpenTelemetrySdkExtension upstreamSdk =

profiler/src/test/java/com/splunk/opentelemetry/profiler/snapshot/InMemoryStagingArea.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 java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+
+/**
+ * In memory implementation of the {@link StagingArea} interface that allows for direct access to
+ * the stockpiled {@link StackTrace}s. Intended for testing use only.
+ */
+class InMemoryStagingArea implements StagingArea {
+  private final Map<String, List<StackTrace>> stackTraces = new ConcurrentHashMap<>();
+
+  @Override
+  public void stage(String traceId, StackTrace stackTrace) {
+    stackTraces.compute(
+        traceId,
+        (id, stackTraces) -> {
+          if (stackTraces == null) {
+            stackTraces = new ArrayList<>();
+          }
+          stackTraces.add(stackTrace);
+          return stackTraces;
+        });
+  }
+
+  @Override
+  public void empty(String traceId) {
+    stackTraces.remove(traceId);
+  }
+
+  public List<StackTrace> allStackTraces() {
+    return stackTraces.values().stream().flatMap(List::stream).collect(Collectors.toList());
+  }
+}