2525import software .amazon .awssdk .core .SplittingTransformerConfiguration ;
2626import software .amazon .awssdk .core .async .AsyncResponseTransformer ;
2727import software .amazon .awssdk .core .async .SdkPublisher ;
28+ import software .amazon .awssdk .core .exception .SdkException ;
29+ import software .amazon .awssdk .core .exception .SdkServiceException ;
30+ import software .amazon .awssdk .core .internal .retry .SdkDefaultRetrySetting ;
31+ import software .amazon .awssdk .core .retry .RetryUtils ;
2832import software .amazon .awssdk .utils .CompletableFutureUtils ;
2933import software .amazon .awssdk .utils .Logger ;
3034import software .amazon .awssdk .utils .Validate ;
@@ -54,16 +58,6 @@ public class SplittingTransformer<ResponseT, ResultT> implements SdkPublisher<As
5458 */
5559 private final AsyncResponseTransformer <ResponseT , ResultT > upstreamResponseTransformer ;
5660
57- /**
58- * Set to true once {@code .prepare()} is called on the upstreamResponseTransformer
59- */
60- private final AtomicBoolean preparedCalled = new AtomicBoolean (false );
61-
62- /**
63- * Set to true once {@code .onResponse()} is called on the upstreamResponseTransformer
64- */
65- private final AtomicBoolean onResponseCalled = new AtomicBoolean (false );
66-
6761 /**
6862 * Set to true once {@code .onStream()} is called on the upstreamResponseTransformer
6963 */
@@ -111,6 +105,24 @@ public class SplittingTransformer<ResponseT, ResultT> implements SdkPublisher<As
111105
112106 private final Object cancelLock = new Object ();
113107
108+ /**
109+ * Keeps track of the upstream future returned by {@code upstreamResponseTransformer.prepare()}. If an error occurs, we
110+ * forward this to the {@code resultFuture}.
111+ */
112+ private volatile CompletableFuture <ResultT > upstreamFuture ;
113+
114+ /**
115+ * Tracks whether an {@code IndividualTransformer} has been created for the first part yet. Errors will only be retried for
116+ * the first part.
117+ */
118+ private final AtomicBoolean isFirstIndividualTransformer = new AtomicBoolean (true );
119+
120+ /**
121+ * Tracks whether an {@code IndividualPartSubscriber} has been created for the first part yet. Errors will only be retried for
122+ * the first part.
123+ */
124+ private final AtomicBoolean isFirstIndividualSubscriber = new AtomicBoolean (true );
125+
114126 private SplittingTransformer (AsyncResponseTransformer <ResponseT , ResultT > upstreamResponseTransformer ,
115127 Long maximumBufferSizeInBytes ,
116128 CompletableFuture <ResultT > resultFuture ) {
@@ -198,7 +210,7 @@ private boolean doEmit() {
198210 }
199211 if (outstandingDemand .get () > 0 ) {
200212 demand = outstandingDemand .decrementAndGet ();
201- downstreamSubscriber .onNext (new IndividualTransformer ());
213+ downstreamSubscriber .onNext (new IndividualTransformer (isFirstIndividualTransformer . compareAndSet ( true , false ) ));
202214 }
203215 }
204216 return false ;
@@ -230,6 +242,7 @@ private void handleSubscriptionCancel() {
230242 } else {
231243 log .trace (() -> "calling downstreamSubscriber.onComplete()" );
232244 downstreamSubscriber .onComplete ();
245+ CompletableFutureUtils .forwardResultTo (upstreamFuture , resultFuture );
233246 }
234247 downstreamSubscriber = null ;
235248 });
@@ -259,28 +272,27 @@ private void handleFutureCancel(Throwable e) {
259272 * body publisher.
260273 */
261274 private class IndividualTransformer implements AsyncResponseTransformer <ResponseT , ResponseT > {
275+ private final boolean isFirstPart ;
262276 private ResponseT response ;
263277 private CompletableFuture <ResponseT > individualFuture ;
264278
279+ IndividualTransformer (boolean isFirstPart ) {
280+ this .isFirstPart = isFirstPart ;
281+ }
282+
265283 @ Override
266284 public CompletableFuture <ResponseT > prepare () {
267285 this .individualFuture = new CompletableFuture <>();
268- if (preparedCalled .compareAndSet (false , true )) {
286+
287+ if (isFirstPart ) {
269288 if (isCancelled .get ()) {
270289 return individualFuture ;
271290 }
272291 log .trace (() -> "calling prepare on the upstream transformer" );
273- CompletableFuture <ResultT > upstreamFuture = upstreamResponseTransformer .prepare ();
274- if (!resultFuture .isDone ()) {
275- CompletableFutureUtils .forwardResultTo (upstreamFuture , resultFuture );
276- }
292+ upstreamFuture = upstreamResponseTransformer .prepare ();
293+
277294 }
278- resultFuture .whenComplete ((r , e ) -> {
279- if (e == null ) {
280- return ;
281- }
282- individualFuture .completeExceptionally (e );
283- });
295+
284296 individualFuture .whenComplete ((r , e ) -> {
285297 if (isCancelled .get ()) {
286298 handleSubscriptionCancel ();
@@ -291,7 +303,7 @@ public CompletableFuture<ResponseT> prepare() {
291303
292304 @ Override
293305 public void onResponse (ResponseT response ) {
294- if (onResponseCalled . compareAndSet ( false , true ) ) {
306+ if (isFirstPart ) {
295307 log .trace (() -> "calling onResponse on the upstream transformer" );
296308 upstreamResponseTransformer .onResponse (response );
297309 }
@@ -304,7 +316,8 @@ public void onStream(SdkPublisher<ByteBuffer> publisher) {
304316 return ;
305317 }
306318 synchronized (cancelLock ) {
307- if (onStreamCalled .compareAndSet (false , true )) {
319+ if (isFirstPart ) {
320+ onStreamCalled .set (true );
308321 log .trace (() -> "calling onStream on the upstream transformer" );
309322 upstreamResponseTransformer .onStream (upstreamSubscriber -> publisherToUpstream .subscribe (
310323 DelegatingBufferingSubscriber .builder ()
@@ -314,29 +327,65 @@ public void onStream(SdkPublisher<ByteBuffer> publisher) {
314327 );
315328 }
316329 }
317- publisher .subscribe (new IndividualPartSubscriber <>(this .individualFuture , response ));
330+
331+ if (!resultFuture .isDone ()) {
332+ CompletableFutureUtils .forwardResultTo (upstreamFuture , resultFuture );
333+ }
334+
335+ publisher .subscribe (new IndividualPartSubscriber <>(this .individualFuture , response ,
336+ isFirstIndividualSubscriber .compareAndSet (true , false )));
318337 }
319338
320339 @ Override
321340 public void exceptionOccurred (Throwable error ) {
322- publisherToUpstream .error (error );
323341 log .trace (() -> "calling exceptionOccurred on the upstream transformer" );
324342 upstreamResponseTransformer .exceptionOccurred (error );
343+
344+ if (!isFirstPart || onStreamCalled .get ()) {
345+ publisherToUpstream .error (error );
346+ }
347+
348+ if (!isRetryableError (error )) {
349+ log .trace (() -> "error is non-retryable, forwarding upstream future to result future" );
350+ CompletableFutureUtils .forwardResultTo (upstreamFuture , resultFuture );
351+ }
325352 }
326353 }
327354
355+ private boolean isRetryableError (Throwable error ) {
356+ if (error instanceof SdkException ) {
357+ SdkException ex = (SdkException ) error ;
358+ return retryOnStatusCodes (ex )
359+ || RetryUtils .isRetryableException (ex )
360+ || RetryUtils .isClockSkewException (ex )
361+ || RetryUtils .isClockSkewException (ex );
362+
363+ }
364+ return false ;
365+ }
366+
367+ private static boolean retryOnStatusCodes (Throwable ex ) {
368+ if (ex instanceof SdkServiceException ) {
369+ SdkServiceException failure = (SdkServiceException ) ex ;
370+ return SdkDefaultRetrySetting .RETRYABLE_STATUS_CODES .contains (failure .statusCode ());
371+ }
372+ return false ;
373+ }
374+
328375 /**
329376 * the Subscriber for each of the individual request's ByteBuffer publisher
330377 */
331378 class IndividualPartSubscriber <T > implements Subscriber <ByteBuffer > {
332379
333380 private final CompletableFuture <T > future ;
334381 private final T response ;
382+ private final boolean isFirstPart ;
335383 private Subscription subscription ;
336384
337- IndividualPartSubscriber (CompletableFuture <T > future , T response ) {
385+ IndividualPartSubscriber (CompletableFuture <T > future , T response , boolean isFirstPart ) {
338386 this .future = future ;
339387 this .response = response ;
388+ this .isFirstPart = isFirstPart ;
340389 }
341390
342391 @ Override
@@ -376,7 +425,9 @@ public void onComplete() {
376425 }
377426
378427 private void handleError (Throwable t ) {
379- publisherToUpstream .error (t );
428+ if (!isFirstPart ) {
429+ publisherToUpstream .error (t );
430+ }
380431 future .completeExceptionally (t );
381432 }
382433 }
0 commit comments