Skip to content

Commit 0d03fe1

Browse files
PDGGKclaude
PDGGK
and
claude
committed
Fix RequestResponseIO parseAndThrow to preserve retryable exception types
Problem: The parseAndThrow method in Call.java was wrapping retryable exceptions (UserCodeTimeoutException, UserCodeRemoteSystemException) in a generic UserCodeExecutionException, which breaks the retry logic that depends on exception.shouldRepeat() returning true. Solution: - Scan the full causal chain using Guava's Throwables.getCausalChain() - Preserve all specific retryable exception types (Quota/Timeout/RemoteSystem) - Prefer specific types over generic UserCodeExecutionException when both exist in the chain to prevent masking of retryable exceptions - Handle circular causal chains gracefully by catching IllegalArgumentException Testing: - Added 10 new unit tests covering: * Direct retryable exceptions (Timeout, RemoteSystem) * Nested exceptions (UncheckedExecutionException wrapping) * Generic UserCodeExecutionException masking specific types * Triple-nested exceptions * Circular reference in causal chain * Non-UserCode exceptions (RuntimeException) - All existing tests pass - Full rrio test suite passes (90 tasks) Fixes #37176 Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
1 parent 4dcc27d commit 0d03fe1

2 files changed

Lines changed: 290 additions & 7 deletions

File tree

  • sdks/java/io/rrio/src
    • main/java/org/apache/beam/io/requestresponse
    • test/java/org/apache/beam/io/requestresponse

‎sdks/java/io/rrio/src/main/java/org/apache/beam/io/requestresponse/Call.java‎

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import org.apache.beam.sdk.values.PCollectionTuple;
4747
import org.apache.beam.sdk.values.TupleTag;
4848
import org.apache.beam.sdk.values.TupleTagList;
49+
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
4950
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
5051
import org.joda.time.Duration;
5152

@@ -613,12 +614,32 @@ private void executeAsync(Callable<Void> callable) throws UserCodeExecutionExcep
613614
private static <T> void parseAndThrow(Future<T> future, ExecutionException e)
614615
throws UserCodeExecutionException {
615616
future.cancel(true);
616-
if (e.getCause() == null) {
617-
throw new UserCodeExecutionException(e);
617+
618+
try {
619+
UserCodeExecutionException genericException = null;
620+
for (Throwable throwable : Throwables.getCausalChain(e)) {
621+
if (throwable instanceof UserCodeQuotaException) {
622+
throw (UserCodeQuotaException) throwable;
623+
} else if (throwable instanceof UserCodeTimeoutException) {
624+
throw (UserCodeTimeoutException) throwable;
625+
} else if (throwable instanceof UserCodeRemoteSystemException) {
626+
throw (UserCodeRemoteSystemException) throwable;
627+
} else if (genericException == null && throwable instanceof UserCodeExecutionException) {
628+
genericException = (UserCodeExecutionException) throwable;
629+
}
630+
}
631+
if (genericException != null) {
632+
throw genericException;
633+
}
634+
} catch (IllegalArgumentException iae) {
635+
// Circular reference detected in causal chain
636+
throw new UserCodeExecutionException(
637+
"circular reference detected in exception causal chain", e);
618638
}
619-
Throwable cause = checkStateNotNull(e.getCause());
620-
if (cause instanceof UserCodeQuotaException) {
621-
throw new UserCodeQuotaException(cause);
639+
640+
Throwable cause = e.getCause();
641+
if (cause == null) {
642+
throw new UserCodeExecutionException(e);
622643
}
623644
throw new UserCodeExecutionException(cause);
624645
}

‎sdks/java/io/rrio/src/test/java/org/apache/beam/io/requestresponse/CallTest.java‎

Lines changed: 264 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,34 @@ public void givenCallerThrowsUserCodeExecutionException_emitsIntoFailurePCollect
104104
pipeline.run();
105105
}
106106

107+
@Test
108+
public void givenCallerThrowsNonUserCodeException_emitsWrappedUserCodeExecutionException() {
109+
Result<Response> result =
110+
pipeline
111+
.apply(Create.of(new Request("a")))
112+
.apply(Call.of(new CallerThrowsRuntimeException(), NON_DETERMINISTIC_RESPONSE_CODER));
113+
114+
PCollection<ApiIOError> failures = result.getFailures();
115+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
116+
.isEqualTo(1L);
117+
118+
pipeline.run();
119+
}
120+
121+
@Test
122+
public void givenCallerThrowsCircularCausalChain_emitsUserCodeExecutionException() {
123+
Result<Response> result =
124+
pipeline
125+
.apply(Create.of(new Request("a")))
126+
.apply(Call.of(new CallerThrowsCircularCause(), NON_DETERMINISTIC_RESPONSE_CODER));
127+
128+
PCollection<ApiIOError> failures = result.getFailures();
129+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
130+
.isEqualTo(1L);
131+
132+
pipeline.run();
133+
}
134+
107135
@Test
108136
public void givenCallerThrowsQuotaException_emitsIntoFailurePCollection() {
109137
Result<Response> result =
@@ -142,22 +170,169 @@ public void givenCallerTimeout_emitsFailurePCollection() {
142170
}
143171

144172
@Test
145-
public void givenCallerThrowsTimeoutException_emitsFailurePCollection() {
173+
public void givenCallerThrowsTimeoutException_thenPreservesExceptionType() {
146174
Result<Response> result =
147175
pipeline
148176
.apply(Create.of(new Request("a")))
149177
.apply(Call.of(new CallerThrowsTimeout(), NON_DETERMINISTIC_RESPONSE_CODER));
150178

151179
PCollection<ApiIOError> failures = result.getFailures();
152180
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
153-
.isEqualTo(1L);
181+
.isEqualTo(0L);
154182
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeQuotaException.class)).isEqualTo(0L);
155183
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeTimeoutException.class))
156184
.isEqualTo(1L);
157185

