Cleanup Hikari instrumentation

- Comment on the overall structure and highlight related instrumentations.
- Avoid using reflection and breakup HikariConcurrentBagInstrumentation.ConstructorAdvice
  into separate, simpler instrumentations: HikariPoolInstrumentation and
  HikariConcurrentBagHandoffQueueInstrumentation.
- Move HikariBlockedTrackingSynchronousQueue into HikariConcurrentBagHandoffQueueInstrumentation.
- Make sure the instrumentationName is consistently "jdbc".

Author: deejgregor

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/HikariBlockedTracker.java

@@ -1,6 +1,6 @@
 package datadog.trace.instrumentation.jdbc;
 
-/** Shared blocked getConnection() tracking ThreadLocking for Hikari. */
+/** Shared blocked getConnection() tracking {@link ThreadLocal} for Hikari. */
 public class HikariBlockedTracker {
   private static final ThreadLocal<Boolean> tracker = ThreadLocal.withInitial(() -> false);
 

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/HikariBlockedTrackingSynchronousQueue.java

@@ -0,0 +1,83 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresField;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.TimeUnit;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * Detect blocking for newer Hikari versions starting with commit f0b3c520c (>=2.6.0) by looking for
+ * calls to <code>handoffQueue.poll(timeout, NANOSECONDS)</code>.
+ */
+@AutoService(InstrumenterModule.class)
+public final class HikariConcurrentBagHandoffQueueInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType,
+        Instrumenter.HasMethodAdvice,
+        Instrumenter.WithTypeStructure {
+
+  public HikariConcurrentBagHandoffQueueInstrumentation() {
+    super("jdbc");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "com.zaxxer.hikari.util.ConcurrentBag";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> structureMatcher() {
+    return declaresField(named("handoffQueue"));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".HikariBlockedTracker",
+      packageName
+          + ".HikariConcurrentBagHandoffQueueInstrumentation$BlockedTrackingSynchronousQueue",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor(),
+        HikariConcurrentBagHandoffQueueInstrumentation.class.getName() + "$ConstructorAdvice");
+  }
+
+  public static class ConstructorAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    static void after(
+        @Advice.FieldValue(value = "handoffQueue", readOnly = false)
+            SynchronousQueue handoffQueue) {
+      handoffQueue = new BlockedTrackingSynchronousQueue<>();
+    }
+  }
+
+  public static class BlockedTrackingSynchronousQueue<T> extends SynchronousQueue<T> {
+    public BlockedTrackingSynchronousQueue() {
+      // This assumes the initialization of the SynchronousQueue in ConcurrentBag doesn't change
+      super(true);
+    }
+
+    @Override
+    public T poll(long timeout, TimeUnit unit) throws InterruptedException {
+      HikariBlockedTracker.setBlocked();
+      return super.poll(timeout, unit);
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/HikariConcurrentBagHandoffQueueInstrumentation.java

@@ -6,24 +6,32 @@
 import static datadog.trace.instrumentation.jdbc.PoolWaitingDecorator.DECORATE;
 import static datadog.trace.instrumentation.jdbc.PoolWaitingDecorator.POOL_WAITING;
 import static java.util.Collections.singletonMap;
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 
 import com.google.auto.service.AutoService;
-import com.zaxxer.hikari.pool.HikariPool;
 import com.zaxxer.hikari.util.ConcurrentBag;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import java.lang.reflect.Field;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import net.bytebuddy.asm.Advice;
 
 /**
  * Instrument Hikari's ConcurrentBag class to detect when blocking occurs trying to get an entry
- * from the connection pool.
+ * from the connection pool. There are two related instrumentations to detect blocking for different
+ * versions of Hikari. The {@link HikariBlockedTracker} contextStore is used to pass blocking state
+ * from the other instrumentations to this class.
+ *
+ * <ul>
+ *   <li>Before commit f0b3c520c (2.4.0 <= version < 2.6.0): calls to <code>
+ *       synchronizer.waitUntilSequenceExceeded(startSeq, timeout)</code> with {@link
+ *       HikariQueuedSequenceSynchronizerInstrumentation}
+ *   <li>Commit f0b3c520c and later (version >= 2.6.0): calls to <code>
+ *       handoffQueue.poll(timeout, NANOSECONDS)</code> with {@link
+ *       HikariQueuedSequenceSynchronizerInstrumentation}
+ * </ul>
  */
 @AutoService(InstrumenterModule.class)
 public final class HikariConcurrentBagInstrumentation extends InstrumenterModule.Tracing
@@ -46,60 +54,26 @@ public String instrumentedType() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".HikariBlockedTrackingSynchronousQueue",
-      packageName + ".HikariBlockedTracker",
-      packageName + ".PoolWaitingDecorator"
+      packageName + ".HikariBlockedTracker", packageName + ".PoolWaitingDecorator"
     };
   }
 
   @Override
   public Map<String, String> contextStore() {
-    // For getting the poolName
+    // The contextStore Map is populated by HikariPoolInstrumentation
     return singletonMap("com.zaxxer.hikari.util.ConcurrentBag", String.class.getName());
   }
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(
-        isConstructor(), HikariConcurrentBagInstrumentation.class.getName() + "$ConstructorAdvice");
     transformer.applyAdvice(
         named("borrow"), HikariConcurrentBagInstrumentation.class.getName() + "$BorrowAdvice");
   }
 
