Capture span stack trace (#21)

custom/build.gradle.kts

@@ -13,4 +13,11 @@ dependencies {
   // needs to be added in order to allow access to AgentListener interface
   // this is currently required because autoconfigure is currently not exposed to the extension API.
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
+
+  // test dependencies
+  testImplementation("io.opentelemetry:opentelemetry-sdk")
+  testImplementation("io.opentelemetry.javaagent:opentelemetry-testing-common")
+  testImplementation("io.opentelemetry:opentelemetry-sdk-testing")
+  testImplementation("org.assertj:assertj-core:3.24.2") // TODO : remove version for assertj
+
 }

custom/src/main/java/co/elastic/otel/ElasticAttributes.java

@@ -21,11 +21,14 @@
 import io.opentelemetry.api.common.AttributeKey;
 
 public interface ElasticAttributes {
-  AttributeKey<Long> SELF_TIME_ATTRIBUTE = AttributeKey.longKey("elastic.span.self_time");
+  AttributeKey<Long> SELF_TIME = AttributeKey.longKey("elastic.span.self_time");
   AttributeKey<String> LOCAL_ROOT_ID = AttributeKey.stringKey("elastic.span.local_root.id");
   AttributeKey<String> LOCAL_ROOT_NAME = AttributeKey.stringKey("elastic.local_root.name");
   AttributeKey<String> LOCAL_ROOT_TYPE = AttributeKey.stringKey("elastic.local_root.type");
   AttributeKey<Boolean> IS_LOCAL_ROOT = AttributeKey.booleanKey("elastic.span.is_local_root");
-  AttributeKey<String> ELASTIC_SPAN_TYPE = AttributeKey.stringKey("elastic.span.type");
-  AttributeKey<String> ELASTIC_SPAN_SUBTYPE = AttributeKey.stringKey("elastic.span.subtype");
+  AttributeKey<String> SPAN_TYPE = AttributeKey.stringKey("elastic.span.type");
+  AttributeKey<String> SPAN_SUBTYPE = AttributeKey.stringKey("elastic.span.subtype");
+
+  // TODO : replace this with semantic conventions v1.24.0 equivalent once released
+  AttributeKey<String> SPAN_STACKTRACE = AttributeKey.stringKey("code.stacktrace");
 }

custom/src/main/java/co/elastic/otel/ElasticBreakdownMetrics.java

@@ -37,21 +37,21 @@
 
 public class ElasticBreakdownMetrics {
 
-  private final ConcurrentHashMap<SpanContext, SpanContextData> elasticSpanData;
+  private final ConcurrentHashMap<SpanContext, BreakdownData> elasticSpanData;
 
   private ElasticSpanExporter spanExporter;
 
   private LongCounter breakDownCounter;
 
   // sidecar object we store for every span
-  public static class SpanContextData {
+  private static class BreakdownData {
 
     private ReadableSpan localRoot;
 
     private final ChildDuration childDuration;
     private long selfTime;
 
-    public SpanContextData(ReadableSpan localRoot, long start) {
+    public BreakdownData(ReadableSpan localRoot, long start) {
       this.localRoot = localRoot;
       this.childDuration = new ChildDuration(start);
       this.selfTime = Long.MIN_VALUE;
@@ -108,7 +108,7 @@ public void onSpanStart(Context parentContext, ReadWriteSpan span) {
       // the span is a local root span
       localRootSpanContext = spanContext;
 
-      elasticSpanData.put(spanContext, new SpanContextData(span, spanStart));
+      elasticSpanData.put(spanContext, new BreakdownData(span, spanStart));
 
     } else {
       ReadableSpan parentSpan = getReadableSpanFromContext(parentContext);
@@ -118,7 +118,7 @@ public void onSpanStart(Context parentContext, ReadWriteSpan span) {
       ReadableSpan localRoot = lookupLocalRootSpan(parentSpan);
       localRootSpanContext = localRoot.getSpanContext();
       if (localRootSpanContext.isValid()) {
-        elasticSpanData.put(spanContext, new SpanContextData(localRoot, spanStart));
+        elasticSpanData.put(spanContext, new BreakdownData(localRoot, spanStart));
       }
 
       // update direct parent span child durations for self-time
@@ -157,11 +157,11 @@ public void onSpanEnd(ReadableSpan span) {
     SpanData spanData = span.toSpanData();
 
     // children duration for current span
-    SpanContextData spanContextData = elasticSpanData.get(spanContext);
+    BreakdownData spanContextData = elasticSpanData.get(spanContext);
     Objects.requireNonNull(spanContextData, "missing elastic span data");
 
     // update children duration for direct parent
-    SpanContextData parentSpanContextData = elasticSpanData.get(span.getParentSpanContext());
+    BreakdownData parentSpanContextData = elasticSpanData.get(span.getParentSpanContext());
 
     if (parentSpanContextData != null) { // parent might be already terminated
       parentSpanContextData.childDuration.endChild(spanData.getEndEpochNanos());
@@ -176,14 +176,15 @@ public void onSpanEnd(ReadableSpan span) {
             // put measured metric as span attribute to allow using an ingest pipeline to alter
             // storage
             // ingest pipelines do not have access to _source and thus can't read the metric as-is.
-            .put(ElasticAttributes.SELF_TIME_ATTRIBUTE, selfTime);
+            .put(ElasticAttributes.SELF_TIME, selfTime);
 
     // unfortunately here we get a read-only span that has already been ended, thus even a cast to
     // ReadWriteSpan
     // does not allow us from adding extra span attributes
     if (spanExporter != null) {
       spanContextData.setSelfTime(selfTime);
-      spanExporter.report(spanContext, spanContextData);
+      spanExporter.addAttribute(
+          spanContext, ElasticAttributes.SELF_TIME, spanContextData.getSelfTime());
     }
 
     breakDownCounter.add(selfTime, metricAttributes.build());
@@ -194,25 +195,25 @@ private static AttributesBuilder buildCounterAttributes(Attributes spanAttribute
     AttributesBuilder builder =
         Attributes.builder()
             // default to app/internal unless other span attributes
-            .put(ElasticAttributes.ELASTIC_SPAN_TYPE, "app")
-            .put(ElasticAttributes.ELASTIC_SPAN_SUBTYPE, "internal");
+            .put(ElasticAttributes.SPAN_TYPE, "app")
+            .put(ElasticAttributes.SPAN_SUBTYPE, "internal");
 
     spanAttributes.forEach(
         (k, v) -> {
           String key = k.getKey();
           if (AttributeType.STRING.equals(k.getType())) {
             int index = key.indexOf(".system");
             if (index > 0) {
-              builder.put(ElasticAttributes.ELASTIC_SPAN_TYPE, key.substring(0, index));
-              builder.put(ElasticAttributes.ELASTIC_SPAN_SUBTYPE, v.toString());
+              builder.put(ElasticAttributes.SPAN_TYPE, key.substring(0, index));
+              builder.put(ElasticAttributes.SPAN_SUBTYPE, v.toString());
             }
           }
         });
     return builder;
   }
 
   private ReadableSpan lookupLocalRootSpan(ReadableSpan span) {
-    SpanContextData spanContextData = elasticSpanData.get(span.getSpanContext());
+    BreakdownData spanContextData = elasticSpanData.get(span.getSpanContext());
     return spanContextData != null ? spanContextData.localRoot : (ReadableSpan) Span.getInvalid();
   }
 

custom/src/main/java/co/elastic/otel/ElasticExtension.java

@@ -79,6 +79,7 @@ public SpanExporter wrapSpanExporter(SpanExporter toWrap) {
     profiler.registerExporter(toWrap);
     spanExporter = new ElasticSpanExporter(toWrap);
     breakdownMetrics.registerSpanExporter(spanExporter);
+    spanProcessor.registerSpanExporter(spanExporter);
     return spanExporter;
   }
 

custom/src/main/java/co/elastic/otel/ElasticSpanExporter.java

@@ -18,7 +18,9 @@
  */
 package co.elastic.otel;
 
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.resources.Resource;
@@ -34,34 +36,35 @@ public class ElasticSpanExporter implements SpanExporter {
 
   private final SpanExporter delegate;
 
-  private ConcurrentHashMap<SpanContext, ElasticBreakdownMetrics.SpanContextData> storage;
+  private final ConcurrentHashMap<SpanContext, AttributesBuilder> attributes;
 
   public ElasticSpanExporter(SpanExporter delegate) {
     this.delegate = delegate;
-    this.storage = new ConcurrentHashMap<>();
+    this.attributes = new ConcurrentHashMap<>();
   }
 
   @Override
   public CompletableResultCode export(Collection<SpanData> spans) {
     // shortcut in the rare case where no filtering is required
-    if (storage.isEmpty()) {
+    if (attributes.isEmpty()) {
       return delegate.export(spans);
     }
 
     List<SpanData> toSend = new ArrayList<>(spans.size());
     for (SpanData span : spans) {
       SpanContext spanContext = span.getSpanContext();
-      ElasticBreakdownMetrics.SpanContextData data = storage.remove(spanContext);
-      if (data == null) {
+      AttributesBuilder extraAttributes = attributes.remove(spanContext);
+      if (extraAttributes == null) {
         toSend.add(span);
       } else {
+        Attributes newAttributes =
+            span.getAttributes().toBuilder().putAll(extraAttributes.build()).build();
+
         toSend.add(
             new DelegatingSpanData(span) {
               @Override
               public Attributes getAttributes() {
-                return span.getAttributes().toBuilder()
-                    .put(ElasticAttributes.SELF_TIME_ATTRIBUTE, data.getSelfTime())
-                    .build();
+                return newAttributes;
               }
 
               @Override
@@ -75,13 +78,21 @@ public Resource getResource() {
     return delegate.export(toSend);
   }
 
-  public void report(SpanContext spanContext, ElasticBreakdownMetrics.SpanContextData data) {
-    this.storage.put(spanContext, data);
+  public <T> void addAttribute(SpanContext spanContext, AttributeKey<T> key, T value) {
+    attributes.compute(
+        spanContext,
+        (k, builder) -> {
+          if (builder == null) {
+            builder = Attributes.builder();
+          }
+          builder.put(key, value);
+          return builder;
+        });
   }
 
   @Override
   public CompletableResultCode flush() {
-    storage.clear();
+    attributes.clear();
     return delegate.flush();
   }
 

custom/src/main/java/co/elastic/otel/ElasticSpanProcessor.java

@@ -23,11 +23,14 @@
 import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SpanProcessor;
+import java.io.PrintWriter;
+import java.io.StringWriter;
 
 public class ElasticSpanProcessor implements SpanProcessor {
 
   private final ElasticProfiler profiler;
   private final ElasticBreakdownMetrics breakdownMetrics;
+  private ElasticSpanExporter spanExporter;
 
   public ElasticSpanProcessor(ElasticProfiler profiler, ElasticBreakdownMetrics breakdownMetrics) {
     this.profiler = profiler;
@@ -49,6 +52,8 @@ public boolean isStartRequired() {
   public void onEnd(ReadableSpan span) {
     profiler.onSpanEnd(span);
     breakdownMetrics.onSpanEnd(span);
+
+    captureStackTrace(span);
   }
 
   @Override
@@ -61,4 +66,28 @@ public CompletableResultCode shutdown() {
     profiler.shutdown();
     return CompletableResultCode.ofSuccess();
   }
+
+  public void registerSpanExporter(ElasticSpanExporter spanExporter) {
+    this.spanExporter = spanExporter;
+  }
+
+  private void captureStackTrace(ReadableSpan span) {
+    if (spanExporter == null) {
+      return;
+    }
+
+    // do not overwrite stacktrace if present
+    if (span.getAttribute(ElasticAttributes.SPAN_STACKTRACE) == null) {
+      Throwable exception = new Throwable();
+      StringWriter stringWriter = new StringWriter();
+      try (PrintWriter printWriter = new PrintWriter(stringWriter)) {
+        exception.printStackTrace(printWriter);
+      }
+
+      // TODO should we filter-out the calling code that is within the agent: at least onEnd +
+      // captureStackTrace will be included here
+      spanExporter.addAttribute(
+          span.getSpanContext(), ElasticAttributes.SPAN_STACKTRACE, stringWriter.toString());
+    }
+  }
 }

custom/src/test/java/co/elastic/otel/ElasticSpanProcessorTest.java

@@ -0,0 +1,91 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 co.elastic.otel;
+
+import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.assertThat;
+import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.satisfies;
+import static org.mockito.Mockito.mock;
+
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import org.assertj.core.api.AbstractCharSequenceAssert;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class ElasticSpanProcessorTest {
+
+  private static final Tracer tracer;
+  private static final InMemorySpanExporter testExporter;
+
+  static {
+    ElasticSpanProcessor elasticSpanProcessor =
+        new ElasticSpanProcessor(mock(ElasticProfiler.class), mock(ElasticBreakdownMetrics.class));
+
+    testExporter = InMemorySpanExporter.create();
+    ElasticSpanExporter elasticSpanExporter = new ElasticSpanExporter(testExporter);
+    elasticSpanProcessor.registerSpanExporter(elasticSpanExporter);
+
+    tracer =
+        SdkTracerProvider.builder()
+            .addSpanProcessor(elasticSpanProcessor)
+            .addSpanProcessor(SimpleSpanProcessor.create(elasticSpanExporter))
+            .build()
+            .get("for-testing");
+  }
+
+  @BeforeEach
+  public void before() {
+    testExporter.reset();
+  }
+
+  @Test
+  void spanStackTraceCapture() {
+    tracer.spanBuilder("span").startSpan().end();
+
+    assertThat(testExporter.getFinishedSpanItems())
+        .hasSize(1)
+        .first()
+        .satisfies(
+            spanData ->
+                assertThat(spanData)
+                    .hasAttributesSatisfying(
+                        satisfies(
+                            ElasticAttributes.SPAN_STACKTRACE,
+                            AbstractCharSequenceAssert::isNotEmpty)));
+  }
+
+  @Test
+  void spanStackTraceCaptureDoesNotOverwrite() {
+    String value = "dummy";
+    tracer
+        .spanBuilder("span")
+        .setAttribute(ElasticAttributes.SPAN_STACKTRACE, value)
+        .startSpan()
+        .end();
+
+    assertThat(testExporter.getFinishedSpanItems())
+        .hasSize(1)
+        .first()
+        .satisfies(
+            spanData ->
+                assertThat(spanData).hasAttribute(ElasticAttributes.SPAN_STACKTRACE, value));
+  }
+}

smoke-tests/src/test/java/com/example/javaagent/smoketest/AgentFeaturesSmokeTest.java

@@ -61,7 +61,8 @@ public void healthcheck() throws InterruptedException {
               .containsKeys(
                   "elastic.span.is_local_root",
                   "elastic.span.local_root.id",
-                  "elastic.span.self_time");
+                  "elastic.span.self_time",
+                  "code.stacktrace");
         });
   }