158186
pipeline.run();
159187
}
160188

189+
@Test
190+
public void givenCallerThrowsRemoteSystemException_thenPreservesExceptionType() {
191+
Result<Response> result =
192+
pipeline
193+
.apply(Create.of(new Request("a")))
194+
.apply(
195+
Call.of(new CallerThrowsRemoteSystemException(), NON_DETERMINISTIC_RESPONSE_CODER));
196+
197+
PCollection<ApiIOError> failures = result.getFailures();
198+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeRemoteSystemException.class))
199+
.isEqualTo(1L);
200+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
201+
.isEqualTo(0L);
202+
203+
pipeline.run();
204+
}
205+
206+
@Test
207+
public void givenNestedExecutionException_thenPreservesExceptionType() {
208+
Result<Response> result =
209+
pipeline
210+
.apply(Create.of(new Request("a")))
211+
.apply(
212+
Call.of(
213+
new CallerThrowsNestedExecutionException(), NON_DETERMINISTIC_RESPONSE_CODER));
214+
215+
PCollection<ApiIOError> failures = result.getFailures();
216+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
217+
.isEqualTo(1L);
218+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeTimeoutException.class))
219+
.isEqualTo(0L);
220+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeRemoteSystemException.class))
221+
.isEqualTo(0L);
222+
223+
pipeline.run();
224+
}
225+
226+
@Test
227+
public void givenCallerThrowsGenericWrappingTimeout_thenPreservesExceptionType() {
228+
Result<Response> result =
229+
pipeline
230+
.apply(Create.of(new Request("a")))
231+
.apply(
232+
Call.of(
233+
new CallerThrowsGenericWrappingTimeout(), NON_DETERMINISTIC_RESPONSE_CODER));
234+
235+
PCollection<ApiIOError> failures = result.getFailures();
236+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeTimeoutException.class))
237+
.isEqualTo(1L);
238+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
239+
.isEqualTo(0L);
240+
241+
pipeline.run();
242+
}
243+
244+
@Test
245+
public void givenCallerThrowsGenericWrappingQuota_thenPreservesExceptionType() {
246+
Result<Response> result =
247+
pipeline
248+
.apply(Create.of(new Request("a")))
249+
.apply(
250+
Call.of(new CallerThrowsGenericWrappingQuota(), NON_DETERMINISTIC_RESPONSE_CODER));
251+
252+
PCollection<ApiIOError> failures = result.getFailures();
253+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeQuotaException.class)).isEqualTo(1L);
254+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
255+
.isEqualTo(0L);
256+
257+
pipeline.run();
258+
}
259+
260+
@Test
261+
public void givenCallerThrowsGenericWrappingRemoteSystem_thenPreservesExceptionType() {
262+
Result<Response> result =
263+
pipeline
264+
.apply(Create.of(new Request("a")))
265+
.apply(
266+
Call.of(
267+
new CallerThrowsGenericWrappingRemoteSystem(),
268+
NON_DETERMINISTIC_RESPONSE_CODER));
269+
270+
PCollection<ApiIOError> failures = result.getFailures();
271+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeRemoteSystemException.class))
272+
.isEqualTo(1L);
273+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
274+
.isEqualTo(0L);
275+
276+
pipeline.run();
277+
}
278+
279+
@Test
280+
public void
281+
givenCallerThrowsUncheckedExecutionExceptionWrappingTimeout_thenPreservesExceptionType() {
282+
Result<Response> result =
283+
pipeline
284+
.apply(Create.of(new Request("a")))
285+
.apply(
286+
Call.of(
287+
new CallerThrowsUncheckedExecutionExceptionWrappingTimeout(),
288+
NON_DETERMINISTIC_RESPONSE_CODER));
289+
290+
PCollection<ApiIOError> failures = result.getFailures();
291+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeTimeoutException.class))
292+
.isEqualTo(1L);
293+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
294+
.isEqualTo(0L);
295+
296+
pipeline.run();
297+
}
298+
299+
@Test
300+
public void
301+
givenCallerThrowsUncheckedExecutionExceptionWrappingRemoteSystem_thenPreservesExceptionType() {
302+
Result<Response> result =
303+
pipeline
304+
.apply(Create.of(new Request("a")))
305+
.apply(
306+
Call.of(
307+
new CallerThrowsUncheckedExecutionExceptionWrappingRemoteSystem(),
308+
NON_DETERMINISTIC_RESPONSE_CODER));
309+
310+
PCollection<ApiIOError> failures = result.getFailures();
311+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeRemoteSystemException.class))
312+
.isEqualTo(1L);
313+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
314+
.isEqualTo(0L);
315+
316+
pipeline.run();
317+
}
318+
319+
@Test
320+
public void givenCallerThrowsTripleNestedTimeout_thenPreservesExceptionType() {
321+
Result<Response> result =
322+
pipeline
323+
.apply(Create.of(new Request("a")))
324+
.apply(
325+
Call.of(new CallerThrowsTripleNestedTimeout(), NON_DETERMINISTIC_RESPONSE_CODER));
326+
327+
PCollection<ApiIOError> failures = result.getFailures();
328+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeTimeoutException.class))
329+
.isEqualTo(1L);
330+
PAssert.thatSingleton(countStackTracesOf(failures, UserCodeExecutionException.class))
331+
.isEqualTo(0L);
332+
333+
pipeline.run();
334+
}
335+
161336
@Test
162337
public void givenSetupThrowsUserCodeExecutionException_throwsError() {
163338
pipeline
@@ -376,6 +551,25 @@ public Response call(Request request) throws UserCodeExecutionException {
376551
}
377552
}
378553