-  public static class ConstructorAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    static void after(@Advice.This ConcurrentBag<?> thiz)
-        throws IllegalAccessException, NoSuchFieldException {
-      try {
-        Field handoffQueueField = thiz.getClass().getDeclaredField("handoffQueue");
-        handoffQueueField.setAccessible(true);
-        handoffQueueField.set(thiz, new HikariBlockedTrackingSynchronousQueue<>());
-      } catch (NoSuchFieldException e) {
-        // ignore -- see HikariQueuedSequenceSynchronizerInstrumentation for older Hikari versions
-      }
-
-      Field hikariPoolField = thiz.getClass().getDeclaredField("listener");
-      hikariPoolField.setAccessible(true);
-      HikariPool hikariPool = (HikariPool) hikariPoolField.get(thiz);
-
-      /*
-       * In earlier versions of Hikari, poolName is directly inside HikariPool, and
-       * in later versions it is in the PoolBase superclass.
-       */
-      final Class<?> hikariPoolSuper = hikariPool.getClass().getSuperclass();
-      final Class<?> poolNameContainingClass;
-      if (!hikariPoolSuper.getName().equals("java.lang.Object")) {
-        poolNameContainingClass = hikariPoolSuper;
-      } else {
-        poolNameContainingClass = hikariPool.getClass();
-      }
-      Field poolNameField = poolNameContainingClass.getDeclaredField("poolName");
-      poolNameField.setAccessible(true);
-      String poolName = (String) poolNameField.get(hikariPool);
-      InstrumentationContext.get(ConcurrentBag.class, String.class).put(thiz, poolName);
-    }
-  }
-
+  /**
+   * Instead of always starting and ending a span, a pool.waiting span is only created if blocking
+   * is detected when attempting to get a connection from the pool.
+   */
   public static class BorrowAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static Long onEnter() {

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/HikariConcurrentBagInstrumentation.java

@@ -0,0 +1,56 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import com.zaxxer.hikari.util.ConcurrentBag;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.InstrumentationContext;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+
+/**
+ * Store the poolName associated with a {@link ConcurrentBag} for later use in {@link
+ * HikariConcurrentBagInstrumentation}.
+ */
+@AutoService(InstrumenterModule.class)
+public final class HikariPoolInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public HikariPoolInstrumentation() {
+    super("jdbc");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "com.zaxxer.hikari.pool.HikariPool";
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("com.zaxxer.hikari.util.ConcurrentBag", String.class.getName());
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor(), HikariPoolInstrumentation.class.getName() + "$ConstructorAdvice");
+  }
+
+  public static class ConstructorAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    static void after(
+        @Advice.FieldValue("connectionBag") ConcurrentBag concurrentBag,
+        @Advice.FieldValue("poolName") String poolName) {
+      InstrumentationContext.get(ConcurrentBag.class, String.class).put(concurrentBag, poolName);
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/HikariPoolInstrumentation.java

@@ -7,21 +7,31 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import net.bytebuddy.asm.Advice;
 
-/** Blocked getConnection() tracking for Hikari starting before commit f0b3c520c. */
+/**
+ * Detect blocking for older Hikari versions before commit f0b3c520c (<2.6.0) by looking for calls
+ * to <code>synchronizer.waitUntilSequenceExceeded(startSeq, timeout)</code>.
+ */
 @AutoService(InstrumenterModule.class)
 public final class HikariQueuedSequenceSynchronizerInstrumentation
     extends InstrumenterModule.Tracing
     implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
 
   public HikariQueuedSequenceSynchronizerInstrumentation() {
-    super("jdbc-datasource");
+    super("jdbc");
   }
 
   @Override
   public String instrumentedType() {
     return "com.zaxxer.hikari.util.QueuedSequenceSynchronizer";
   }
 
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".HikariBlockedTracker",
+    };
+  }
+
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(