Skip to content

Commit fd3ad14

Browse files
alextwoodsalextwoods
authored
Use CRT response file for downloadFile in S3 TransferManager (#7230)
* Use CRT response file for downloadFile in S3 TransferManager * Add javadoc to pause mentioning blocking until requests are cleaned up. * Fix changelog name
1 parent b213aab commit fd3ad14

19 files changed

Lines changed: 1347 additions & 45 deletions

File tree

‎.changes/next-release/feature-S3TransferManager-9c1f4ab.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "S3 Transfer Manager",
4+
"contributor": "",
5+
"description": "Use CRT's response file when downloading to a file with the CRT-based S3TransferManager - CRT now writes the object to the destination file directly instead of streaming the body back through the SDK. Pausing such a download uses CRT's new asynchronous pause, so the resume token reflects a fully flushed file."
6+
}

‎build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,7 @@
333333
<Class name="~software\.amazon\.awssdk\.http\.nio\.netty\.internal\.AwaitCloseChannelPoolMap" />
334334
<Class name="~software\.amazon\.awssdk\.http\.nio\.netty\.internal\.ChannelAttributeKey" />
335335
<Class name="~software\.amazon\.awssdk\.transfer\.s3\.internal\.progress\.ResumeTransferProgress" />
336+
<Class name="~software\.amazon\.awssdk\.transfer\.s3\.internal\.model\.CrtFileDownload" />
336337
<Class name="~software\.amazon\.awssdk\.utils\.CompletableFutureUtils" />
337338
<Class name="~software\.amazon\.awssdk\.metrics\.publishers\.cloudwatch\.CloudWatchMetricPublisher" />
338339
<Class name="~software\.amazon\.awssdk\.http\.apache\.internal\.conn\.IdleConnectionReaper\$ReaperTask" />

‎services-custom/s3-transfer-manager/src/it/java/software/amazon/awssdk/transfer/s3/S3TransferManagerDownloadPauseResumeIntegrationTest.java‎

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import java.nio.file.Files;
2727
import java.nio.file.Path;
2828
import java.time.Duration;
29-
import java.util.Optional;
3029
import org.apache.commons.lang3.RandomStringUtils;
3130
import org.apache.logging.log4j.Level;
3231
import org.assertj.core.api.Assertions;
@@ -35,12 +34,11 @@
3534
import org.junit.jupiter.api.BeforeAll;
3635
import org.junit.jupiter.params.ParameterizedTest;
3736
import org.junit.jupiter.params.provider.MethodSource;
38-
import software.amazon.awssdk.core.SdkResponse;
3937
import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
4038
import software.amazon.awssdk.core.sync.RequestBody;
4139
import software.amazon.awssdk.core.waiters.Waiter;
4240
import software.amazon.awssdk.core.waiters.WaiterAcceptor;
43-
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
41+
import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
4442
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
4543
import software.amazon.awssdk.services.s3.model.PutObjectResponse;
4644
import software.amazon.awssdk.testutils.LogCaptor;
@@ -94,7 +92,10 @@ void pauseAndResume_ObjectEtagChange_shouldRestartDownload(S3TransferManager tm)
9492
ResumableFileDownload resumableFileDownload = download.pause();
9593
log.debug(() -> "Paused: " + resumableFileDownload);
9694

97-
String originalEtag = testDownloadListener.getObjectResponse.eTag();
95+
// The ETag has to be read off the resume token rather than the progress snapshot: when CRT writes the body directly
96+
// to the file, the SDK never sees the GetObjectResponse until the transfer finishes.
97+
String originalEtag = resumableFileDownload.s3ObjectEtag().orElse(null);
98+
assertThat(originalEtag).isNotNull();
9899

