Skip to content

Commit 9b642bd

Browse files
amarzialiamarziali
authored
Support proper async propagation for GAX retrying futures (#11629)
Support proper async propagation for GAX retrying futures prevent possible race more tests Add versioned alias Co-authored-by: andrea.marziali <andrea.marziali@datadoghq.com>
1 parent 677e9da commit 9b642bd

7 files changed

Lines changed: 316 additions & 5 deletions

File tree

‎dd-java-agent/instrumentation/gax-1.4/build.gradle‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
muzzle {
2+
pass {
3+
group = "com.google.api"
4+
module = "gax"
5+
// CallbackChainRetryingFuture (with setAttemptFuture + attemptFutureCompletionListener) first
6+
// appears in gax 1.4.0. No need to assert the low bound hence since the instrumentation is not applied.
7+
versions = "[1.4.0,]"
8+
}
9+
}
10+
11+
apply from: "$rootDir/gradle/java.gradle"
12+
13+
addTestSuiteForDir('latestDepTest', 'test')
14+
15+
dependencies {
16+
compileOnly group: 'com.google.api', name: 'gax', version: '1.4.0'
17+
testImplementation project(":dd-java-agent:instrumentation:guava-10.0")
18+
testImplementation group: 'com.google.api', name: 'gax', version: '1.4.0'
19+
20+
latestDepTestImplementation group: 'com.google.api', name: 'gax', version: '+'
21+
}

‎dd-java-agent/instrumentation/gax-1.4/src/main/java/datadog/trace/instrumentation/gax/CallbackChainRetryingFutureInstrumentation.java‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package datadog.trace.instrumentation.gax;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static java.util.Collections.singletonMap;
5+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
7+
8+
import com.google.auto.service.AutoService;
9+
import datadog.trace.agent.tooling.Instrumenter;
10+
import datadog.trace.agent.tooling.InstrumenterModule;
11+
import datadog.trace.bootstrap.ContextStore;
12+
import datadog.trace.bootstrap.InstrumentationContext;
13+
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
14+
import java.util.Map;
15+
import net.bytebuddy.asm.Advice;
16+
17+
/**
18+
* Cancels the trace continuation captured on a gax retry attempt's completion listener when that
19+
* listener is superseded.
20+
*
21+
* <p>{@code AttemptCallable.call()} calls {@code setAttemptFuture} twice per attempt: first with a
22+
* {@code NonCancellableFuture} placeholder (a reservation that is never completed), then with the
23+
* real RPC future. {@code CallbackChainRetryingFuture.setAttemptFuture} registers a fresh {@code
24+
* AttemptCompletionListener} each time and stores it in {@code attemptFutureCompletionListener}.
25+
* The guava {@code AbstractFuture} instrumentation captures a continuation onto each listener. The
26+
* placeholder's listener never runs (its future never completes), so its continuation would never
27+
* be activated or canceled.
28+
*
29+
* <p>Each {@code setAttemptFuture} overwrites the previous listener: that overwrite is the
30+
* abandonment signal. We capture the previous listener on entry but cancel its continuation on
31+
* exit, once the field has actually been replaced, so a listener GAX still treats as active (early
32+
* return, or a completion racing the overwrite) keeps its continuation.
33+
*/
34+
@AutoService(InstrumenterModule.class)
35+
public class CallbackChainRetryingFutureInstrumentation extends InstrumenterModule.ContextTracking
36+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
37+
38+
public CallbackChainRetryingFutureInstrumentation() {
39+
super("gax", "gax-1.4");
40+
}
41+
42+
@Override
43+
public String instrumentedType() {
44+
return "com.google.api.gax.retrying.CallbackChainRetryingFuture";
45+
}
46+
47+
@Override
48+
public Map<String, String> contextStore() {
49+
return singletonMap(Runnable.class.getName(), State.class.getName());
50+
}
51+
52+
@Override
53+
public void methodAdvice(MethodTransformer transformer) {
54+
transformer.applyAdvice(
55+
named("setAttemptFuture")
56+
.and(takesArguments(1))
57+
.and(takesArgument(0, named("com.google.api.core.ApiFuture"))),
58+
CallbackChainRetryingFutureInstrumentation.class.getName() + "$SetAttemptFutureAdvice");
59+
}
60+
61+
public static class SetAttemptFutureAdvice {
62+
@Advice.OnMethodEnter(suppress = Throwable.class)
63+
public static Runnable capturePrevious(
64+
@Advice.FieldValue("attemptFutureCompletionListener") final Runnable previousListener) {
65+
return previousListener;
66+
}
67+
68+
@Advice.OnMethodExit(suppress = Throwable.class)
69+
public static void cancelSuperseded(
70+
@Advice.Enter final Runnable previousListener,
71+
@Advice.FieldValue("attemptFutureCompletionListener") final Runnable newListener) {
72+
// Only cancel once the field has actually been replaced: GAX may return early without
73+
// reassigning, and a listener still treated as active must keep its continuation.
74+
if (previousListener != null && previousListener != newListener) {
75+
final ContextStore<Runnable, State> contextStore =
76+
InstrumentationContext.get(Runnable.class, State.class);
77+
final State state = contextStore.remove(previousListener);
78+
if (state != null) {
79+
state.closeContinuation();
80+
}
81+
}
82+
}
83+
}
84+
}

