Context propagation to elasticsearch-transport callbacks

instrumentation/elasticsearch/elasticsearch-transport-5.0/javaagent/build.gradle.kts

@@ -33,6 +33,7 @@ dependencies {
   testInstrumentation(project(":instrumentation:apache-httpasyncclient-4.1:javaagent"))
   testInstrumentation(project(":instrumentation:netty:netty-4.1:javaagent"))
 
+  testImplementation(project(":instrumentation:elasticsearch:elasticsearch-transport-testing"))
   testImplementation("org.apache.logging.log4j:log4j-core:2.11.0")
   testImplementation("org.apache.logging.log4j:log4j-api:2.11.0")
 

instrumentation/elasticsearch/elasticsearch-transport-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/transport/v5_0/AbstractClientInstrumentation.java

@@ -54,11 +54,13 @@ public static void onEnter(
         @Advice.Argument(value = 2, readOnly = false)
             ActionListener<ActionResponse> actionListener) {
 
-      context = tracer().startSpan(currentContext(), null, action);
+      Context parentContext = currentContext();
+      context = tracer().startSpan(parentContext, null, action);
       scope = context.makeCurrent();
 
       tracer().onRequest(context, action.getClass(), actionRequest.getClass());
-      actionListener = new TransportActionListener<>(actionRequest, actionListener, context);
+      actionListener =
+          new TransportActionListener<>(actionRequest, actionListener, context, parentContext);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/elasticsearch/elasticsearch-transport-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/transport/v5_0/TransportActionListener.java

@@ -9,6 +9,7 @@
 
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.config.Config;
 import io.opentelemetry.instrumentation.api.tracer.net.NetPeerAttributes;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
@@ -34,11 +35,16 @@ public class TransportActionListener<T extends ActionResponse> implements Action
 
   private final ActionListener<T> listener;
   private final Context context;
+  private final Context parentContext;
 
   public TransportActionListener(
-      ActionRequest<?> actionRequest, ActionListener<T> listener, Context context) {
+      ActionRequest<?> actionRequest,
+      ActionListener<T> listener,
+      Context context,
+      Context parentContext) {
     this.listener = listener;
     this.context = context;
+    this.parentContext = parentContext;
     onRequest(actionRequest);
   }
 
@@ -124,12 +130,16 @@ public void onResponse(T response) {
     }
 
     tracer().end(context);
-    listener.onResponse(response);
+    try (Scope ignored = parentContext.makeCurrent()) {
+      listener.onResponse(response);
+    }
   }
 
   @Override
   public void onFailure(Exception e) {
     tracer().endExceptionally(context, e);
-    listener.onFailure(e);
+    try (Scope ignored = parentContext.makeCurrent()) {
+      listener.onFailure(e);
+    }
   }
 }

instrumentation/elasticsearch/elasticsearch-transport-5.0/javaagent/src/test/groovy/Elasticsearch5NodeClientTest.groovy

@@ -4,12 +4,12 @@
  */
 
 import static io.opentelemetry.api.trace.SpanKind.CLIENT
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL
 import static io.opentelemetry.api.trace.StatusCode.ERROR
 import static org.elasticsearch.cluster.ClusterName.CLUSTER_NAME_SETTING
 
-import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes
-import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest
+import org.elasticsearch.client.Client
 import org.elasticsearch.common.io.FileSystemUtils
 import org.elasticsearch.common.settings.Settings
 import org.elasticsearch.env.Environment
@@ -18,8 +18,9 @@ import org.elasticsearch.node.Node
 import org.elasticsearch.node.internal.InternalSettingsPreparer
 import org.elasticsearch.transport.Netty3Plugin
 import spock.lang.Shared
+import spock.lang.Unroll
 
