Add support for JDBC executeBatch (elastic#747)

noctisma · Jul 25, 2019 · 113167e · 113167e
1 parent b796ce2
commit 113167e
Show file tree

Hide file tree

Showing 8 changed files with 182 additions and 177 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@
 ## Bug Fixes
  * Some JMS Consumers and Producers are filtered due to class name filtering in instrumentation matching
  * Jetty: When no display name is set and context path is "/" transaction service names will now correctly fall back to configured values
+ * JDBC's `executeBatch` is not traced
  * Drops non-String labels when connected to APM Server < 6.7 to avoid validation errors (#687)
 
 # 1.7.0

diff --git a/...ns/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/ConnectionInstrumentation.java b/...ns/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/ConnectionInstrumentation.java
@@ -26,14 +26,13 @@
 
 import co.elastic.apm.agent.bci.ElasticApmInstrumentation;
 import co.elastic.apm.agent.bci.VisibleForAdvice;
-import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
+import co.elastic.apm.agent.jdbc.helper.JdbcHelper;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.NamedElement;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
-import javax.annotation.Nullable;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.util.Collection;
@@ -55,28 +54,12 @@
  */
 public class ConnectionInstrumentation extends ElasticApmInstrumentation {
 
-    @VisibleForAdvice
-    public static final WeakConcurrentMap<Object, String> statementSqlMap = new WeakConcurrentMap<Object, String>(true);
     static final String JDBC_INSTRUMENTATION_GROUP = "jdbc";
 
     @VisibleForAdvice
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void storeSql(@Advice.Return final PreparedStatement statement, @Advice.Argument(0) String sql) {
-        statementSqlMap.putIfAbsent(statement, sql);
-    }
-
-    /**
-     * Returns the SQL statement belonging to provided {@link PreparedStatement}.
-     * <p>
-     * Might return {@code null} when the provided {@link PreparedStatement} is a wrapper of the actual statement.
-     * </p>
-     *
-     * @return the SQL statement belonging to provided {@link PreparedStatement}, or {@code null}
-     */
-    @Nullable
-    @VisibleForAdvice
-    public static String getSqlForStatement(Object statement) {
-        return statementSqlMap.get(statement);
+        JdbcHelper.mapStatementToSql(statement, sql);
     }
 
     @Override

diff --git a/...jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/PreparedStatementInstrumentation.java b/...jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/PreparedStatementInstrumentation.java
diff --git a/...ins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/StatementInstrumentation.java b/...ins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/StatementInstrumentation.java
@@ -51,45 +51,28 @@
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 /**
  * Creates spans for JDBC calls
  */
-public class StatementInstrumentation extends ElasticApmInstrumentation {
+public abstract class StatementInstrumentation extends ElasticApmInstrumentation {
 
+    @SuppressWarnings("WeakerAccess")
     @Nullable
     @VisibleForAdvice
     public static HelperClassManager<JdbcHelper> jdbcHelperManager;
 
-    public StatementInstrumentation(ElasticApmTracer tracer) {
-        jdbcHelperManager = HelperClassManager.ForSingleClassLoader.of(tracer, "co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl",
+    private final ElementMatcher<? super MethodDescription> methodMatcher;
+
+    StatementInstrumentation(ElasticApmTracer tracer, ElementMatcher<? super MethodDescription> methodMatcher) {
+        this.methodMatcher = methodMatcher;
+        jdbcHelperManager = HelperClassManager.ForSingleClassLoader.of(tracer,
+            "co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl",
             "co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl$1",
             "co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl$ConnectionMetaData");
     }
 
-    @Nullable
-    @VisibleForAdvice
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static Span onBeforeExecute(@Advice.This Statement statement, @Advice.Argument(0) String sql) throws SQLException {
-        if (tracer != null && jdbcHelperManager != null) {
-            JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
-            if (helperImpl != null) {
-                return helperImpl.createJdbcSpan(sql, statement.getConnection(), tracer.getActive(), false);
-            }
-        }
-        return null;
-    }
-
-    @VisibleForAdvice
-    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void onAfterExecute(@Advice.Enter @Nullable Span span, @Advice.Thrown Throwable t) {
-        if (span != null) {
-            span.captureException(t)
-                .deactivate()
-                .end();
-        }
-    }
-
     @Override
     public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
         return nameContains("Statement");
@@ -103,13 +86,98 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
 
     @Override
     public ElementMatcher<? super MethodDescription> getMethodMatcher() {
-        return nameStartsWith("execute")
-            .and(takesArgument(0, String.class))
-            .and(isPublic());
+        return methodMatcher;
     }
 
     @Override
     public Collection<String> getInstrumentationGroupNames() {
         return Collections.singleton(JDBC_INSTRUMENTATION_GROUP);
     }
+
+    public static class ExecuteWithQueryInstrumentation extends StatementInstrumentation {
+
+        public ExecuteWithQueryInstrumentation(ElasticApmTracer tracer) {
+            super(tracer,
+                nameStartsWith("execute")
+                    .and(takesArgument(0, String.class))
+                    .and(isPublic())
+            );
+        }
+
+        @Nullable
+        @VisibleForAdvice
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static Span onBeforeExecute(@Advice.This Statement statement, @Advice.Argument(0) String sql) throws SQLException {
+            if (tracer != null && jdbcHelperManager != null) {
+                JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
+                if (helperImpl != null) {
+                    return helperImpl.createJdbcSpan(sql, statement.getConnection(), tracer.getActive(), false);
+                }
+            }
+            return null;
+        }
+
+        @VisibleForAdvice
+        @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+        public static void onAfterExecute(@Advice.Enter @Nullable Span span, @Advice.Thrown Throwable t) {
+            if (span != null) {
+                span.captureException(t)
+                    .deactivate()
+                    .end();
+            }
+        }
+    }
+
+    public static class AddBatchInstrumentation extends StatementInstrumentation {
+
+        public AddBatchInstrumentation(ElasticApmTracer tracer) {
+            super(tracer,
+                nameStartsWith("addBatch")
+                    .and(takesArgument(0, String.class))
+                    .and(isPublic())
+            );
+        }
+
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void storeSql(@Advice.This Statement statement, @Advice.Argument(0) String sql) {
+            if (jdbcHelperManager != null) {
+                JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
+                if (helperImpl != null) {
+                    helperImpl.mapStatementToSql(statement, sql);
+                }
+            }
+        }
+    }
+
+    public static class ExecuteWithoutQueryInstrumentation extends StatementInstrumentation {
+        public ExecuteWithoutQueryInstrumentation(ElasticApmTracer tracer) {
+            super(tracer,
+                nameStartsWith("execute")
+                    .and(takesArguments(0))
+                    .and(isPublic())
+            );
+        }
+
+        @Nullable
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static Span onBeforeExecute(@Advice.This Statement statement) throws SQLException {
+            if (tracer != null && jdbcHelperManager != null) {
+                JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
+                if (helperImpl != null) {
+                    final @Nullable String sql = helperImpl.retrieveSqlForStatement(statement);
+                    return helperImpl.createJdbcSpan(sql, statement.getConnection(), tracer.getActive(), true);
+                }
+            }
+            return null;
+        }
+
+        @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+        public static void onAfterExecute(@Advice.Enter @Nullable Span span, @Advice.Thrown Throwable t) {
+            if (span != null) {
+                span.captureException(t)
+                    .deactivate()
+                    .end();
+            }
+        }
+    }
 }
diff --git a/...nt-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/helper/JdbcHelper.java b/...nt-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/helper/JdbcHelper.java
@@ -24,13 +24,43 @@
  */
 package co.elastic.apm.agent.jdbc.helper;
 
+import co.elastic.apm.agent.bci.VisibleForAdvice;
 import co.elastic.apm.agent.impl.transaction.Span;
 import co.elastic.apm.agent.impl.transaction.TraceContextHolder;
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 
 import javax.annotation.Nullable;
 import java.sql.Connection;
 
-public interface JdbcHelper {
+public abstract class JdbcHelper {
+    @SuppressWarnings("WeakerAccess")
+    @VisibleForAdvice
+    public static final WeakConcurrentMap<Object, String> statementSqlMap = new WeakConcurrentMap<>(true);
+
+    /**
+     * Maps the provided sql to the provided Statement object
+     *
+     * @param statement javax.sql.Statement object
+     * @param sql       query string
+     */
+    public static void mapStatementToSql(Object statement, String sql) {
+        statementSqlMap.putIfAbsent(statement, sql);
+    }
+
+    /**
+     * Returns the SQL statement belonging to provided Statement.
+     * <p>
+     * Might return {@code null} when the provided Statement is a wrapper of the actual statement.
+     * </p>
+     *
+     * @return the SQL statement belonging to provided Statement, or {@code null}
+     */
+    @Nullable
+    public static String retrieveSqlForStatement(Object statement) {
+        return statementSqlMap.get(statement);
+    }
+
+
     @Nullable
-    Span createJdbcSpan(@Nullable String sql, Connection connection, @Nullable TraceContextHolder<?> parent, boolean preparedStatement);
+    public abstract Span createJdbcSpan(@Nullable String sql, Connection connection, @Nullable TraceContextHolder<?> parent, boolean preparedStatement);
 }
diff --git a/...lugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/helper/JdbcHelperImpl.java b/...lugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/helper/JdbcHelperImpl.java
@@ -37,7 +37,7 @@
 import java.sql.DatabaseMetaData;
 import java.sql.SQLException;
 
-public class JdbcHelperImpl implements JdbcHelper {
+public class JdbcHelperImpl extends JdbcHelper {
     public static final String DB_SPAN_TYPE = "db";
     public static final String DB_SPAN_ACTION = "query";
 

diff --git a/...n/src/main/resources/META-INF/services/co.elastic.apm.agent.bci.ElasticApmInstrumentation b/...n/src/main/resources/META-INF/services/co.elastic.apm.agent.bci.ElasticApmInstrumentation
@@ -1,3 +1,4 @@
 co.elastic.apm.agent.jdbc.ConnectionInstrumentation
-co.elastic.apm.agent.jdbc.StatementInstrumentation
-co.elastic.apm.agent.jdbc.PreparedStatementInstrumentation
+co.elastic.apm.agent.jdbc.StatementInstrumentation$ExecuteWithQueryInstrumentation
+co.elastic.apm.agent.jdbc.StatementInstrumentation$AddBatchInstrumentation
+co.elastic.apm.agent.jdbc.StatementInstrumentation$ExecuteWithoutQueryInstrumentation