‎dd-java-agent/instrumentation/gax-1.4/src/test/java/datadog/trace/instrumentation/gax/GaxRetryContinuationTest.java‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
package datadog.trace.instrumentation.gax;
2+
3+
import static datadog.trace.agent.test.assertions.SpanMatcher.span;
4+
import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME;
5+
import static datadog.trace.agent.test.assertions.TraceMatcher.trace;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
7+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
8+
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertThrows;
10+
11+
import com.google.api.core.ApiClock;
12+
import com.google.api.core.NanoClock;
13+
import com.google.api.core.SettableApiFuture;
14+
import com.google.api.gax.retrying.BasicResultRetryAlgorithm;
15+
import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
16+
import com.google.api.gax.retrying.RetryAlgorithm;
17+
import com.google.api.gax.retrying.RetrySettings;
18+
import com.google.api.gax.retrying.RetryingFuture;
19+
import com.google.api.gax.retrying.ScheduledRetryingExecutor;
20+
import datadog.trace.agent.test.AbstractInstrumentationTest;
21+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
22+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
23+
import java.util.concurrent.CountDownLatch;
24+
import java.util.concurrent.ExecutionException;
25+
import java.util.concurrent.Executors;
26+
import java.util.concurrent.ScheduledExecutorService;
27+
import java.util.concurrent.TimeUnit;
28+
import java.util.concurrent.atomic.AtomicInteger;
29+
import org.junit.jupiter.api.Test;
30+
import org.threeten.bp.Duration;
31+
32+
class GaxRetryContinuationTest extends AbstractInstrumentationTest {
33+
34+
@Test
35+
void supersededAttemptListenerDoesNotLeak() throws Exception {
36+
ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
37+
try {
38+
ScheduledRetryingExecutor<String> executor =
39+
new ScheduledRetryingExecutor<>(retryAlgorithm(1), scheduler);
40+
RetryingFuture<String> retryingFuture = executor.createFuture(() -> "ok");
41+
42+
AgentSpan span = startSpan("gax", "publish");
43+
try (AgentScope scope = activateSpan(span)) {
44+
// reserve the attempt slot with a placeholder that is never completed
45+
SettableApiFuture<String> placeholder = SettableApiFuture.create();
46+
retryingFuture.setAttemptFuture(placeholder);
47+
// supersede it with the real attempt future -> must cancel the placeholder's continuation
48+
SettableApiFuture<String> attempt = SettableApiFuture.create();
49+
retryingFuture.setAttemptFuture(attempt);
50+
// real attempt completes -> its listener runs -> resolves normally
51+
attempt.set("ok");
52+
}
53+
span.finish();
54+
// one trace that will be dropped if the placeholder's continuation is never cancelled
55+
assertTraces(trace(span().root().operationName("publish")));
56+
} finally {
57+
scheduler.shutdownNow();
58+
}
59+
}
60+
61+
@Test
62+
void singleAttemptSucceedsAndContextNotLeaked() throws Exception {
63+
ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
64+
try {
65+
ScheduledRetryingExecutor<String> executor =
66+
new ScheduledRetryingExecutor<>(retryAlgorithm(3), scheduler);
67+
AgentSpan parent = startSpan("gax", "publish");
68+
CountDownLatch allAttemptsDone = new CountDownLatch(1);
69+
RetryingFuture<String> future =
70+
executor.createFuture(
71+
() -> {
72+
AgentSpan attempt = startSpan("gax", "attempt");
73+
try {
74+
return "ok";
75+
} finally {
76+
attempt.finish();
77+
allAttemptsDone.countDown();
78+
}
79+
});
80+
81+
try (AgentScope scope = activateSpan(parent)) {
82+
future.setAttemptFuture(executor.submit(future));
83+
assertEquals("ok", future.get(5, TimeUnit.SECONDS));
84+
}
85+
allAttemptsDone.await(5, TimeUnit.SECONDS);
86+
parent.finish();
87+
assertTraces(
88+
trace(
89+
SORT_BY_START_TIME,
90+
span().root().operationName("publish"),
91+
span().childOf(parent.getSpanId()).operationName("attempt")));
92+
} finally {
93+
scheduler.shutdownNow();
94+
}
95+
}
96+
97+
@Test
98+
void retriedAttemptsSucceedAndContextNotLeaked() throws Exception {
99+
ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
100+
try {
101+
ScheduledRetryingExecutor<String> executor =
102+
new ScheduledRetryingExecutor<>(retryAlgorithm(3), scheduler);
103+
AtomicInteger count = new AtomicInteger(0);
104+
AgentSpan parent = startSpan("gax", "publish");
105+
CountDownLatch allAttemptsDone = new CountDownLatch(3);
106+
RetryingFuture<String> future =
107+
executor.createFuture(
108+
() -> {
109+
AgentSpan attempt = startSpan("gax", "attempt");
110+
try {
111+
if (count.incrementAndGet() < 3) {
112+
throw new RuntimeException("transient");
113+
}
114+
return "ok";
115+
} finally {
116+
attempt.finish();
117+
allAttemptsDone.countDown();
118+
}
119+
});
120+
121+
try (AgentScope scope = activateSpan(parent)) {
122+
future.setAttemptFuture(executor.submit(future));
123+
assertEquals("ok", future.get(5, TimeUnit.SECONDS));
124+
}
125+
allAttemptsDone.await(5, TimeUnit.SECONDS);
126+
parent.finish();
127+
assertTraces(
128+
trace(
129+
SORT_BY_START_TIME,
130+
span().root().operationName("publish"),
131+
span().childOf(parent.getSpanId()).operationName("attempt"),
132+
span().childOf(parent.getSpanId()).operationName("attempt"),
133+
span().childOf(parent.getSpanId()).operationName("attempt")));
134+
} finally {
135+
scheduler.shutdownNow();
136+
}
137+
}
138+
139+
@Test
140+
void exhaustedRetriesContextNotLeaked() throws Exception {
141+
ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
142+
try {
143+
ScheduledRetryingExecutor<String> executor =
144+
new ScheduledRetryingExecutor<>(retryAlgorithm(3), scheduler);
145+
AgentSpan parent = startSpan("gax", "publish");
146+
CountDownLatch allAttemptsDone = new CountDownLatch(3);
147+
RetryingFuture<String> future =
148+
executor.createFuture(
149+
() -> {
150+
AgentSpan attempt = startSpan("gax", "attempt");
151+
try {
152+
throw new RuntimeException("always fails");
153+
} finally {
154+
attempt.finish();
155+
allAttemptsDone.countDown();
156+
}
157+
});
158+
159+
try (AgentScope scope = activateSpan(parent)) {
160+
future.setAttemptFuture(executor.submit(future));
161+
assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS));
162+
}
163+
allAttemptsDone.await(5, TimeUnit.SECONDS);
164+
parent.finish();
165+
assertTraces(
166+
trace(
167+
SORT_BY_START_TIME,
168+
span().root().operationName("publish"),
169+
span().childOf(parent.getSpanId()).operationName("attempt"),
170+
span().childOf(parent.getSpanId()).operationName("attempt"),
171+
span().childOf(parent.getSpanId()).operationName("attempt")));
172+
} finally {
173+
scheduler.shutdownNow();
174+
}
175+
}
176+
177+
private static RetryAlgorithm<String> retryAlgorithm(int maxAttempts) {
178+
RetrySettings settings =
179+
RetrySettings.newBuilder()
180+
.setMaxAttempts(maxAttempts)
181+
.setInitialRetryDelay(Duration.ofMillis(1))
182+
.setRetryDelayMultiplier(1.0)
183+
.setMaxRetryDelay(Duration.ofMillis(10))
184+
.setInitialRpcTimeout(Duration.ofSeconds(5))
185+
.setRpcTimeoutMultiplier(1.0)
186+
.setMaxRpcTimeout(Duration.ofSeconds(5))
187+
.setTotalTimeout(Duration.ofSeconds(30))
188+
.build();
189+
ApiClock clock = NanoClock.getDefaultClock();
190+
return new RetryAlgorithm<>(
191+
new BasicResultRetryAlgorithm<>(), new ExponentialRetryAlgorithm(settings, clock));
192+
}
193+
}