554+
private static class CallerThrowsRuntimeException implements Caller<Request, Response> {
555+
556+
@Override
557+
public Response call(Request request) {
558+
throw new RuntimeException("unexpected error");
559+
}
560+
}
561+
562+
private static class CallerThrowsCircularCause implements Caller<Request, Response> {
563+
564+
@Override
565+
public Response call(Request request) {
566+
Exception a = new Exception("a");
567+
Exception b = new Exception("b", a);
568+
a.initCause(b); // a -> b -> a (circular reference)
569+
throw new RuntimeException("boom", a);
570+
}
571+
}
572+
379573
private static class CallerThrowsTimeout implements Caller<Request, Response> {
380574

381575
@Override
@@ -384,6 +578,74 @@ public Response call(Request request) throws UserCodeExecutionException {
384578
}
385579
}
386580

581+
private static class CallerThrowsRemoteSystemException implements Caller<Request, Response> {
582+
583+
@Override
584+
public Response call(Request request) throws UserCodeExecutionException {
585+
throw new UserCodeRemoteSystemException("");
586+
}
587+
}
588+
589+
private static class CallerThrowsNestedExecutionException implements Caller<Request, Response> {
590+
591+
@Override
592+
public Response call(Request request) throws UserCodeExecutionException {
593+
throw new UncheckedExecutionException(new UserCodeExecutionException("nested"));
594+
}
595+
}
596+
597+
private static class CallerThrowsGenericWrappingTimeout implements Caller<Request, Response> {
598+
599+
@Override
600+
public Response call(Request request) throws UserCodeExecutionException {
601+
throw new UserCodeExecutionException("generic", new UserCodeTimeoutException("timeout"));
602+
}
603+
}
604+
605+
private static class CallerThrowsGenericWrappingQuota implements Caller<Request, Response> {
606+
607+
@Override
608+
public Response call(Request request) throws UserCodeExecutionException {
609+
throw new UserCodeExecutionException("generic", new UserCodeQuotaException("quota"));
610+
}
611+
}
612+
613+
private static class CallerThrowsGenericWrappingRemoteSystem
614+
implements Caller<Request, Response> {
615+
616+
@Override
617+
public Response call(Request request) throws UserCodeExecutionException {
618+
throw new UserCodeExecutionException("generic", new UserCodeRemoteSystemException("remote"));
619+
}
620+
}
621+
622+
private static class CallerThrowsUncheckedExecutionExceptionWrappingTimeout
623+
implements Caller<Request, Response> {
624+
625+
@Override
626+
public Response call(Request request) throws UserCodeExecutionException {
627+
throw new UncheckedExecutionException(new UserCodeTimeoutException("timeout"));
628+
}
629+
}
630+
631+
private static class CallerThrowsUncheckedExecutionExceptionWrappingRemoteSystem
632+
implements Caller<Request, Response> {
633+
634+
@Override
635+
public Response call(Request request) throws UserCodeExecutionException {
636+
throw new UncheckedExecutionException(new UserCodeRemoteSystemException("remote"));
637+
}
638+
}
639+
640+
private static class CallerThrowsTripleNestedTimeout implements Caller<Request, Response> {
641+
642+
@Override
643+
public Response call(Request request) throws UserCodeExecutionException {
644+
throw new UncheckedExecutionException(
645+
new RuntimeException(new UserCodeTimeoutException("deep timeout")));
646+
}
647+
}
648+
387649
private static class CallerInvokesQuotaException implements Caller<Request, Response> {
388650

389651
@Override

0 commit comments

Comments
 (0)