Skip to content

Commit b673261

Browse files
deejgregordeejgregor
committed
HikariCP: Add a span when waiting on the pool
1 parent 7d20566 commit b673261

File tree

6 files changed

+398
-0
lines changed

6 files changed

+398
-0
lines changed

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package datadog.trace.instrumentation.jdbc;
2+
3+
/**
4+
* Shared blocked getConnection() tracking ThreadLocking for Hikari.
5+
*/
6+
public class HikariBlockedTracker {
7+
private static final ThreadLocal<Boolean> tracker = ThreadLocal.withInitial(() -> false);
8+
9+
public static void clearBlocked() {
10+
tracker.set(false);
11+
}
12+
13+
public static void setBlocked() {
14+
tracker.set(true);
15+
}
16+
public static boolean wasBlocked() {
17+
return tracker.get();
18+
}
19+
}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package datadog.trace.instrumentation.jdbc;
2+
3+
import java.util.concurrent.SynchronousQueue;
4+
import java.util.concurrent.TimeUnit;
5+
6+
/**
7+
* Blocked getConnection() tracking for Hikari starting with commit f0b3c520c.
8+
*/
9+
public class HikariBlockedTrackingSynchronousQueue<T> extends SynchronousQueue<T> {
10+
public HikariBlockedTrackingSynchronousQueue() {
11+
// This assumes the initialization of the SynchronousQueue in ConcurrentBag doesn't change
12+
super(true);
13+
}
14+
15+
@Override
16+
public T poll(long timeout, TimeUnit unit) throws InterruptedException {
17+
HikariBlockedTracker.setBlocked();
18+
return super.poll(timeout, unit);
19+
}
20+
}

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

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
package datadog.trace.instrumentation.jdbc;
2+
3+
import com.google.auto.service.AutoService;
4+
import com.zaxxer.hikari.pool.HikariPool;
5+
import com.zaxxer.hikari.util.ConcurrentBag;
6+
import datadog.trace.agent.tooling.Instrumenter;
7+
import datadog.trace.agent.tooling.InstrumenterModule;
8+
import datadog.trace.bootstrap.InstrumentationContext;
9+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
10+
import net.bytebuddy.asm.Advice;
11+
import org.slf4j.Logger;
12+
import org.slf4j.LoggerFactory;
13+
14+
import java.lang.reflect.Field;
15+
import java.util.Map;
16+
import java.util.concurrent.TimeUnit;
17+
18+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
19+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
20+
import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_POOL_NAME;
21+
import static java.util.Collections.singletonMap;
22+
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
23+
24+
/**
25+
* Instrument Hikari's ConcurrentBag class to detect when blocking occurs trying to get
26+
* an entry from the connection pool.
27+
*/
28+
@AutoService(InstrumenterModule.class)
29+
public final class HikariConcurrentBagInstrumentation extends InstrumenterModule.Tracing
30+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
31+
32+
private static final Logger log = LoggerFactory.getLogger(HikariConcurrentBagInstrumentation.class);
33+
34+
public HikariConcurrentBagInstrumentation() {
35+
super("jdbc-datasource");
36+
}
37+
38+
@Override
39+
public String instrumentedType() {
40+
return "com.zaxxer.hikari.util.ConcurrentBag";
41+
}
42+
43+
@Override
44+
public String[] helperClassNames() {
45+
return new String[]{packageName + ".HikariBlockedTrackingSynchronousQueue",
46+
packageName + ".HikariBlockedTracker"};
47+
}
48+
49+
@Override
50+
public Map<String, String> contextStore() {
51+
// For getting the poolName
52+
return singletonMap("com.zaxxer.hikari.util.ConcurrentBag", String.class.getName());
53+
}
54+
55+
@Override
56+
public void methodAdvice(MethodTransformer transformer) {
57+
transformer.applyAdvice(
58+
isConstructor(), HikariConcurrentBagInstrumentation.class.getName() + "$ConstructorAdvice");
59+
transformer.applyAdvice(
60+
named("borrow"),
61+
HikariConcurrentBagInstrumentation.class.getName() + "$BorrowAdvice");
62+
}
63+
64+
public static class ConstructorAdvice {
65+
@Advice.OnMethodExit(suppress = Throwable.class)
66+
static void after(@Advice.This ConcurrentBag<?> thiz)
67+
throws IllegalAccessException, NoSuchFieldException {
68+
try {
69+
Field handoffQueueField = thiz.getClass().getDeclaredField("handoffQueue");
70+
handoffQueueField.setAccessible(true);
71+
handoffQueueField.set(thiz, new HikariBlockedTrackingSynchronousQueue<>());
72+
} catch (NoSuchFieldException e) {
73+
// ignore -- see HikariQueuedSequenceSynchronizerInstrumentation for older Hikari versions
74+
}
75+
76+
Field hikariPoolField = thiz.getClass().getDeclaredField("listener");
77+
hikariPoolField.setAccessible(true);
78+
HikariPool hikariPool = (HikariPool) hikariPoolField.get(thiz);
79+
80+
/*
81+
* In earlier versions of Hikari, poolName is directly inside HikariPool, and
82+
* in later versions it is in the PoolBase superclass.
83+
*/
84+
final Class<?> hikariPoolSuper = hikariPool.getClass().getSuperclass();
85+
final Class<?> poolNameContainingClass;
86+
if (!hikariPoolSuper.getName().equals("java.lang.Object")) {
87+
poolNameContainingClass = hikariPoolSuper;
88+
} else {
89+
poolNameContainingClass = hikariPool.getClass();
90+
}
91+
Field poolNameField = poolNameContainingClass.getDeclaredField("poolName");
92+
poolNameField.setAccessible(true);
93+
String poolName = (String) poolNameField.get(hikariPool);
94+
InstrumentationContext.get(ConcurrentBag.class, String.class).put(thiz, poolName);
95+
}
96+
}
97+
98+
public static class BorrowAdvice {
99+
private static final String POOL_WAITING = "pool.waiting";
100+
101+
@Advice.OnMethodEnter(suppress = Throwable.class)
102+
public static Long onEnter() {
103+
HikariBlockedTracker.clearBlocked();
104+
return System.currentTimeMillis();
105+
}
106+
107+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
108+
public static void stopSpan(@Advice.This ConcurrentBag thiz,
109+
@Advice.Enter final Long startTimeMillis, @Advice.Thrown final Throwable throwable) {
110+
if (HikariBlockedTracker.wasBlocked()) {
111+
final AgentSpan span = startSpan("hikari", POOL_WAITING, TimeUnit.MILLISECONDS.toMicros(startTimeMillis));
112+
final String poolName = InstrumentationContext.get(ConcurrentBag.class, String.class).get(thiz);
113+
if (poolName != null) {
114+
span.setTag(DB_POOL_NAME, poolName);
115+
}
116+
// XXX should we do anything with the throwable?
117+
span.finish();
118+
}
119+
HikariBlockedTracker.clearBlocked();
120+
}
121+
}
122+
}

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package datadog.trace.instrumentation.jdbc;
2+
3+
import com.google.auto.service.AutoService;
4+
import datadog.trace.agent.tooling.Instrumenter;
5+
import datadog.trace.agent.tooling.InstrumenterModule;
6+
import net.bytebuddy.asm.Advice;
7+
import org.slf4j.Logger;
8+
import org.slf4j.LoggerFactory;
9+
10+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
11+
12+
/**
13+
* Blocked getConnection() tracking for Hikari starting before commit f0b3c520c.
14+
*/
15+
@AutoService(InstrumenterModule.class)
16+
public final class HikariQueuedSequenceSynchronizerInstrumentation extends InstrumenterModule.Tracing
17+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
18+
19+
private static final Logger log = LoggerFactory.getLogger(HikariQueuedSequenceSynchronizerInstrumentation.class);
20+
21+
public HikariQueuedSequenceSynchronizerInstrumentation() {
22+
super("jdbc-datasource");
23+
}
24+
25+
@Override
26+
public String instrumentedType() {
27+
return "com.zaxxer.hikari.util.QueuedSequenceSynchronizer";
28+
}
29+
30+
@Override
31+
public void methodAdvice(MethodTransformer transformer) {
32+
transformer.applyAdvice(
33+
named("waitUntilSequenceExceeded"),
34+
HikariQueuedSequenceSynchronizerInstrumentation.class.getName()
35+
+ "$WaitUntilSequenceExceededAdvice");
36+
}
37+
38+
public static class WaitUntilSequenceExceededAdvice {
39+
@Advice.OnMethodEnter(suppress = Throwable.class)
40+
public static void onEnter() {
41+
HikariBlockedTracker.setBlocked();
42+
}
43+
}
44+
}

