Port span correlation (#2251)

* Add a custom ContextStorage to track the currently active span for traces selected for snapshotting.

* Configure active span tracking in SnapshotProfilingSdkCustomizer.

* Add the SpanTracker interface.

* Include the active span details on collected stack traces when available.

* Activate span stracking in SnapshotProfilingSdkCustomizer.

* Remove missed debugging statements in test.

* Remove unnecessary log record span context setting from InMemoryOtelLogger.

* Include the currently active span ID on collected stack traces.

* Introduce the SpanTrackingActivator interface to add a mechanism for tests to reliably inject the snapshot profiling active span tracker into the OpenTelemetry SDK runtime.

* Introduce a SpanTrackerProvider to remove the need for ActiveSpanTracker to be implemented as a Singleton.

* Remove unnecessary 'null' check in ActiveSpanTracker.

* Add tests for InterceptingContextStorageSpanTrackingActivator.

* Add tests for SpanTrackerProvider.

* Consider whether a span is sampled, not recording, when deciding whether or not to track it.

* Remove unnecessary fake ContextStorage in ActiveSpanTrackerTest.

* Add a builder subclass for FakeSpan in ActiveSpanTrackerTest.

* Move SpanContext builder into Snapshotting test helper class.

* Remove helper method for generating random trace IDs from the Snapshotting class in favor of direct use of OpenTelemetry's IdGenerator.

* Restore unnecessary changes to InMemoryOtelLogger.

* Remove FakeSpan builder in ActiveSpanTrackerTest in favor of the SpanContextBuilder obtained from Snapshotting.

* Add easier way to build SpanContext instances using trace IDs from other spans.

* Do not use traceId from the active span tracker when creating StackTrace instances so to always have a trace ID associated even if span tracking is offline.

* Remove the possibility that the SpanTracker in SpanTrackerProvider can ever be null.

* Add ActiveSpanTrackerTest asserting that the active span is restored to a parent span upon scope closing.

* Renaming class.

* Expect span ID on exported logs.

* Apply spotless code formatting.

* Track active spans per thread rather than trace ID to account for traces that span multiple threads.

* Replace FakeSpan in ActiveSpanTrackerTest with Span.wrap.

* Add ActiveSpanTracker test to ensure that invalid spans are not tracked.

* Track active spans by Thread rather than thread ID and use a map of weak reference keys in ActiveSpanTracker to ensure than map entries are automatically removed if there are a thread terminates without the OpenTelemetry scope being closed properly.

* Add a CountDownLatch to allow tests to explicitly release background tasks to ensure those tasks execute on multiple threads.

* Apply spotless code formatting.

Author: tduncan

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

@@ -0,0 +1,77 @@
+/*
+ * 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.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextStorage;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.internal.cache.Cache;
+import java.util.Optional;
+import javax.annotation.Nullable;
+
+class ActiveSpanTracker implements ContextStorage, SpanTracker {
+  private final Cache<Thread, SpanContext> cache = Cache.weak();
+
+  private final ContextStorage delegate;
+  private final TraceRegistry registry;
+
+  ActiveSpanTracker(ContextStorage delegate, TraceRegistry registry) {
+    this.delegate = delegate;
+    this.registry = registry;
+  }
+
+  @Override
+  public Scope attach(Context toAttach) {
+    Scope scope = delegate.attach(toAttach);
+    SpanContext newSpanContext = Span.fromContext(toAttach).getSpanContext();
+    if (doNotTrack(newSpanContext)) {
+      return scope;
+    }
+
+    Thread thread = Thread.currentThread();
+    SpanContext oldSpanContext = cache.get(thread);
+    if (oldSpanContext == newSpanContext) {
+      return scope;
+    }
+
+    cache.put(thread, newSpanContext);
+    return () -> {
+      if (oldSpanContext != null) {
+        cache.put(thread, oldSpanContext);
+      } else {
+        cache.remove(thread);
+      }
+      scope.close();
+    };
+  }
+
+  private boolean doNotTrack(SpanContext spanContext) {
+    return !spanContext.isSampled() || !registry.isRegistered(spanContext);
+  }
+
+  @Nullable
+  @Override
+  public Context current() {
+    return delegate.current();
+  }
+
+  public Optional<SpanContext> getActiveSpan(Thread thread) {
+    return Optional.ofNullable(cache.get(thread));
+  }
+}

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

@@ -65,7 +65,7 @@ private Runnable pprofExporter(Logger otelLogger, List<StackTrace> stackTraces)
               stackTrace.getStackFrames(),
               stackTrace.getTimestamp(),
               stackTrace.getTraceId(),
-              null);
+              stackTrace.getSpanId());
         }
         cpuEventExporter.flush();
       } catch (Exception e) {

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

@@ -0,0 +1,46 @@
+/*
+ * 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.context.ContextStorage;
+import java.util.function.Consumer;
+import java.util.function.UnaryOperator;
+
+class InterceptingContextStorageSpanTrackingActivator implements SpanTrackingActivator {
+  private final Consumer<UnaryOperator<ContextStorage>> contextStorageWrappingFunction;
+
+  InterceptingContextStorageSpanTrackingActivator() {
+    this(ContextStorage::addWrapper);
+  }
+
+  @VisibleForTesting
+  InterceptingContextStorageSpanTrackingActivator(
+      Consumer<UnaryOperator<ContextStorage>> contextStorageWrappingFunction) {
+    this.contextStorageWrappingFunction = contextStorageWrappingFunction;
+  }
+
+  @Override
+  public void activate(TraceRegistry registry) {
+    contextStorageWrappingFunction.accept(
+        contextStorage -> {
+          ActiveSpanTracker tracker = new ActiveSpanTracker(contextStorage, registry);
+          SpanTrackerProvider.INSTANCE.configure(tracker);
+          return tracker;
+        });
+  }
+}

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

@@ -40,10 +40,13 @@ class ScheduledExecutorStackTraceSampler implements StackTraceSampler {
       new ConcurrentHashMap<>();
   private final ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
   private final StagingArea stagingArea;
+  private final Supplier<SpanTracker> spanTracker;
   private final Duration samplingPeriod;
 
-  ScheduledExecutorStackTraceSampler(StagingArea stagingArea, Duration samplingPeriod) {
+  ScheduledExecutorStackTraceSampler(
+      StagingArea stagingArea, Supplier<SpanTracker> spanTracker, Duration samplingPeriod) {
     this.stagingArea = stagingArea;
+    this.spanTracker = spanTracker;
     this.samplingPeriod = samplingPeriod;
   }
 
@@ -55,7 +58,7 @@ public void start(SpanContext spanContext) {
           ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
           scheduler.scheduleAtFixedRate(
               new StackTraceGatherer(
-                  samplingPeriod, spanContext.getTraceId(), Thread.currentThread().getId()),
+                  samplingPeriod, spanContext.getTraceId(), Thread.currentThread()),
               SCHEDULER_INITIAL_DELAY,
               samplingPeriod.toMillis(),
               TimeUnit.MILLISECONDS);
@@ -75,26 +78,32 @@ public void stop(SpanContext spanContext) {
   class StackTraceGatherer implements Runnable {
     private final Duration samplingPeriod;
     private final String traceId;
-    private final long threadId;
+    private final Thread thread;
 
-    StackTraceGatherer(Duration samplingPeriod, String traceId, long threadId) {
+    StackTraceGatherer(Duration samplingPeriod, String traceId, Thread thread) {
       this.samplingPeriod = samplingPeriod;
       this.traceId = traceId;
-      this.threadId = threadId;
+      this.thread = thread;
     }
 
     @Override
     public void run() {
       try {
         Instant now = Instant.now();
-        ThreadInfo threadInfo = threadMXBean.getThreadInfo(threadId, Integer.MAX_VALUE);
-        StackTrace stackTrace = StackTrace.from(now, samplingPeriod, traceId, threadInfo);
+        ThreadInfo threadInfo = threadMXBean.getThreadInfo(thread.getId(), Integer.MAX_VALUE);
+        SpanContext spanContext = retrieveActiveSpan(thread);
+        StackTrace stackTrace =
+            StackTrace.from(now, samplingPeriod, threadInfo, traceId, spanContext.getSpanId());
         stagingArea.stage(traceId, stackTrace);
       } catch (Exception e) {
-        logger.log(Level.SEVERE, e, samplerErrorMessage(traceId, threadId));
+        logger.log(Level.SEVERE, e, samplerErrorMessage(traceId, thread.getId()));
       }
     }
 
+    private SpanContext retrieveActiveSpan(Thread thread) {
+      return spanTracker.get().getActiveSpan(thread).orElse(SpanContext.getInvalid());
+    }
+
     private Supplier<String> samplerErrorMessage(String traceId, long threadId) {
       return () ->
           "Exception thrown attempting to stage callstacks for trace ID ' "

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

@@ -35,35 +35,46 @@
 public class SnapshotProfilingSdkCustomizer implements AutoConfigurationCustomizerProvider {
   private final TraceRegistry registry;
   private final Function<ConfigProperties, StackTraceSampler> samplerProvider;
+  private final SpanTrackingActivator spanTrackingActivator;
 
   public SnapshotProfilingSdkCustomizer() {
-    this(new TraceRegistry(), stackTraceSamplerProvider());
+    this(
+        new TraceRegistry(),
+        stackTraceSamplerProvider(),
+        new InterceptingContextStorageSpanTrackingActivator());
   }
 
   private static Function<ConfigProperties, StackTraceSampler> stackTraceSamplerProvider() {
     return properties -> {
       Duration samplingPeriod = Configuration.getSnapshotProfilerSamplingInterval(properties);
       return new ScheduledExecutorStackTraceSampler(
-          new AccumulatingStagingArea(StackTraceExporterProvider.INSTANCE), samplingPeriod);
+          new AccumulatingStagingArea(StackTraceExporterProvider.INSTANCE),
+          SpanTrackerProvider.INSTANCE,
+          samplingPeriod);
     };
   }
 
   @VisibleForTesting
-  SnapshotProfilingSdkCustomizer(TraceRegistry registry, StackTraceSampler samplerProvider) {
-    this(registry, properties -> samplerProvider);
+  SnapshotProfilingSdkCustomizer(
+      TraceRegistry registry, StackTraceSampler sampler, SpanTrackingActivator activator) {
+    this(registry, properties -> sampler, activator);
   }
 
   private SnapshotProfilingSdkCustomizer(
-      TraceRegistry registry, Function<ConfigProperties, StackTraceSampler> samplerProvider) {
+      TraceRegistry registry,
+      Function<ConfigProperties, StackTraceSampler> samplerProvider,
+      SpanTrackingActivator spanTrackingActivator) {
     this.registry = registry;
     this.samplerProvider = samplerProvider;
+    this.spanTrackingActivator = spanTrackingActivator;
   }
 
   @Override
   public void customize(AutoConfigurationCustomizer autoConfigurationCustomizer) {
     autoConfigurationCustomizer
         .addPropertiesCustomizer(autoConfigureSnapshotVolumePropagator())
-        .addTracerProviderCustomizer(snapshotProfilingSpanProcessor(registry));
+        .addTracerProviderCustomizer(snapshotProfilingSpanProcessor(registry))
+        .addPropertiesCustomizer(startTrackingActiveSpans(registry));
   }
 
   private BiFunction<SdkTracerProviderBuilder, ConfigProperties, SdkTracerProviderBuilder>
@@ -111,6 +122,16 @@ private boolean includeTraceContextPropagator(Set<String> configuredPropagators)
     return configuredPropagators.isEmpty();
   }
 
+  private Function<ConfigProperties, Map<String, String>> startTrackingActiveSpans(
+      TraceRegistry registry) {
+    return properties -> {
+      if (snapshotProfilingEnabled(properties)) {
+        spanTrackingActivator.activate(registry);
+      }
+      return Collections.emptyMap();
+    };
+  }
+
   private boolean snapshotProfilingEnabled(ConfigProperties properties) {
     return Configuration.isSnapshotProfilingEnabled(properties);
   }

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

@@ -0,0 +1,26 @@
+/*
+ * 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;
+import java.util.Optional;
+
+interface SpanTracker {
+  SpanTracker NOOP = thread -> Optional.empty();
+
+  Optional<SpanContext> getActiveSpan(Thread thread);
+}

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

@@ -0,0 +1,37 @@
+/*
+ * 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.Objects;
+import java.util.function.Supplier;
+
+class SpanTrackerProvider implements Supplier<SpanTracker> {
+  public static final SpanTrackerProvider INSTANCE = new SpanTrackerProvider();
+
+  private SpanTracker tracker = SpanTracker.NOOP;
+
+  @Override
+  public SpanTracker get() {
+    return tracker;
+  }
+
+  void configure(SpanTracker tracker) {
+    this.tracker = Objects.requireNonNull(tracker);
+  }
+
+  private SpanTrackerProvider() {}
+}

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

@@ -0,0 +1,21 @@
+/*
+ * 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;
+
+interface SpanTrackingActivator {
+  void activate(TraceRegistry registry);
+}