99100
File newSourceFile = new RandomTempFile(OBJ_SIZE);
100101
PutObjectResponse putResponse = s3.putObject(PutObjectRequest.builder()
@@ -149,12 +150,13 @@ void pauseAndResume_ObjectNotChanged_shouldResumeDownload(S3TransferManager tm)
149150
long bytesTransferred = resumableFileDownload.bytesTransferred();
150151
log.debug(() -> "Paused: " + resumableFileDownload);
151152
assertEqualsBySdkFields(resumableFileDownload.downloadFileRequest(), request);
152-
assertThat(testDownloadListener.getObjectResponse).isNotNull();
153153

154154
// Skip the test if everything has been downloaded.
155155
Assumptions.assumeTrue(resumableFileDownload.bytesTransferred() < sourceFile.length());
156156

157-
assertThat(resumableFileDownload.s3ObjectLastModified()).hasValue(testDownloadListener.getObjectResponse.lastModified());
157+
HeadObjectResponse headObjectResponse = s3.headObject(b -> b.bucket(BUCKET).key(KEY));
158+
assertThat(resumableFileDownload.s3ObjectLastModified()).hasValue(headObjectResponse.lastModified());
159+
assertThat(resumableFileDownload.s3ObjectEtag()).hasValue(headObjectResponse.eTag());
158160
// Request may not be cancelled right away when pause is invoked, so there may be more bytes written to the file
159161
assertThat(bytesTransferred).isLessThanOrEqualTo(path.toFile().length());
160162
assertThat(resumableFileDownload.totalSizeInBytes()).hasValue(sourceFile.length());
@@ -254,20 +256,11 @@ private static void waitUntilFirstByteBufferDelivered(FileDownload download) {
254256

255257
private static final class TestDownloadListener implements TransferListener {
256258
private int transferInitiatedCount = 0;
257-
private GetObjectResponse getObjectResponse;
258259

259260
@Override
260261
public void transferInitiated(Context.TransferInitiated context) {
261262
transferInitiatedCount++;
262263
}
263-
264-
@Override
265-
public void bytesTransferred(Context.BytesTransferred context) {
266-
Optional<SdkResponse> sdkResponse = context.progressSnapshot().sdkResponse();
267-
if (sdkResponse.isPresent() && sdkResponse.get() instanceof GetObjectResponse) {
268-
getObjectResponse = (GetObjectResponse) sdkResponse.get();
269-
}
270-
}
271264
}
272265

273266
}

‎services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/CrtS3TransferManager.java‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,38 +18,58 @@
1818
import static software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute.SDK_HTTP_EXECUTION_ATTRIBUTES;
1919
import static software.amazon.awssdk.services.s3.crt.S3CrtSdkHttpExecutionAttribute.CRT_PROGRESS_LISTENER;
2020
import static software.amazon.awssdk.services.s3.crt.S3CrtSdkHttpExecutionAttribute.METAREQUEST_PAUSE_OBSERVABLE;
21+
import static software.amazon.awssdk.services.s3.internal.crt.DefaultS3CrtAsyncClient.RESPONSE_FILE_OPTION;
22+
import static software.amazon.awssdk.services.s3.internal.crt.DefaultS3CrtAsyncClient.RESPONSE_FILE_PATH;
2123
import static software.amazon.awssdk.services.s3.internal.crt.S3InternalSdkHttpExecutionAttribute.CRT_PAUSE_RESUME_TOKEN;
24+
import static software.amazon.awssdk.transfer.s3.internal.utils.ResumableRequestConverter.canResumeDownload;
25+
import static software.amazon.awssdk.transfer.s3.internal.utils.ResumableRequestConverter.toCrtDownloadFileRequest;
2226

2327
import java.util.concurrent.CompletableFuture;
2428
import java.util.function.Consumer;
2529
import software.amazon.awssdk.annotations.SdkInternalApi;
2630
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
2731
import software.amazon.awssdk.core.async.AsyncRequestBody;
32+
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
2833
import software.amazon.awssdk.crt.s3.ResumeToken;
34+
import software.amazon.awssdk.crt.s3.S3MetaRequestOptions.ResponseFileOption;
2935
import software.amazon.awssdk.http.SdkHttpExecutionAttributes;
3036
import software.amazon.awssdk.services.s3.S3AsyncClient;
37+
import software.amazon.awssdk.services.s3.internal.crt.CrtResponseFileResponseTransformer;
3138
import software.amazon.awssdk.services.s3.internal.crt.S3MetaRequestPauseObservable;
39+
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
40+
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
41+
import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
3242
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
3343
import software.amazon.awssdk.services.s3.model.PutObjectResponse;
3444
import software.amazon.awssdk.transfer.s3.S3TransferManager;
45+
import software.amazon.awssdk.transfer.s3.internal.model.CrtFileDownload;
3546
import software.amazon.awssdk.transfer.s3.internal.model.CrtFileUpload;
3647
import software.amazon.awssdk.transfer.s3.internal.model.DefaultUpload;
48+
import software.amazon.awssdk.transfer.s3.internal.progress.ResumeTransferProgress;
3749
import software.amazon.awssdk.transfer.s3.internal.progress.TransferProgressUpdater;
50+
import software.amazon.awssdk.transfer.s3.model.CompletedFileDownload;
3851
import software.amazon.awssdk.transfer.s3.model.CompletedFileUpload;
3952
import software.amazon.awssdk.transfer.s3.model.CompletedUpload;
53+
import software.amazon.awssdk.transfer.s3.model.DownloadFileRequest;
54+
import software.amazon.awssdk.transfer.s3.model.FileDownload;
4055
import software.amazon.awssdk.transfer.s3.model.FileUpload;
56+
import software.amazon.awssdk.transfer.s3.model.ResumableFileDownload;
4157
import software.amazon.awssdk.transfer.s3.model.ResumableFileUpload;
4258
import software.amazon.awssdk.transfer.s3.model.Upload;
4359
import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
4460
import software.amazon.awssdk.transfer.s3.model.UploadRequest;
61+
import software.amazon.awssdk.transfer.s3.progress.TransferProgress;
4562
import software.amazon.awssdk.utils.CompletableFutureUtils;
63+
import software.amazon.awssdk.utils.Logger;
4664
import software.amazon.awssdk.utils.Validate;
4765

4866
/**
4967
* An implementation of {@link S3TransferManager} that uses CRT-based S3 client under the hood.
5068
*/
5169
@SdkInternalApi
5270
class CrtS3TransferManager extends GenericS3TransferManager {
71+
private static final Logger log = Logger.loggerFor(S3TransferManager.class);
72+
5373
private final S3AsyncClient s3AsyncClient;
5474

5575
CrtS3TransferManager(TransferManagerConfiguration transferConfiguration, S3AsyncClient s3AsyncClient,
@@ -137,6 +157,137 @@ FileUpload doResumeUpload(ResumableFileUpload resumableFileUpload) {
137157
.build());
138158
}
139159

160+
/**
161+
* Downloads the object by handing the destination file to CRT, which writes the response body to it directly rather than
162+
* streaming the body back through the SDK.
163+
*/
164+
@Override
165+
FileDownload doDownloadFile(DownloadFileRequest downloadRequest) {
166+
S3MetaRequestPauseObservable observable = new S3MetaRequestPauseObservable();
167+
TransferProgressUpdater progressUpdater = new TransferProgressUpdater(downloadRequest, null);
168+
169+
DownloadFileRequest crtDownloadRequest = crtDownloadFileRequest(downloadRequest, observable, progressUpdater,
170+
ResponseFileOption.CREATE_OR_REPLACE);
171+
172+
CompletableFuture<CompletedFileDownload> returnFuture = new CompletableFuture<>();
173+
initiateCrtDownload(crtDownloadRequest, progressUpdater, returnFuture);
174+
175+
return new CrtFileDownload(returnFuture, progressUpdater.progress(), observable, () -> crtDownloadRequest, null);
176+
}
177+
178+
@Override
179+
FileDownload doResumeDownloadFile(ResumableFileDownload resumableFileDownload) {
180+
DownloadFileRequest originalDownloadRequest = resumableFileDownload.downloadFileRequest();
181+
GetObjectRequest getObjectRequest = originalDownloadRequest.getObjectRequest();
182+
183+
CompletableFuture<CompletedFileDownload> returnFuture = new CompletableFuture<>();
184+
CompletableFuture<TransferProgress> progressFuture = new CompletableFuture<>();
185+
CompletableFuture<DownloadFileRequest> newDownloadFileRequestFuture = new CompletableFuture<>();
186+
S3MetaRequestPauseObservable observable = new S3MetaRequestPauseObservable();
187+
188+
CompletableFuture<HeadObjectResponse> headFuture =
189+
s3AsyncClient.headObject(b -> b.bucket(getObjectRequest.bucket()).key(getObjectRequest.key()));
190+
191+
// Ensure cancellations are forwarded to the head future
192+
CompletableFutureUtils.forwardExceptionTo(returnFuture, headFuture);
193+
194+
headFuture.thenAccept(headObjectResponse -> {
195+
boolean restartFromBeginning = !canResumeDownload(resumableFileDownload, headObjectResponse)
196+
|| hasCompletedParts(resumableFileDownload);
197+
198+
DownloadFileRequest newDownloadFileRequest = toCrtDownloadFileRequest(resumableFileDownload, headObjectResponse,
199+
originalDownloadRequest,
200+
restartFromBeginning);
201+
202+
// CRT appends to whatever the destination file already holds, so appending is only correct when the download is
203+
// genuinely being continued. Otherwise the file has to be replaced.
204+
ResponseFileOption responseFileOption = restartFromBeginning ? ResponseFileOption.CREATE_OR_REPLACE
205+
: ResponseFileOption.CREATE_OR_APPEND;
206+
207+
TransferProgressUpdater progressUpdater = new TransferProgressUpdater(newDownloadFileRequest, null);
208+
DownloadFileRequest crtDownloadRequest = crtDownloadFileRequest(newDownloadFileRequest, observable,
209+
progressUpdater, responseFileOption);
210+
211+
newDownloadFileRequestFuture.complete(crtDownloadRequest);
212+
log.debug(() -> "Sending downloadFileRequest " + crtDownloadRequest);
213+
214+
initiateCrtDownload(crtDownloadRequest, progressUpdater, returnFuture);
215+
progressFuture.complete(progressUpdater.progress());
216+
}).exceptionally(throwable -> {
217+
handleException(returnFuture, progressFuture, newDownloadFileRequestFuture, throwable);
218+
return null;
219+
});
220+
221+
return new CrtFileDownload(returnFuture,
222+
new ResumeTransferProgress(progressFuture),
223+
observable,
224+
() -> newOrOriginalRequestForPause(newDownloadFileRequestFuture,
225+
originalDownloadRequest),
226+
resumableFileDownload);
227+
}
228+
229+
/**
230+
* A download that was paused while fetching individual parts can leave gaps in the destination file, because parts are
231+
* written at their own offsets rather than in order. The CRT-based client only ever appends to the end of the file, so
232+
* such a download cannot be continued and has to start over. This only arises when a {@link ResumableFileDownload}
233+
* produced by the Java-based transfer manager is resumed with the CRT-based one.
234+
*/
235+
private static boolean hasCompletedParts(ResumableFileDownload resumableFileDownload) {
236+
if (resumableFileDownload.completedParts().isEmpty()) {
237+
return false;
238+
}
239+
log.debug(() -> "The paused download had completed individual parts, which the CRT-based S3 client cannot continue "
240+
+ "from. The SDK will download the S3 object from the beginning.");
241+
return true;
242+
}
243+
244+
private void initiateCrtDownload(DownloadFileRequest downloadRequest,
245+
TransferProgressUpdater progressUpdater,
246+
CompletableFuture<CompletedFileDownload> returnFuture) {
247+
try {
248+
progressUpdater.transferInitiated();
249+
AsyncResponseTransformer<GetObjectResponse, GetObjectResponse> responseTransformer =
250+
progressUpdater.wrapCrtResponseFileTransformer(new CrtResponseFileResponseTransformer<>());
251+
progressUpdater.registerCompletion(returnFuture);
252+
253+
assertNotUnsupportedArn(downloadRequest.getObjectRequest().bucket(), "download");
254+
255+
CompletableFuture<GetObjectResponse> crtFuture =
256+
s3AsyncClient.getObject(downloadRequest.getObjectRequest(), responseTransformer);
257+
258+
// Forward download cancellation to CRT future
259+
CompletableFutureUtils.forwardExceptionTo(returnFuture, crtFuture);
260+
261+
CompletableFutureUtils.forwardTransformedResultTo(crtFuture, returnFuture,
262+
res -> CompletedFileDownload.builder()
263+
.response(res)
264+
.build());
265+
} catch (Throwable throwable) {
266+
returnFuture.completeExceptionally(throwable);
267+
}
268+
}
269+
270+
/**
271+
* Attaches the execution attributes that tell the CRT-based client to write the response body straight to the destination
272+
* file, together with the progress listener and pause observable needed to report progress on and pause that transfer.
273+
*/
274+
private DownloadFileRequest crtDownloadFileRequest(DownloadFileRequest downloadRequest,
275+
S3MetaRequestPauseObservable observable,
276+
TransferProgressUpdater progressUpdater,
277+
ResponseFileOption responseFileOption) {
278+
GetObjectRequest getObjectRequest = attachSdkAttribute(
279+
downloadRequest.getObjectRequest(),
280+
b -> b.putExecutionAttribute(RESPONSE_FILE_PATH, downloadRequest.destination())
281+
.putExecutionAttribute(RESPONSE_FILE_OPTION, responseFileOption));
282+
283+
GetObjectRequest crtGetObjectRequest = attachCrtSdkAttribute(
284+
getObjectRequest,
285+
b -> b.put(METAREQUEST_PAUSE_OBSERVABLE, observable)
286+
.put(CRT_PROGRESS_LISTENER, progressUpdater.crtProgressListener()));
287+
288+
return downloadRequest.copy(r -> r.getObjectRequest(crtGetObjectRequest));
289+
}
290+
140291
private static ResumeToken crtResumeToken(ResumableFileUpload resumableFileUpload) {
141292
return new ResumeToken(new ResumeToken.PutResumeTokenBuilder()
142293
.withNumPartsCompleted(resumableFileUpload.transferredParts().orElse(0L))
@@ -166,4 +317,20 @@ private PutObjectRequest attachCrtSdkAttribute(PutObjectRequest putObjectRequest
166317
.overrideConfiguration(modifiedRequestOverrideConfig)
167318
.build();
168319
}
320+
321+
private GetObjectRequest attachCrtSdkAttribute(GetObjectRequest getObjectRequest,
322+
Consumer<SdkHttpExecutionAttributes.Builder> builderMutation) {
323+
SdkHttpExecutionAttributes existingAttributes =
324+
getObjectRequest.overrideConfiguration()
325+
.map(o -> o.executionAttributes().getAttribute(SDK_HTTP_EXECUTION_ATTRIBUTES))
326+
.orElse(null);
327+
328+
SdkHttpExecutionAttributes modifiedAttributes =
329+
(existingAttributes == null ? SdkHttpExecutionAttributes.builder() : existingAttributes.toBuilder())
330+
.applyMutation(builderMutation)
331+
.build();
332+
333+
return attachSdkAttribute(getObjectRequest,
334+
b -> b.putExecutionAttribute(SDK_HTTP_EXECUTION_ATTRIBUTES, modifiedAttributes));
335+
}
169336
}

0 commit comments

Comments
 (0)