-class Elasticsearch5NodeClientTest extends AgentInstrumentationSpecification {
+class Elasticsearch5NodeClientTest extends AbstractElasticsearchNodeClientTest {
   public static final long TIMEOUT = 10000 // 10 seconds
 
   @Shared
@@ -29,7 +30,8 @@ class Elasticsearch5NodeClientTest extends AgentInstrumentationSpecification {
   @Shared
   String clusterName = UUID.randomUUID().toString()
 
-  def client = testNode.client()
+  @Shared
+  Client client
 
   def setupSpec() {
 
@@ -48,10 +50,11 @@ class Elasticsearch5NodeClientTest extends AgentInstrumentationSpecification {
       .build()
     testNode = new Node(new Environment(InternalSettingsPreparer.prepareSettings(settings)), [Netty3Plugin])
     testNode.start()
+    client = testNode.client()
     runWithSpan("setup") {
       // this may potentially create multiple requests and therefore multiple spans, so we wrap this call
       // into a top level trace to get exactly one trace in the result.
-      testNode.client().admin().cluster().prepareHealth().setWaitForYellowStatus().execute().actionGet(TIMEOUT)
+      client.admin().cluster().prepareHealth().setWaitForYellowStatus().execute().actionGet(TIMEOUT)
     }
     ignoreTracesAndClear(1)
   }
@@ -64,46 +67,79 @@ class Elasticsearch5NodeClientTest extends AgentInstrumentationSpecification {
     }
   }
 
-  def "test elasticsearch status"() {
-    setup:
-    def result = client.admin().cluster().health(new ClusterHealthRequest())
+  @Override
+  Client client() {
+    client
+  }
 
-    def clusterHealthStatus = result.get().status
+  @Unroll
+  def "test elasticsearch status #callKind"() {
+    setup:
+    def clusterHealthStatus = runWithSpan("parent") {
+      call.call()
+    }
 
     expect:
     clusterHealthStatus.name() == "GREEN"
 
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
         span(0) {
+          name "parent"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "ClusterHealthAction"
           kind CLIENT
+          childOf(span(0))
           attributes {
             "${SemanticAttributes.DB_SYSTEM.key}" "elasticsearch"
             "${SemanticAttributes.DB_OPERATION.key}" "ClusterHealthAction"
             "elasticsearch.action" "ClusterHealthAction"
             "elasticsearch.request" "ClusterHealthRequest"
           }
         }
+        span(2) {
+          name "callback"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
+
+    where:
+    callKind | call
+    "sync"   | { clusterHealthSync() }
+    "async"  | { clusterHealthAsync() }
   }
 
-  def "test elasticsearch error"() {
+  @Unroll
+  def "test elasticsearch error #callKind"() {
     when:
-    client.prepareGet(indexName, indexType, id).get()
+    runWithSpan("parent") {
+      call.call(indexName, indexType, id)
+    }
 
     then:
     thrown IndexNotFoundException
 
     and:
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
         span(0) {
+          name "parent"
+          status ERROR
+          errorEvent IndexNotFoundException, "no such index"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "GetAction"
           status ERROR
           errorEvent IndexNotFoundException, "no such index"
           kind CLIENT
+          childOf(span(0))
           attributes {
             "${SemanticAttributes.DB_SYSTEM.key}" "elasticsearch"
             "${SemanticAttributes.DB_OPERATION.key}" "GetAction"
@@ -112,13 +148,21 @@ class Elasticsearch5NodeClientTest extends AgentInstrumentationSpecification {
             "elasticsearch.request.indices" indexName
           }
         }
+        span(2) {
+          name "callback"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
 
     where:
     indexName = "invalid-index"
     indexType = "test-type"
     id = "1"
+    callKind | call
+    "sync"   | { indexName, indexType, id -> prepareGetSync(indexName, indexType, id) }
+    "async"  | { indexName, indexType, id -> prepareGetAsync(indexName, indexType, id) }
   }
 
   def "test elasticsearch get"() {

instrumentation/elasticsearch/elasticsearch-transport-5.0/javaagent/src/test/groovy/Elasticsearch5TransportClientTest.groovy

@@ -4,12 +4,11 @@
  */
 
 import static io.opentelemetry.api.trace.SpanKind.CLIENT
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL
 import static io.opentelemetry.api.trace.StatusCode.ERROR
 import static org.elasticsearch.cluster.ClusterName.CLUSTER_NAME_SETTING
 
-import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes
-import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest
 import org.elasticsearch.client.transport.TransportClient
 import org.elasticsearch.common.io.FileSystemUtils
 import org.elasticsearch.common.settings.Settings
@@ -23,8 +22,9 @@ import org.elasticsearch.transport.RemoteTransportException
 import org.elasticsearch.transport.TransportService
 import org.elasticsearch.transport.client.PreBuiltTransportClient
 import spock.lang.Shared
+import spock.lang.Unroll
 
-class Elasticsearch5TransportClientTest extends AgentInstrumentationSpecification {
+class Elasticsearch5TransportClientTest extends AbstractElasticsearchTransportClientTest {
   public static final long TIMEOUT = 10000 // 10 seconds
 
   @Shared
@@ -80,18 +80,29 @@ class Elasticsearch5TransportClientTest extends AgentInstrumentationSpecificatio
     }
   }
 
-  def "test elasticsearch status"() {
-    setup:
-    def result = client.admin().cluster().health(new ClusterHealthRequest())
+  @Override
+  TransportClient client() {
+    client
+  }
 
-    def clusterHealthStatus = result.get().status
+  @Unroll
+  def "test elasticsearch status #callKind"() {
+    setup:
+    def clusterHealthStatus = runWithSpan("parent") {
+      call.call()
+    }
 
     expect:
     clusterHealthStatus.name() == "GREEN"
 
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
         span(0) {
+          name "parent"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "ClusterHealthAction"
           kind CLIENT
           attributes {
@@ -104,24 +115,44 @@ class Elasticsearch5TransportClientTest extends AgentInstrumentationSpecificatio
             "elasticsearch.request" "ClusterHealthRequest"
           }
         }
+        span(2) {
+          name "callback"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
+
+    where:
+    callKind | call
+    "sync"   | { clusterHealthSync() }
+    "async"  | { clusterHealthAsync() }
   }
 
-  def "test elasticsearch error"() {
+  def "test elasticsearch error #callKind"() {
     when:
-    client.prepareGet(indexName, indexType, id).get()
+    runWithSpan("parent") {
+      call.call(indexName, indexType, id)
+    }
 
     then:
     thrown IndexNotFoundException
 
     and:
     assertTraces(1) {
-      trace(0, 1) {
+      trace(0, 3) {
         span(0) {
+          name "parent"
+          status ERROR
+          errorEvent IndexNotFoundException, "no such index"
+          kind INTERNAL
+          hasNoParent()
+        }
+        span(1) {
           name "GetAction"
           kind CLIENT
           status ERROR
+          childOf(span(0))
           errorEvent RemoteTransportException, String
           attributes {
             "${SemanticAttributes.DB_SYSTEM.key}" "elasticsearch"
@@ -131,13 +162,21 @@ class Elasticsearch5TransportClientTest extends AgentInstrumentationSpecificatio
             "elasticsearch.request.indices" indexName
           }
         }
+        span(2) {
+          name "callback"
+          kind INTERNAL
+          childOf(span(0))
+        }
       }
     }
 
     where:
     indexName = "invalid-index"
     indexType = "test-type"
     id = "1"
+    callKind | call
+    "sync"   | { indexName, indexType, id -> prepareGetSync(indexName, indexType, id) }
+    "async"  | { indexName, indexType, id -> prepareGetAsync(indexName, indexType, id) }
   }
 
   def "test elasticsearch get"() {

instrumentation/elasticsearch/elasticsearch-transport-5.3/javaagent/build.gradle.kts

@@ -38,6 +38,7 @@ dependencies {
   testInstrumentation(project(":instrumentation:netty:netty-4.1:javaagent"))
   testInstrumentation(project(":instrumentation:spring:spring-data-1.8:javaagent"))
 
+  testImplementation(project(":instrumentation:elasticsearch:elasticsearch-transport-testing"))
   testImplementation("org.apache.logging.log4j:log4j-core:2.11.0")
   testImplementation("org.apache.logging.log4j:log4j-api:2.11.0")
 

instrumentation/elasticsearch/elasticsearch-transport-5.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/transport/v5_3/AbstractClientInstrumentation.java

@@ -54,11 +54,13 @@ public static void onEnter(
         @Advice.Argument(value = 2, readOnly = false)
             ActionListener<ActionResponse> actionListener) {
 
-      context = tracer().startSpan(currentContext(), null, action);
+      Context parentContext = currentContext();
+      context = tracer().startSpan(parentContext, null, action);
       scope = context.makeCurrent();
 
       tracer().onRequest(context, action.getClass(), actionRequest.getClass());
-      actionListener = new TransportActionListener<>(actionRequest, actionListener, context);
+      actionListener =
+          new TransportActionListener<>(actionRequest, actionListener, context, parentContext);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)