‎dd-java-agent/instrumentation/google-pubsub-1.116/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ dependencies {
1818
testImplementation group: 'com.google.cloud', name: 'google-cloud-pubsub', version: '1.116.0'
1919
testImplementation project(":dd-java-agent:instrumentation:grpc-1.5")
2020
testImplementation project(":dd-java-agent:instrumentation:guava-10.0")
21+
testImplementation project(":dd-java-agent:instrumentation:gax-1.4")
2122
latestDepTestImplementation group: 'com.google.cloud', name: 'google-cloud-pubsub', version: '+'
2223
}
2324

‎dd-java-agent/instrumentation/google-pubsub-1.116/src/test/groovy/PubSubTest.groovy‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,6 @@ abstract class PubSubTest extends VersionedNamingTestBase {
6767
null
6868
}
6969

70-
@Override
71-
boolean useStrictTraceWrites() {
72-
false
73-
}
74-
7570
boolean shadowGrpcSpans() {
7671
true
7772
}

‎metadata/supported-configurations.json‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5777,6 +5777,22 @@
57775777
"aliases": ["DD_TRACE_INTEGRATION_FREEMARKER_ENABLED", "DD_INTEGRATION_FREEMARKER_ENABLED"]
57785778
}
57795779
],
5780+
"DD_TRACE_GAX_1_4_ENABLED": [
5781+
{
5782+
"version": "A",
5783+
"type": "boolean",
5784+
"default": "true",
5785+
"aliases": ["DD_TRACE_INTEGRATION_GAX_1_4_ENABLED", "DD_INTEGRATION_GAX_1_4_ENABLED"]
5786+
}
5787+
],
5788+
"DD_TRACE_GAX_ENABLED": [
5789+
{
5790+
"version": "A",
5791+
"type": "boolean",
5792+
"default": "true",
5793+
"aliases": ["DD_TRACE_INTEGRATION_GAX_ENABLED", "DD_INTEGRATION_GAX_ENABLED"]
5794+
}
5795+
],
57805796
"DD_TRACE_GIT_METADATA_ENABLED": [
57815797
{
57825798
"version": "A",

‎settings.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,7 @@ include(
359359
":dd-java-agent:instrumentation:finatra-2.9",
360360
":dd-java-agent:instrumentation:freemarker:freemarker-2.3.24",
361361
":dd-java-agent:instrumentation:freemarker:freemarker-2.3.9",
362+
":dd-java-agent:instrumentation:gax-1.4",
362363
":dd-java-agent:instrumentation:glassfish-3.0",
363364
":dd-java-agent:instrumentation:google-http-client-1.19",
364365
":dd-java-agent:instrumentation:google-pubsub-1.116",

0 commit comments

Comments
 (0)