feat(redis): enhance connection handling with optional connection reuse
ZhenYi committed
Modified lib/infra/storage/src/s3.rs lib/infra/storage/src/s3.rs +58 −46 View lib/infra/storage/src/s3.rs
@@ -4,7 +4,7 @@
44use futures::StreamExt as _;
55use http::HeaderMap;
66use s3::Bucket;
7use s3::bucket::{CHUNK_SIZE, Part};
7use s3::bucket::CHUNK_SIZE;
88use s3::creds::Credentials;
99use s3::region::Region;
1010use s3::serde_types::{DateTime as S3DateTime, Object as S3Object};
@@ -167,56 +167,21 @@
167167 let path = msg.key;
168168 let upload_id = msg.upload_id;
169169
170 let mut chunk = first;
171 let mut part_number: u32 = 0;
172 let mut parts: Vec<Part> = Vec::new();
173 let mut total: u64 = 0;
174 loop {
175 total += chunk.len() as u64;
176 part_number += 1;
177 match self
178 .bucket
170 let uploaded = upload_parts(reader, first, async |chunk, part_number| {
171 self.bucket
179172 .put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)
180173 .await
181 {
182 Ok(part) => parts.push(part),
183 Err(e) => {
184 abort_upload(&self.bucket, &path, &upload_id).await;
185 return Err(StorageError::from(e));
186 }
187 }
174 .map_err(StorageError::from)
175 })
176 .await;
188177
189 chunk = match read_chunk(reader).await {
190 Ok(c) => c,
191 Err(e) => {
192 abort_upload(&self.bucket, &path, &upload_id).await;
193 return Err(e);
194 }
195 };
196 if chunk.is_empty() {
178 let (total, parts) = match uploaded {
179 Ok(result) => result,
180 Err(error) => {
197181 abort_upload(&self.bucket, &path, &upload_id).await;
198 return Err(StorageError::UnexpectedStatus {
199 operation: "put_stream",
200 status: 400,
201 });
182 return Err(error);
202183 }
203 if chunk.len() < CHUNK_SIZE {
204 total += chunk.len() as u64;
205 part_number += 1;
206 match self
207 .bucket
208 .put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)
209 .await
210 {
211 Ok(part) => parts.push(part),
212 Err(e) => {
213 abort_upload(&self.bucket, &path, &upload_id).await;
214 return Err(StorageError::from(e));
215 }
216 }
217 break;
218 }
219 }
184 };
220185
221186 match self
222187 .bucket
@@ -423,6 +388,33 @@
423388 Ok(chunk)
424389}
425390
391/// Upload the first chunk plus every following chunk as a multipart part.
392/// A zero-byte read is EOF and simply ends the upload, so a body whose size is
393/// an exact multiple of `CHUNK_SIZE` completes instead of being aborted.
394async fn upload_parts<R, F, T>(
395 reader: &mut R,
396 first: Vec<u8>,
397 mut put_part: F,
398) -> Result<(u64, Vec<T>), StorageError>
399where
400 R: AsyncRead + Unpin,
401 F: AsyncFnMut(Vec<u8>, u32) -> Result<T, StorageError>,
402{
403 let mut chunk = first;
404 let mut part_number: u32 = 0;
405 let mut parts = Vec::new();
406 let mut total: u64 = 0;
407 loop {
408 total += chunk.len() as u64;
409 part_number += 1;
410 parts.push(put_part(chunk, part_number).await?);
411 chunk = read_chunk(reader).await?;
412 if chunk.is_empty() {
413 return Ok((total, parts));
414 }
415 }
416}
417
426418/// Fire-and-forget multipart abort (errors are intentionally swallowed: the
427419/// original failure is what the caller cares about).
428420async fn abort_upload(bucket: &Bucket, path: &str, upload_id: &str) {
@@ -452,4 +444,24 @@
452444 let first = read_chunk(&mut cursor).await.expect("read");
453445 assert_eq!(first.len(), CHUNK_SIZE - 1);
454446 }
447
448 #[tokio::test]
449 async fn multipart_upload_handles_exact_chunk_multiple() {
450 let data = vec![7u8; CHUNK_SIZE * 2];
451 let mut cursor = io::Cursor::new(&data);
452 let first = read_chunk(&mut cursor).await.expect("read");
453 assert_eq!(first.len(), CHUNK_SIZE);
454
455 let mut uploaded: Vec<usize> = Vec::new();
456 let (total, part_sizes) = upload_parts(&mut cursor, first, async |chunk, _part_number| {
457 uploaded.push(chunk.len());
458 Ok::<usize, StorageError>(chunk.len())
459 })
460 .await
461 .expect("an exact chunk multiple must complete");
462
463 assert_eq!(uploaded, vec![CHUNK_SIZE, CHUNK_SIZE]);
464 assert_eq!(part_sizes, vec![CHUNK_SIZE, CHUNK_SIZE]);
465 assert_eq!(total, (CHUNK_SIZE * 2) as u64);
466 }
455467}