Skip to content

Commit

Permalink
aws#4801 Cancelling origin body to prevent stuck calling thread
Browse files Browse the repository at this point in the history
  • Loading branch information
Kirill Chaykin committed Jan 5, 2024
1 parent 5c2868f commit a85a651
Showing 1 changed file with 48 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,13 @@
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;
import software.amazon.awssdk.annotations.SdkPublicApi;
import software.amazon.awssdk.core.exception.NonRetryableException;
import software.amazon.awssdk.core.internal.async.SplittingPublisher;
import software.amazon.awssdk.core.internal.io.SdkLengthAwareInputStream;
import software.amazon.awssdk.core.internal.util.NoopSubscription;
import software.amazon.awssdk.utils.async.DelegatingSubscriber;
import software.amazon.awssdk.utils.async.InputStreamConsumingPublisher;

/**
Expand Down Expand Up @@ -104,6 +107,11 @@ public void subscribe(Subscriber<? super ByteBuffer> s) {
}
}

@Override
public SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration splitConfiguration) {
return new BlockingSplittingPublisher(this, splitConfiguration);
}

private void waitForSubscriptionIfNeeded() throws InterruptedException {
long timeoutSeconds = subscribeTimeout.getSeconds();
if (!subscribedLatch.await(timeoutSeconds, TimeUnit.SECONDS)) {
Expand All @@ -112,4 +120,44 @@ private void waitForSubscriptionIfNeeded() throws InterruptedException {
+ "BEFORE invoking doBlockingWrite if your caller is single-threaded.");
}
}

private class BlockingSplittingPublisher extends SplittingPublisher {

public BlockingSplittingPublisher(AsyncRequestBody asyncRequestBody,
AsyncRequestBodySplitConfiguration splitConfiguration) {
super(asyncRequestBody, splitConfiguration);
}

@Override
public void subscribe(Subscriber<? super AsyncRequestBody> downstreamSubscriber) {
Subscriber<? super AsyncRequestBody> delegatingSubscriber = new DelegatingSubscriber<AsyncRequestBody, AsyncRequestBody>(
downstreamSubscriber) {
@Override
public void onSubscribe(Subscription subscription) {
Subscription delegatingSubscription = new Subscription() {
@Override
public void request(long n) {
subscription.request(n);
}

@Override
public void cancel() {
subscription.cancel();

//Cancel origin body to prevent stuck calling thread
BlockingInputStreamAsyncRequestBody.this.cancel();
}
};
super.onSubscribe(delegatingSubscription);
}

@Override
public void onNext(AsyncRequestBody body) {
subscriber.onNext(body);
}
};

super.subscribe(delegatingSubscriber);
}
}
}

0 comments on commit a85a651

Please sign in to comment.