Fixed OpenTelemetry span events (#9864)

* Fixed OpenTelemetry span events

* Missing interface

* Refactoring

* Added comment

* Spotless

* Exclude SpanWrapper from coverage tests

Author: ValentinZakharov

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelSpan.java

@@ -17,6 +17,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
 import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
+import datadog.trace.bootstrap.instrumentation.api.SpanWrapper;
 import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
@@ -31,7 +32,7 @@
 import javax.annotation.ParametersAreNonnullByDefault;
 
 @ParametersAreNonnullByDefault
-public class OtelSpan implements Span, WithAgentSpan {
+public class OtelSpan implements Span, WithAgentSpan, SpanWrapper {
   private final AgentSpan delegate;
   private StatusCode statusCode;
   private boolean recording;
@@ -142,16 +143,12 @@ public Span updateName(String name) {
   @Override
   public void end() {
     this.recording = false;
-    applyNamingConvention(this.delegate);
-    setEventsAsTag(this.delegate, this.events);
     this.delegate.finish();
   }
 
   @Override
   public void end(long timestamp, TimeUnit unit) {
     this.recording = false;
-    applyNamingConvention(this.delegate);
-    setEventsAsTag(this.delegate, this.events);
     this.delegate.finish(unit.toMicros(timestamp));
   }
 
@@ -178,6 +175,12 @@ public AgentSpan asAgentSpan() {
     return this.delegate;
   }
 
+  @Override
+  public void onSpanFinished() {
+    applyNamingConvention(this.delegate);
+    setEventsAsTag(this.delegate, this.events);
+  }
+
   private static class NoopSpan implements Span {
     private static final Span INSTANCE = new NoopSpan();
 

dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/main/java/datadog/trace/instrumentation/opentelemetry/OtelSpan.java

@@ -3,6 +3,7 @@
 import datadog.trace.api.DDTags;
 import datadog.trace.api.interceptor.MutableSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.SpanWrapper;
 import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan;
 import io.opentelemetry.common.AttributeValue;
 import io.opentelemetry.trace.EndSpanOptions;
@@ -13,7 +14,7 @@
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
-public class OtelSpan implements Span, MutableSpan, WithAgentSpan {
+public class OtelSpan implements Span, MutableSpan, WithAgentSpan, SpanWrapper {
   private final AgentSpan delegate;
   private final TypeConverter converter;
 

dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/main/java/datadog/trace/instrumentation/opentelemetry/TypeConverter.java

@@ -41,7 +41,7 @@ public Span toSpan(final AgentSpan agentSpan) {
       if (wrapper instanceof Span) {
         return (Span) wrapper;
       }
-      Span spanWrapper = new OtelSpan(agentSpan, this);
+      OtelSpan spanWrapper = new OtelSpan(agentSpan, this);
       attachableSpanWrapper.attachWrapper(spanWrapper);
       return spanWrapper;
     }

dd-java-agent/instrumentation/opentracing/api-0.31/src/main/java/datadog/trace/instrumentation/opentracing31/OTSpan.java

@@ -4,6 +4,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
+import datadog.trace.bootstrap.instrumentation.api.SpanWrapper;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan;
 import datadog.trace.instrumentation.opentracing.LogHandler;
@@ -15,7 +16,7 @@
  * This class should be castable to MutableSpan since that is the way we've encouraged users to
  * interact with non-ot parts of our API.
  */
-class OTSpan implements Span, MutableSpan, WithAgentSpan {
+class OTSpan implements Span, MutableSpan, WithAgentSpan, SpanWrapper {
   private final AgentSpan delegate;
   private final TypeConverter converter;
   private final LogHandler logHandler;

dd-java-agent/instrumentation/opentracing/api-0.32/src/main/java/datadog/trace/instrumentation/opentracing32/OTSpan.java

@@ -4,6 +4,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
+import datadog.trace.bootstrap.instrumentation.api.SpanWrapper;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan;
 import datadog.trace.instrumentation.opentracing.LogHandler;
@@ -16,7 +17,7 @@
  * This class should be castable to MutableSpan since that is the way we've encouraged users to
  * interact with non-ot parts of our API.
  */
-class OTSpan implements Span, MutableSpan, WithAgentSpan {
+class OTSpan implements Span, MutableSpan, WithAgentSpan, SpanWrapper {
   private final AgentSpan delegate;
   private final TypeConverter converter;
   private final LogHandler logHandler;

dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java

@@ -28,8 +28,10 @@
 import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
 import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
+import datadog.trace.bootstrap.instrumentation.api.SpanWrapper;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.core.util.StackTraces;
+import edu.umd.cs.findbugs.annotations.NonNull;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -94,9 +96,9 @@ static DDSpan create(
   private volatile EndpointTracker endpointTracker;
 
   // Cached OT/OTel wrapper to avoid multiple allocations, e.g. when span is activated
-  private volatile Object wrapper;
-  private static final AtomicReferenceFieldUpdater<DDSpan, Object> WRAPPER_FIELD_UPDATER =
-      AtomicReferenceFieldUpdater.newUpdater(DDSpan.class, Object.class, "wrapper");
+  private volatile SpanWrapper wrapper;
+  private static final AtomicReferenceFieldUpdater<DDSpan, SpanWrapper> WRAPPER_FIELD_UPDATER =
+      AtomicReferenceFieldUpdater.newUpdater(DDSpan.class, SpanWrapper.class, "wrapper");
 
   // the request is to be blocked (AppSec)
   private volatile Flow.Action.RequestBlockingAction requestBlockingAction;
@@ -148,6 +150,10 @@ private void finishAndAddToTrace(final long durationNano) {
     // ensure a min duration of 1
     if (DURATION_NANO_UPDATER.compareAndSet(this, 0, Math.max(1, durationNano))) {
       setLongRunningVersion(-this.longRunningVersion);
+      SpanWrapper wrapper = getWrapper();
+      if (wrapper != null) {
+        wrapper.onSpanFinished();
+      }
       this.metrics.onSpanFinished();
       TraceCollector.PublishState publishState = context.getTraceCollector().onPublish(this);
       log.debug("Finished span ({}): {}", publishState, this);
@@ -804,12 +810,12 @@ public String toString() {
   }
 
   @Override
-  public void attachWrapper(Object wrapper) {
+  public void attachWrapper(@NonNull SpanWrapper wrapper) {
     WRAPPER_FIELD_UPDATER.compareAndSet(this, null, wrapper);
   }
 
   @Override
-  public Object getWrapper() {
+  public SpanWrapper getWrapper() {
     return WRAPPER_FIELD_UPDATER.get(this);
   }
 

dd-trace-ot/src/main/java/datadog/opentracing/OTSpan.java

@@ -4,6 +4,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
+import datadog.trace.bootstrap.instrumentation.api.SpanWrapper;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan;
 import io.opentracing.Span;
@@ -15,7 +16,7 @@
  * This class should be castable to MutableSpan since that is the way we've encouraged users to
  * interact with non-ot parts of our API.
  */
-class OTSpan implements Span, MutableSpan, WithAgentSpan {
+class OTSpan implements Span, MutableSpan, WithAgentSpan, SpanWrapper {
   private final AgentSpan delegate;
   private final TypeConverter converter;
   private final LogHandler logHandler;

internal-api/build.gradle.kts

@@ -110,6 +110,7 @@ val excludedClassesCoverage by extra(
     "datadog.trace.bootstrap.instrumentation.api.ScopeSource",
     "datadog.trace.bootstrap.instrumentation.api.ScopedContext",
     "datadog.trace.bootstrap.instrumentation.api.ScopedContextKey",
+    "datadog.trace.bootstrap.instrumentation.api.SpanWrapper",
     "datadog.trace.bootstrap.instrumentation.api.TagContext",
     "datadog.trace.bootstrap.instrumentation.api.TagContext.HttpHeaders",
     "datadog.trace.api.civisibility.config.TestIdentifier",

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AttachableWrapper.java

@@ -7,9 +7,9 @@
 public interface AttachableWrapper {
 
   /** Attaches a OT/OTel wrapper to a tracer object. */
-  void attachWrapper(@Nonnull Object wrapper);
+  void attachWrapper(@Nonnull SpanWrapper wrapper);
 
   /** Returns an attached OT/OTel wrapper or null. */
   @Nullable
-  Object getWrapper();
+  SpanWrapper getWrapper();
 }

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanWrapper.java

@@ -0,0 +1,10 @@
+package datadog.trace.bootstrap.instrumentation.api;
+
+/**
+ * Interface implemented by span wrappers such as OpenTelemetry and OpenTracing spans. Provides a
+ * callback that is invoked when the underlying span is finished. This guarantees that wrapper
+ * specific logic runs when spans are auto-finished.
+ */
+public interface SpanWrapper {
+  default void onSpanFinished() {}
+}