[java] Allow custom tracing to integrate with existing OpenTelemetry tracing

llatinov · pujagani · web-flow · commit 6c2e8440a331 · 2021-11-23T20:09:04.000+05:30
Fixes #9943 Co-authored-by: Puja Jagani <puja.jagani93@gmail.com>
diff --git a/java/src/org/openqa/selenium/remote/tracing/Tracer.java b/java/src/org/openqa/selenium/remote/tracing/Tracer.java
@@ -17,10 +17,14 @@
 
 package org.openqa.selenium.remote.tracing;
 
+import io.opentelemetry.context.Context;
+
 public interface Tracer {
 
   TraceContext getCurrentContext();
 
   Propagator getPropagator();
 
+  void setOpenTelemetryContext(Context context);
+
 }
diff --git a/java/src/org/openqa/selenium/remote/tracing/empty/BUILD.bazel b/java/src/org/openqa/selenium/remote/tracing/empty/BUILD.bazel
@@ -1,3 +1,4 @@
+load("@rules_jvm_external//:defs.bzl", "artifact")
 load("//java:defs.bzl", "java_library")
 
 java_library(
@@ -9,5 +10,6 @@ java_library(
     deps = [
         "//java/src/org/openqa/selenium:core",
         "//java/src/org/openqa/selenium/remote/tracing:tracing-lib",
+        artifact("io.opentelemetry:opentelemetry-context")
     ],
 )
diff --git a/java/src/org/openqa/selenium/remote/tracing/empty/NullTracer.java b/java/src/org/openqa/selenium/remote/tracing/empty/NullTracer.java
@@ -17,6 +17,7 @@
 
 package org.openqa.selenium.remote.tracing.empty;
 
+import io.opentelemetry.context.Context;
 import org.openqa.selenium.remote.tracing.Propagator;
 import org.openqa.selenium.remote.tracing.TraceContext;
 import org.openqa.selenium.remote.tracing.Tracer;
@@ -31,4 +32,8 @@ public TraceContext getCurrentContext() {
   public Propagator getPropagator() {
     return new NullPropagator();
   }
+
+  @Override
+  public void setOpenTelemetryContext(Context context) {
+  }
 }
diff --git a/java/src/org/openqa/selenium/remote/tracing/opentelemetry/OpenTelemetryTracer.java b/java/src/org/openqa/selenium/remote/tracing/opentelemetry/OpenTelemetryTracer.java
@@ -83,6 +83,7 @@ private static OpenTelemetryTracer createTracer() {
 
   private final Tracer tracer;
   private final OpenTelemetryPropagator telemetryPropagator;
+  private Context context;
 
   public OpenTelemetryTracer(Tracer tracer, TextMapPropagator propagator) {
     this.tracer = Require.nonNull("Tracer", tracer);
@@ -93,11 +94,16 @@ public OpenTelemetryTracer(Tracer tracer, TextMapPropagator propagator) {
 
   @Override
   public TraceContext getCurrentContext() {
-    return new OpenTelemetryContext(tracer, Context.current());
+    return new OpenTelemetryContext(tracer, context != null ? context : Context.current());
   }
 
   @Override
   public Propagator getPropagator() {
     return telemetryPropagator;
   }
+
+  @Override
+  public void setOpenTelemetryContext(Context context) {
+    this.context = context;
+  }
 }
diff --git a/java/test/org/openqa/selenium/remote/tracing/opentelemetry/TracerTest.java b/java/test/org/openqa/selenium/remote/tracing/opentelemetry/TracerTest.java
@@ -23,6 +23,7 @@
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.ContextPropagators;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.common.CompletableResultCode;
@@ -102,7 +103,6 @@ public void shouldBeAbleToCreateATracer() {
         StatusCode.ERROR);
     assertThat(values).element(0)
         .extracting(el -> el.getAttributes().get(AttributeKey.stringKey("cheese"))).isEqualTo("gouda");
-
   }
 
   @Test
@@ -587,6 +587,33 @@ public void cleverShenanigansRepresentingWhatWeSeeInTheRouter() {
     routable.execute(new HttpRequest(GET, "/"));
   }
 
+  @Test
+  public void shouldBeAbleToSetExternalContextAndCreatedSpansAreItsChildren() {
+    List<SpanData> allSpans = new ArrayList<>();
+    Tracer tracer = createTracer(allSpans);
+
+    OpenTelemetrySdk openTelemetrySdk = OpenTelemetrySdk.builder().build();
+    io.opentelemetry.api.trace.Span externalSpan = openTelemetrySdk
+        .getTracer("externalTracer")
+        .spanBuilder("externalSpan")
+        .startSpan();
+    Context parentContext = Context.current().with(externalSpan);
+    tracer.setOpenTelemetryContext(parentContext);
+
+    Span parent = tracer.getCurrentContext().createSpan("parent");
+    try (Span child = parent.createSpan("child")) {
+    }
+    parent.close();
+
+    assertThat(allSpans).hasSize(2);
+    assertThat(allSpans.get(0).getName()).isEqualTo("child");
+    assertThat(allSpans.get(0).getParentSpanId())
+        .isEqualTo(parent.getId());
+    assertThat(allSpans.get(1).getName()).isEqualTo("parent");
+    assertThat(allSpans.get(1).getParentSpanId())
+        .isEqualTo(externalSpan.getSpanContext().getSpanId());
+  }
+
   private Tracer createTracer(List<SpanData> exportTo) {
     ContextPropagators propagators =
       ContextPropagators.create((W3CTraceContextPropagator.getInstance()));

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@`
`17`	`17`
`18`	`18`	`package org.openqa.selenium.remote.tracing.empty;`
`19`	`19`
	`20`	`+import io.opentelemetry.context.Context;`
`20`	`21`	`import org.openqa.selenium.remote.tracing.Propagator;`
`21`	`22`	`import org.openqa.selenium.remote.tracing.TraceContext;`
`22`	`23`	`import org.openqa.selenium.remote.tracing.Tracer;`
`@@ -31,4 +32,8 @@ public TraceContext getCurrentContext() {`
`31`	`32`	`public Propagator getPropagator() {`
`32`	`33`	`return new NullPropagator();`
`33`	`34`	`}`
	`35`	`+`
	`36`	`+ @Override`
	`37`	`+ public void setOpenTelemetryContext(Context context) {`
	`38`	`+ }`
`34`	`39`	`}`