‎dd-java-agent/instrumentation/jdbc/src/test/groovy/SaturatedPoolBlockingTest.groovy

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import datadog.trace.agent.test.AgentTestRunner
2+
import com.zaxxer.hikari.HikariConfig
3+
import com.zaxxer.hikari.HikariDataSource
4+
import test.TestDataSource
5+
6+
import java.sql.SQLTimeoutException
7+
import java.sql.SQLTransientConnectionException
8+
9+
/**
10+
* Ideas taken from Hikari's com.zaxxer.hikari.pool.TestSaturatedPool830.
11+
*/
12+
class SaturatedPoolBlockingTest extends AgentTestRunner {
13+
def "saturated pool test"(int connectionTimeout, Long exhaustPoolForMillis, int expectedWaitingSpans, boolean expectedTimeout) {
14+
setup:
15+
TEST_WRITER.setFilter((trace) -> trace.get(0).getOperationName() == "test.when")
16+
17+
final HikariConfig config = new HikariConfig()
18+
config.setPoolName("testPool")
19+
config.setMaximumPoolSize(1)
20+
config.setConnectionTimeout(connectionTimeout)
21+
config.setDataSourceClassName(TestDataSource.class.getName())
22+
final HikariDataSource ds = new HikariDataSource(config)
23+
24+
when:
25+
if (exhaustPoolForMillis != null) {
26+
def saturatedConnection = ds.getConnection()
27+
new Thread(() -> {
28+
Thread.sleep(exhaustPoolForMillis)
29+
saturatedConnection.close()
30+
}, "saturated connection closer").start()
31+
}
32+
33+
def timedOut = false
34+
def span = TEST_TRACER.startSpan("test", "test.when")
35+
try (def ignore = TEST_TRACER.activateSpan(span)) {
36+
def connection = ds.getConnection()
37+
connection.close()
38+
} catch (SQLTransientConnectionException e) {
39+
if (e.getMessage().contains("request timed out after")) { // Hikari, newer
40+
timedOut = true
41+
} else {
42+
throw e
43+
}
44+
} catch (SQLTimeoutException ignored) { // Hikari, older
45+
timedOut = true
46+
}
47+
span.finish()
48+
49+
then:
50+
def waiting = TEST_WRITER.firstTrace().findAll {
51+
element -> element.getOperationName() == "pool.waiting" }
52+
53+
print(TEST_WRITER.firstTrace())
54+
55+
verifyAll {
56+
TEST_WRITER.size() == 1
57+
waiting.size() == expectedWaitingSpans
58+
timedOut == expectedTimeout
59+
}
60+
61+
where:
62+
connectionTimeout | exhaustPoolForMillis | expectedWaitingSpans | expectedTimeout
63+
1000 | null | 0 | false
64+
1000 | null | 0 | false
65+
1000 | 500 | 1 | false
66+
1000 | 1500 | 1 | true
67+
}
68+
}

0 commit comments

Comments
 (0)