Skip to content

Commit 45ced32

Browse files
zoewanggzoewangg
committed
Refactoring
1 parent 9b661d8 commit 45ced32

File tree

13 files changed

+435
-409
lines changed

13 files changed

+435
-409
lines changed

‎core/sdk-core/src/main/java/software/amazon/awssdk/core/async/AsyncRequestBody.java‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import java.nio.file.Path;
2525
import java.util.Arrays;
2626
import java.util.Optional;
27-
import java.util.concurrent.CompletableFuture;
2827
import java.util.concurrent.ExecutorService;
2928
import org.reactivestreams.Publisher;
3029
import org.reactivestreams.Subscriber;
@@ -420,24 +419,20 @@ static AsyncRequestBody empty() {
420419
* @param maxMemoryUsageInBytes the max memory the SDK will use to buffer the content
421420
* @return SplitAsyncRequestBodyResult
422421
*/
423-
default SplitAsyncRequestBodyResponse split(long chunkSizeInBytes, long maxMemoryUsageInBytes) {
422+
default SdkPublisher<AsyncRequestBody> split(long chunkSizeInBytes, long maxMemoryUsageInBytes) {
424423
Validate.isPositive(chunkSizeInBytes, "chunkSizeInBytes");
425424
Validate.isPositive(maxMemoryUsageInBytes, "maxMemoryUsageInBytes");
426425

427-
if (!this.contentLength().isPresent()) {
426+
if (!contentLength().isPresent()) {
428427
Validate.isTrue(maxMemoryUsageInBytes >= chunkSizeInBytes,
429428
"maxMemoryUsageInBytes must be larger than or equal to " +
430429
"chunkSizeInBytes if the content length is unknown");
431430
}
432431

433-
CompletableFuture<Void> future = new CompletableFuture<>();
434-
SplittingPublisher splittingPublisher = SplittingPublisher.builder()
435-
.asyncRequestBody(this)
436-
.chunkSizeInBytes(chunkSizeInBytes)
437-
.maxMemoryUsageInBytes(maxMemoryUsageInBytes)
438-
.resultFuture(future)
439-
.build();
440-
441-
return SplitAsyncRequestBodyResponse.create(splittingPublisher, future);
432+
return SplittingPublisher.builder()
433+
.asyncRequestBody(this)
434+
.chunkSizeInBytes(chunkSizeInBytes)
435+
.maxMemoryUsageInBytes(maxMemoryUsageInBytes)
436+
.build();
442437
}
443438
}

‎core/sdk-core/src/main/java/software/amazon/awssdk/core/async/SplitAsyncRequestBodyResponse.java‎

Lines changed: 0 additions & 80 deletions
This file was deleted.

‎core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/SplittingPublisher.java‎

Lines changed: 21 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717

1818
import java.nio.ByteBuffer;
1919
import java.util.Optional;
20-
import java.util.concurrent.CompletableFuture;
2120
import java.util.concurrent.atomic.AtomicBoolean;
2221
import java.util.concurrent.atomic.AtomicInteger;
2322
import java.util.concurrent.atomic.AtomicLong;
@@ -45,24 +44,12 @@ public class SplittingPublisher implements SdkPublisher<AsyncRequestBody> {
4544
private final SimplePublisher<AsyncRequestBody> downstreamPublisher = new SimplePublisher<>();
4645
private final long chunkSizeInBytes;
4746
private final long maxMemoryUsageInBytes;
48-
private final CompletableFuture<Void> future;
4947

5048
private SplittingPublisher(Builder builder) {
5149
this.upstreamPublisher = Validate.paramNotNull(builder.asyncRequestBody, "asyncRequestBody");
5250
this.chunkSizeInBytes = Validate.isPositive(builder.chunkSizeInBytes, "chunkSizeInBytes");
5351
this.splittingSubscriber = new SplittingSubscriber(upstreamPublisher.contentLength().orElse(null));
5452
this.maxMemoryUsageInBytes = Validate.isPositive(builder.maxMemoryUsageInBytes, "maxMemoryUsageInBytes");
55-
this.future = builder.future;
56-
57-
// We need to cancel upstream subscription if the future gets cancelled.
58-
future.whenComplete((r, t) -> {
59-
if (t != null) {
60-
if (splittingSubscriber.upstreamSubscription != null) {
61-
log.trace(() -> "Cancelling subscription because return future completed exceptionally ", t);
62-
splittingSubscriber.upstreamSubscription.cancel();
63-
}
64-
}
65-
});
6653
}
6754

6855
public static Builder builder() {
@@ -117,16 +104,20 @@ public void onNext(ByteBuffer byteBuffer) {
117104
byteBufferSizeHint = byteBuffer.remaining();
118105

119106
while (true) {
107+
108+
if (!byteBuffer.hasRemaining()) {
109+
break;
110+
}
111+
120112
int amountRemainingInChunk = amountRemainingInChunk();
121113

122114
// If we have fulfilled this chunk,
123115
// complete the current body
124116
if (amountRemainingInChunk == 0) {
125117
completeCurrentBodyAndCreateNewIfNeeded(byteBuffer);
118+
amountRemainingInChunk = amountRemainingInChunk();
126119
}
127120

128-
amountRemainingInChunk = amountRemainingInChunk();
129-
130121
// If the current ByteBuffer < this chunk, send it as-is
131122
if (amountRemainingInChunk > byteBuffer.remaining()) {
132123
currentBody.send(byteBuffer.duplicate());
@@ -154,29 +145,28 @@ public void onNext(ByteBuffer byteBuffer) {
154145

155146
private void completeCurrentBodyAndCreateNewIfNeeded(ByteBuffer byteBuffer) {
156147
completeCurrentBody();
148+
int currentChunk = chunkNumber.incrementAndGet();
149+
boolean shouldCreateNewDownstreamRequestBody;
150+
Long dataRemaining = totalDataRemaining();
157151

158-
if (shouldCreateNewDownstreamRequestBody(byteBuffer)) {
159-
int currentChunk = chunkNumber.incrementAndGet();
160-
long chunkSize = calculateChunkSize(totalDataRemaining());
161-
currentBody = initializeNextDownstreamBody(upstreamSize != null, chunkSize, currentChunk);
152+
if (upstreamSize == null) {
153+
shouldCreateNewDownstreamRequestBody = !upstreamComplete || byteBuffer.hasRemaining();
154+
} else {
155+
shouldCreateNewDownstreamRequestBody = dataRemaining != null && dataRemaining > 0;
162156
}
163-
}
164-
165157

166-
/**
167-
* If content length is known, we should create new DownstreamRequestBody if there's remaining data.
168-
* If content length is unknown, we should create new DownstreamRequestBody if upstream is not completed yet.
169-
*/
170-
private boolean shouldCreateNewDownstreamRequestBody(ByteBuffer byteBuffer) {
171-
return !upstreamComplete || byteBuffer.remaining() > 0;
158+
if (shouldCreateNewDownstreamRequestBody) {
159+
long chunkSize = calculateChunkSize(dataRemaining);
160+
currentBody = initializeNextDownstreamBody(upstreamSize != null, chunkSize, currentChunk);
161+
}
172162
}
173163

174164
private int amountRemainingInChunk() {
175165
return Math.toIntExact(currentBody.maxLength - currentBody.transferredLength);
176166
}
177167

178168
private void completeCurrentBody() {
179-
log.debug(() -> "completeCurrentBody");
169+
log.debug(() -> "completeCurrentBody for chunk " + chunkNumber.get());
180170
currentBody.complete();
181171
if (upstreamSize == null) {
182172
sendCurrentBody(currentBody);
@@ -188,16 +178,16 @@ public void onComplete() {
188178
upstreamComplete = true;
189179
log.trace(() -> "Received onComplete()");
190180
completeCurrentBody();
191-
downstreamPublisher.complete().thenRun(() -> future.complete(null));
181+
downstreamPublisher.complete();
192182
}
193183

194184
@Override
195185
public void onError(Throwable t) {
196-
currentBody.error(t);
186+
log.trace(() -> "Received onError()", t);
187+
downstreamPublisher.error(t);
197188
}
198189

199190
private void sendCurrentBody(AsyncRequestBody body) {
200-
log.debug(() -> "sendCurrentBody");
201191
downstreamPublisher.send(body).exceptionally(t -> {
202192
downstreamPublisher.error(t);
203193
return null;
@@ -300,7 +290,6 @@ public static final class Builder {
300290
private AsyncRequestBody asyncRequestBody;
301291
private Long chunkSizeInBytes;
302292
private Long maxMemoryUsageInBytes;
303-
private CompletableFuture<Void> future;
304293

305294
/**
306295
* Configures the asyncRequestBody to split
@@ -339,18 +328,6 @@ public Builder maxMemoryUsageInBytes(long maxMemoryUsageInBytes) {
339328
return this;
340329
}
341330

342-
/**
343-
* Sets the result future. The future will be completed when all request bodies
344-
* have been sent.
345-
*
346-
* @param future The new future value.
347-
* @return This object for method chaining.
348-
*/
349-
public Builder resultFuture(CompletableFuture<Void> future) {
350-
this.future = future;
351-
return this;
352-
}
353-
354331
public SplittingPublisher build() {
355332
return new SplittingPublisher(this);
356333
}

‎core/sdk-core/src/test/java/software/amazon/awssdk/core/async/SplitAsyncRequestBodyResponseTest.java‎

Lines changed: 0 additions & 29 deletions
This file was deleted.

‎core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/SplittingPublisherTest.java‎

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,7 @@
3131
import java.util.Optional;
3232
import java.util.UUID;
3333
import java.util.concurrent.CompletableFuture;
34-
import java.util.concurrent.ExecutionException;
35-
import java.util.concurrent.ThreadLocalRandom;
3634
import java.util.concurrent.TimeUnit;
37-
import java.util.concurrent.TimeoutException;
3835
import java.util.concurrent.atomic.AtomicInteger;
3936
import org.apache.commons.lang3.RandomStringUtils;
4037
import org.junit.jupiter.api.AfterAll;
@@ -45,7 +42,6 @@
4542
import org.reactivestreams.Subscriber;
4643
import org.reactivestreams.Subscription;
4744
import software.amazon.awssdk.core.async.AsyncRequestBody;
48-
import software.amazon.awssdk.testutils.RandomTempFile;
4945
import software.amazon.awssdk.utils.BinaryUtils;
5046

5147
public class SplittingPublisherTest {
@@ -87,26 +83,6 @@ void differentChunkSize_byteArrayShouldSplitAsyncRequestBodyCorrectly(int chunkS
8783
verifySplitContent(AsyncRequestBody.fromBytes(CONTENT), chunkSize);
8884
}
8985

90-
91-
@Test
92-
void cancelFuture_shouldCancelUpstream() throws IOException {
93-
CompletableFuture<Void> future = new CompletableFuture<>();
94-
TestAsyncRequestBody asyncRequestBody = new TestAsyncRequestBody();
95-
SplittingPublisher splittingPublisher = SplittingPublisher.builder()
96-
.resultFuture(future)
97-
.asyncRequestBody(asyncRequestBody)
98-
.chunkSizeInBytes(CHUNK_SIZE)
99-
.maxMemoryUsageInBytes(10L)
100-
.build();
101-
102-
OnlyRequestOnceSubscriber downstreamSubscriber = new OnlyRequestOnceSubscriber();
103-
splittingPublisher.subscribe(downstreamSubscriber);
104-
105-
future.completeExceptionally(new RuntimeException("test"));
106-
assertThat(asyncRequestBody.cancelled).isTrue();
107-
assertThat(downstreamSubscriber.asyncRequestBodies.size()).isEqualTo(1);
108-
}
109-
11086
@Test
11187
void contentLengthNotPresent_shouldHandle() throws Exception {
11288
CompletableFuture<Void> future = new CompletableFuture<>();
@@ -117,7 +93,6 @@ public Optional<Long> contentLength() {
11793
}
11894
};
11995
SplittingPublisher splittingPublisher = SplittingPublisher.builder()
120-
.resultFuture(future)
12196
.asyncRequestBody(asyncRequestBody)
12297
.chunkSizeInBytes(CHUNK_SIZE)
12398
.maxMemoryUsageInBytes(10L)
@@ -159,11 +134,8 @@ public Optional<Long> contentLength() {
159134

160135

161136
private static void verifySplitContent(AsyncRequestBody asyncRequestBody, int chunkSize) throws Exception {
162-
CompletableFuture<Void> future = new CompletableFuture<>();
163137
SplittingPublisher splittingPublisher = SplittingPublisher.builder()
164-
.resultFuture(future)
165138
.asyncRequestBody(asyncRequestBody)
166-
.resultFuture(future)
167139
.chunkSizeInBytes(chunkSize)
168140
.maxMemoryUsageInBytes((long) chunkSize * 4)
169141
.build();
@@ -194,7 +166,6 @@ private static void verifySplitContent(AsyncRequestBody asyncRequestBody, int ch
194166
assertThat(actualBytes).isEqualTo(expected);
195167
};
196168
}
197-
assertThat(future).isCompleted();
198169
}
199170

200171
private static class TestAsyncRequestBody implements AsyncRequestBody {

‎services/s3/src/it/java/software/amazon/awssdk/services/s3/multipart/S3MultipartClientPutObjectIntegrationTest.java‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public class S3MultipartClientPutObjectIntegrationTest extends S3IntegrationTest
5252

5353
private static final String TEST_BUCKET = temporaryBucketName(S3MultipartClientPutObjectIntegrationTest.class);
5454
private static final String TEST_KEY = "testfile.dat";
55-
private static final int OBJ_SIZE = 31 * 1024 * 1024;
55+
private static final int OBJ_SIZE = 19 * 1024 * 1024;
5656

5757
private static File testFile;
5858
private static S3AsyncClient mpuS3Client;
@@ -90,7 +90,6 @@ void putObject_fileRequestBody_objectSentCorrectly() throws Exception {
9090
}
9191

9292
@Test
93-
@Timeout(value = 30, unit = SECONDS)
9493
void putObject_byteAsyncRequestBody_objectSentCorrectly() throws Exception {
9594
byte[] bytes = RandomStringUtils.randomAscii(OBJ_SIZE).getBytes(Charset.defaultCharset());
9695
AsyncRequestBody body = AsyncRequestBody.fromBytes(bytes);
@@ -105,11 +104,9 @@ void putObject_byteAsyncRequestBody_objectSentCorrectly() throws Exception {
105104
}
106105

107106
@Test
108-
@Timeout(value = 30, unit = SECONDS)
109107
void putObject_unknownContentLength_objectSentCorrectly() throws Exception {
110108
AsyncRequestBody body = FileAsyncRequestBody.builder()
111109
.path(testFile.toPath())
112-
//.chunkSizeInBytes(2 * 1024 * 1024)
113110
.build();
114111
mpuS3Client.putObject(r -> r.bucket(TEST_BUCKET).key(TEST_KEY), new AsyncRequestBody() {
115112
@Override

0 commit comments

Comments
 (0)