Add response type parameter to db attributes extractor (open-telemetr…

…y#3093)

* Add response type parameter to DbAttributesExtractor

* AutoValue

* ResultSet

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/db/DbAttributesExtractor.java

@@ -19,7 +19,8 @@
  * methods, but implement as many as possible for best compliance with the OpenTelemetry
  * specification.
  */
-public abstract class DbAttributesExtractor<REQUEST> extends AttributesExtractor<REQUEST, Void> {
+public abstract class DbAttributesExtractor<REQUEST, RESPONSE>
+    extends AttributesExtractor<REQUEST, RESPONSE> {
   @Override
   protected void onStart(AttributesBuilder attributes, REQUEST request) {
     set(attributes, SemanticAttributes.DB_SYSTEM, system(request));
@@ -32,7 +33,7 @@ protected void onStart(AttributesBuilder attributes, REQUEST request) {
 
   @Override
   protected final void onEnd(
-      AttributesBuilder attributes, REQUEST request, @Nullable Void unused) {}
+      AttributesBuilder attributes, REQUEST request, @Nullable RESPONSE response) {}
 
   @Nullable
   protected abstract String system(REQUEST request);

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/db/DbSpanNameExtractor.java

@@ -18,15 +18,15 @@ public final class DbSpanNameExtractor<REQUEST> implements SpanNameExtractor<REQ
    * @see SqlAttributesExtractor#table(Object) used to extract {@code <db.table>}.
    */
   public static <REQUEST> SpanNameExtractor<REQUEST> create(
-      DbAttributesExtractor<REQUEST> attributesExtractor) {
+      DbAttributesExtractor<REQUEST, ?> attributesExtractor) {
     return new DbSpanNameExtractor<>(attributesExtractor);
   }
 
   private static final String DEFAULT_SPAN_NAME = "DB Query";
 
-  private final DbAttributesExtractor<REQUEST> attributesExtractor;
+  private final DbAttributesExtractor<REQUEST, ?> attributesExtractor;
 
-  private DbSpanNameExtractor(DbAttributesExtractor<REQUEST> attributesExtractor) {
+  private DbSpanNameExtractor(DbAttributesExtractor<REQUEST, ?> attributesExtractor) {
     this.attributesExtractor = attributesExtractor;
   }
 
@@ -59,7 +59,7 @@ public String extract(REQUEST request) {
   @Nullable
   private String getTableName(REQUEST request) {
     if (attributesExtractor instanceof SqlAttributesExtractor) {
-      return ((SqlAttributesExtractor<REQUEST>) attributesExtractor).table(request);
+      return ((SqlAttributesExtractor<REQUEST, ?>) attributesExtractor).table(request);
     }
     return null;
   }

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/db/SqlAttributesExtractor.java

@@ -22,7 +22,8 @@
  *   <li>It enables adding the table name extracted by the sanitizer as a parameter.
  * </ul>
  */
-public abstract class SqlAttributesExtractor<REQUEST> extends DbAttributesExtractor<REQUEST> {
+public abstract class SqlAttributesExtractor<REQUEST, RESPONSE>
+    extends DbAttributesExtractor<REQUEST, RESPONSE> {
 
   @Override
   protected final void onStart(AttributesBuilder attributes, REQUEST request) {

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/db/DbAttributesExtractorTest.java

@@ -17,8 +17,8 @@
 import org.junit.jupiter.api.Test;
 
 class DbAttributesExtractorTest {
-  DbAttributesExtractor<Map<String, String>> underTest =
-      new DbAttributesExtractor<Map<String, String>>() {
+  DbAttributesExtractor<Map<String, String>, Void> underTest =
+      new DbAttributesExtractor<Map<String, String>, Void>() {
         @Override
         protected String system(Map<String, String> map) {
           return map.get("db.system");

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/db/DbSpanNameExtractorTest.java

@@ -16,8 +16,8 @@
 
 @ExtendWith(MockitoExtension.class)
 class DbSpanNameExtractorTest {
-  @Mock DbAttributesExtractor<DbRequest> dbAttributesExtractor;
-  @Mock SqlAttributesExtractor<DbRequest> sqlAttributesExtractor;
+  @Mock DbAttributesExtractor<DbRequest, Void> dbAttributesExtractor;
+  @Mock SqlAttributesExtractor<DbRequest, Void> sqlAttributesExtractor;
 
   @Test
   void shouldExtractFullSpanName() {

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/db/SqlAttributesExtractorTest.java

@@ -19,8 +19,8 @@
 
 class SqlAttributesExtractorTest {
   AttributeKey<String> dbTableAttribute;
-  SqlAttributesExtractor<Map<String, String>> underTest =
-      new SqlAttributesExtractor<Map<String, String>>() {
+  SqlAttributesExtractor<Map<String, String>, Void> underTest =
+      new SqlAttributesExtractor<Map<String, String>, Void>() {
 
         @Override
         protected AttributeKey<String> dbTableAttribute() {

instrumentation/cassandra/cassandra-3.0/javaagent/cassandra-3.0-javaagent.gradle

@@ -29,6 +29,9 @@ muzzle {
 dependencies {
   library "com.datastax.cassandra:cassandra-driver-core:3.0.0"
 
+  compileOnly deps.autoValueAnnotations
+  annotationProcessor deps.autoValue
+
   testLibrary "com.datastax.cassandra:cassandra-driver-core:3.2.0"
   testInstrumentation project(':instrumentation:guava-10.0:javaagent')
 

instrumentation/cassandra/cassandra-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/cassandra/v3_0/CassandraInstrumenters.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.cassandra.v3_0;
 
+import com.datastax.driver.core.ResultSet;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
@@ -15,23 +16,23 @@
 public final class CassandraInstrumenters {
   private static final String INSTRUMENTATION_NAME = "io.opentelemetry.javaagent.cassandra-3.0";
 
-  private static final Instrumenter<CassandraRequest, Void> INSTRUMENTER;
+  private static final Instrumenter<CassandraRequest, ResultSet> INSTRUMENTER;
 
   static {
-    DbAttributesExtractor<CassandraRequest> attributesExtractor =
+    DbAttributesExtractor<CassandraRequest, ResultSet> attributesExtractor =
         new CassandraSqlAttributesExtractor();
     SpanNameExtractor<CassandraRequest> spanName = DbSpanNameExtractor.create(attributesExtractor);
 
     INSTRUMENTER =
-        Instrumenter.<CassandraRequest, Void>newBuilder(
+        Instrumenter.<CassandraRequest, ResultSet>newBuilder(
                 GlobalOpenTelemetry.get(), INSTRUMENTATION_NAME, spanName)
             .addAttributesExtractor(attributesExtractor)
             .addAttributesExtractor(new CassandraNetAttributesExtractor())
             .addAttributesExtractor(new CassandraKeyspaceExtractor())
             .newInstrumenter(SpanKindExtractor.alwaysClient());
   }
 
-  public static Instrumenter<CassandraRequest, Void> instrumenter() {
+  public static Instrumenter<CassandraRequest, ResultSet> instrumenter() {
     return INSTRUMENTER;
   }
 

instrumentation/cassandra/cassandra-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/cassandra/v3_0/CassandraKeyspaceExtractor.java

@@ -5,11 +5,12 @@
 
 package io.opentelemetry.javaagent.instrumentation.cassandra.v3_0;
 
+import com.datastax.driver.core.ResultSet;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 
-final class CassandraKeyspaceExtractor extends AttributesExtractor<CassandraRequest, Void> {
+final class CassandraKeyspaceExtractor extends AttributesExtractor<CassandraRequest, ResultSet> {
 
   @Override
   protected void onStart(AttributesBuilder attributes, CassandraRequest request) {
@@ -18,5 +19,6 @@ protected void onStart(AttributesBuilder attributes, CassandraRequest request) {
   }
 
   @Override
-  protected void onEnd(AttributesBuilder attributes, CassandraRequest request, Void unused) {}
+  protected void onEnd(
+      AttributesBuilder attributes, CassandraRequest request, ResultSet response) {}
 }

instrumentation/cassandra/cassandra-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/cassandra/v3_0/CassandraNetAttributesExtractor.java

@@ -5,13 +5,13 @@
 
 package io.opentelemetry.javaagent.instrumentation.cassandra.v3_0;
 
-import com.datastax.driver.core.ExecutionInfo;
+import com.datastax.driver.core.ResultSet;
 import io.opentelemetry.instrumentation.api.instrumenter.net.InetSocketAddressNetAttributesExtractor;
 import java.net.InetSocketAddress;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class CassandraNetAttributesExtractor
-    extends InetSocketAddressNetAttributesExtractor<CassandraRequest, Void> {
+    extends InetSocketAddressNetAttributesExtractor<CassandraRequest, ResultSet> {
 
   @Override
   @Nullable
@@ -20,8 +20,10 @@ public String transport(CassandraRequest request) {
   }
 
   @Override
-  public @Nullable InetSocketAddress getAddress(CassandraRequest request, @Nullable Void response) {
-    ExecutionInfo executionInfo = request.getExecutionInfo();
-    return executionInfo == null ? null : executionInfo.getQueriedHost().getSocketAddress();
+  public @Nullable InetSocketAddress getAddress(
+      CassandraRequest request, @Nullable ResultSet response) {
+    return response == null
+        ? null
+        : response.getExecutionInfo().getQueriedHost().getSocketAddress();
   }
 }

instrumentation/cassandra/cassandra-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/cassandra/v3_0/CassandraRequest.java

@@ -5,34 +5,17 @@
 
 package io.opentelemetry.javaagent.instrumentation.cassandra.v3_0;
 
-import com.datastax.driver.core.ExecutionInfo;
 import com.datastax.driver.core.Session;
+import com.google.auto.value.AutoValue;
 
-public final class CassandraRequest {
+@AutoValue
+public abstract class CassandraRequest {
 
-  private final Session session;
-  private final String statement;
-  // volatile is not needed here as this field is set and get from the same thread
-  private ExecutionInfo executionInfo;
+  public abstract Session getSession();
 
-  public CassandraRequest(Session session, String statement) {
-    this.session = session;
-    this.statement = statement;
-  }
-
-  public Session getSession() {
-    return session;
-  }
-
-  public String getStatement() {
-    return statement;
-  }
-
-  public void setExecutionInfo(ExecutionInfo executionInfo) {
-    this.executionInfo = executionInfo;
-  }
+  public abstract String getStatement();
 
-  public ExecutionInfo getExecutionInfo() {
-    return executionInfo;
+  public static CassandraRequest create(Session session, String statement) {
+    return new AutoValue_CassandraRequest(session, statement);
   }
 }

instrumentation/cassandra/cassandra-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/cassandra/v3_0/CassandraSqlAttributesExtractor.java

@@ -5,12 +5,14 @@
 
 package io.opentelemetry.javaagent.instrumentation.cassandra.v3_0;
 
+import com.datastax.driver.core.ResultSet;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.instrumentation.api.instrumenter.db.SqlAttributesExtractor;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
-final class CassandraSqlAttributesExtractor extends SqlAttributesExtractor<CassandraRequest> {
+final class CassandraSqlAttributesExtractor
+    extends SqlAttributesExtractor<CassandraRequest, ResultSet> {
 
   @Override
   protected String system(CassandraRequest request) {

instrumentation/cassandra/cassandra-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/cassandra/v3_0/TracingSession.java

@@ -48,7 +48,7 @@ public ListenableFuture<Session> initAsync() {
 
   @Override
   public ResultSet execute(String query) {
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     ResultSet resultSet;
     try (Scope ignored = context.makeCurrent()) {
@@ -57,14 +57,13 @@ public ResultSet execute(String query) {
       instrumenter().end(context, request, null, t);
       throw t;
     }
-    request.setExecutionInfo(resultSet.getExecutionInfo());
-    instrumenter().end(context, request, null, null);
+    instrumenter().end(context, request, resultSet, null);
     return resultSet;
   }
 
   @Override
   public ResultSet execute(String query, Object... values) {
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     ResultSet resultSet;
     try (Scope ignored = context.makeCurrent()) {
@@ -73,14 +72,13 @@ public ResultSet execute(String query, Object... values) {
       instrumenter().end(context, request, null, t);
       throw t;
     }
-    request.setExecutionInfo(resultSet.getExecutionInfo());
-    instrumenter().end(context, request, null, null);
+    instrumenter().end(context, request, resultSet, null);
     return resultSet;
   }
 
   @Override
   public ResultSet execute(String query, Map<String, Object> values) {
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     ResultSet resultSet;
     try (Scope ignored = context.makeCurrent()) {
@@ -89,15 +87,14 @@ public ResultSet execute(String query, Map<String, Object> values) {
       instrumenter().end(context, request, null, t);
       throw t;
     }
-    request.setExecutionInfo(resultSet.getExecutionInfo());
-    instrumenter().end(context, request, null, null);
+    instrumenter().end(context, request, resultSet, null);
     return resultSet;
   }
 
   @Override
   public ResultSet execute(Statement statement) {
     String query = getQuery(statement);
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     ResultSet resultSet;
     try (Scope ignored = context.makeCurrent()) {
@@ -106,14 +103,13 @@ public ResultSet execute(Statement statement) {
       instrumenter().end(context, request, null, t);
       throw t;
     }
-    request.setExecutionInfo(resultSet.getExecutionInfo());
-    instrumenter().end(context, request, null, null);
+    instrumenter().end(context, request, resultSet, null);
     return resultSet;
   }
 
   @Override
   public ResultSetFuture executeAsync(String query) {
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     try (Scope ignored = context.makeCurrent()) {
       ResultSetFuture future = session.executeAsync(query);
@@ -124,7 +120,7 @@ public ResultSetFuture executeAsync(String query) {
 
   @Override
   public ResultSetFuture executeAsync(String query, Object... values) {
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     try (Scope ignored = context.makeCurrent()) {
       ResultSetFuture future = session.executeAsync(query, values);
@@ -135,7 +131,7 @@ public ResultSetFuture executeAsync(String query, Object... values) {
 
   @Override
   public ResultSetFuture executeAsync(String query, Map<String, Object> values) {
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     try (Scope ignored = context.makeCurrent()) {
       ResultSetFuture future = session.executeAsync(query, values);
@@ -147,7 +143,7 @@ public ResultSetFuture executeAsync(String query, Map<String, Object> values) {
   @Override
   public ResultSetFuture executeAsync(Statement statement) {
     String query = getQuery(statement);
-    CassandraRequest request = new CassandraRequest(session, query);
+    CassandraRequest request = CassandraRequest.create(session, query);
     Context context = instrumenter().start(Context.current(), request);
     try (Scope ignored = context.makeCurrent()) {
       ResultSetFuture future = session.executeAsync(statement);
@@ -219,8 +215,7 @@ private void addCallbackToEndSpan(
         new FutureCallback<ResultSet>() {
           @Override
           public void onSuccess(ResultSet resultSet) {
-            request.setExecutionInfo(resultSet.getExecutionInfo());
-            instrumenter().end(context, request, null, null);
+            instrumenter().end(context, request, resultSet, null);
           }
 
           @Override

instrumentation/geode-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/geode/GeodeDbAttributesExtractor.java

@@ -10,7 +10,7 @@
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
-final class GeodeDbAttributesExtractor extends DbAttributesExtractor<GeodeRequest> {
+final class GeodeDbAttributesExtractor extends DbAttributesExtractor<GeodeRequest, Void> {
   @Override
   protected String system(GeodeRequest request) {
     return SemanticAttributes.DbSystemValues.GEODE;

instrumentation/geode-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/geode/GeodeInstrumenters.java

@@ -18,7 +18,8 @@ public final class GeodeInstrumenters {
   private static final Instrumenter<GeodeRequest, Void> INSTRUMENTER;
 
   static {
-    DbAttributesExtractor<GeodeRequest> attributesExtractor = new GeodeDbAttributesExtractor();
+    DbAttributesExtractor<GeodeRequest, Void> attributesExtractor =
+        new GeodeDbAttributesExtractor();
     SpanNameExtractor<GeodeRequest> spanName = DbSpanNameExtractor.create(attributesExtractor);
 
     INSTRUMENTER =

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/JdbcAttributesExtractor.java

@@ -10,7 +10,7 @@
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
-final class JdbcAttributesExtractor extends SqlAttributesExtractor<DbRequest> {
+final class JdbcAttributesExtractor extends SqlAttributesExtractor<DbRequest, Void> {
   @Nullable
   @Override
   protected String system(DbRequest dbRequest) {

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/JdbcInstrumenters.java

@@ -19,7 +19,7 @@ public final class JdbcInstrumenters {
   private static final Instrumenter<DbRequest, Void> INSTRUMENTER;
 
   static {
-    DbAttributesExtractor<DbRequest> dbAttributesExtractor = new JdbcAttributesExtractor();
+    DbAttributesExtractor<DbRequest, Void> dbAttributesExtractor = new JdbcAttributesExtractor();
     SpanNameExtractor<DbRequest> spanName = DbSpanNameExtractor.create(dbAttributesExtractor);
     JdbcNetAttributesExtractor netAttributesExtractor = new JdbcNetAttributesExtractor();