use std::io::{Error, ErrorKind, Result};
use log::{error, info, warn};
use bytes::Buf;
use aws_sdk_s3::Client;
use aws_sdk_s3::error::{SdkError, ProvideErrorMetadata};
use aws_sdk_s3::primitives::SdkBody;
use aws_sdk_s3::operation::head_object::HeadObjectOutput;
use aws_sdk_s3::types::{Object, ObjectIdentifier};
#[cfg(feature = "wal")]
use aws_sdk_s3::types::CommonPrefix;
use crate::inode::OnDiskState;
pub fn is_condition_failure<E: ProvideErrorMetadata>(err: &SdkError<E>) -> bool {
matches!(err.raw_response().map(|r| r.status().as_u16()), Some(409 | 412))
|| matches!(err.code(), Some("PreconditionFailed" | "ConditionalRequestConflict"))
}
pub fn s3_error_kind<E: ProvideErrorMetadata>(err: &SdkError<E>) -> ErrorKind {
if is_condition_failure(err) {
return ErrorKind::AlreadyExists;
}
match err {
SdkError::TimeoutError(_) => ErrorKind::TimedOut,
SdkError::DispatchFailure(_) => ErrorKind::ConnectionReset,
SdkError::ConstructionFailure(_) => ErrorKind::InvalidInput,
SdkError::ResponseError(_) => ErrorKind::ConnectionReset,
_ => match err.raw_response().map(|r| r.status().as_u16()) {
Some(404) => ErrorKind::NotFound,
Some(403) => ErrorKind::PermissionDenied,
Some(408) => ErrorKind::TimedOut,
Some(409 | 412) => ErrorKind::AlreadyExists,
Some(429 | 500 | 502 | 503 | 504) => ErrorKind::ResourceBusy,
Some(400 | 405 | 411 | 416 | 501) => ErrorKind::InvalidInput,
_ => ErrorKind::Other,
},
}
}
pub fn is_retryable(kind: ErrorKind) -> bool {
matches!(kind, ErrorKind::ConnectionReset | ErrorKind::TimedOut | ErrorKind::ResourceBusy)
}
pub fn s3_code_kind(code: Option<&str>) -> ErrorKind {
match code {
Some("AccessDenied" | "AllAccessDisabled" | "AccountProblem" | "InvalidAccessKeyId"
| "SignatureDoesNotMatch" | "ExpiredToken" | "InvalidToken") => ErrorKind::PermissionDenied,
Some("NoSuchKey" | "NoSuchBucket" | "NoSuchVersion" | "NoSuchUpload") => ErrorKind::NotFound,
Some("PreconditionFailed" | "ConditionalRequestConflict" | "OperationAborted") => ErrorKind::AlreadyExists,
Some("SlowDown" | "InternalError" | "ServiceUnavailable" | "Throttling" | "ThrottlingException"
| "RequestLimitExceeded" | "TooManyRequests" | "BadGateway" | "GatewayTimeout") => ErrorKind::ResourceBusy,
Some("RequestTimeout") => ErrorKind::TimedOut,
Some("InvalidArgument" | "InvalidRequest" | "MalformedXML" | "KeyTooLongError"
| "NotImplemented") => ErrorKind::InvalidInput,
_ => ErrorKind::Other,
}
}
fn body_error(e: aws_sdk_s3::primitives::ByteStreamError, bucket: &str, key: &str, range: Option<&str>) -> Error {
let what = match range {
Some(r) => format!("GetObject s3://{}/{} by range {}", bucket, key, r),
None => format!("GetObject s3://{}/{}", bucket, key),
};
let msg = format!("{}: reading the response failed: {}", what, std::error::Error::source(&e)
.map(|s| format!("{} ({})", e, s)).unwrap_or_else(|| e.to_string()));
warn!("{}", msg);
Error::new(ErrorKind::ConnectionReset, msg)
}
const WRITE_MARK: &str = "hf-write";
fn write_mark() -> String {
ulid::Ulid::new().to_string()
}
pub(crate) struct S3Ops;
impl S3Ops {
async fn own_write_or(client: &Client, bucket: &str, key: &str, mark: &str, conflict: Error)
-> Result<Option<OnDiskState>>
{
match client.head_object().bucket(bucket).key(key).send().await {
Ok(head) => {
if head.metadata().and_then(|m| m.get(WRITE_MARK)).map(String::as_str) == Some(mark) {
info!("a conditional write to s3://{}/{} was refused after its own earlier \
attempt landed, its answer lost; it succeeded", bucket, key);
return Ok(head.e_tag.map(|t| OnDiskState { checksum: t.replace("\"", ""), timestamp: 0 }));
}
Err(conflict)
},
Err(sdk_err) if sdk_err.as_service_error().is_some_and(|e| e.is_not_found()) => Err(conflict),
Err(sdk_err) => {
let err_str = format!("{}; and whether it was this write's own earlier attempt \
could not be told, HeadObject failed: {}", conflict, sdk_err);
error!("{}", err_str);
Err(Error::new(s3_error_kind(&sdk_err), err_str))
},
}
}
#[cfg(feature = "wal")]
pub(crate) async fn do_list_directory(client: &Client, bucket: &str, prefix: &str, mut f: impl FnMut(&CommonPrefix)) -> Result<()> {
let mut stream = client
.list_objects_v2()
.bucket(bucket)
.prefix(prefix)
.delimiter("/")
.into_paginator()
.send();
while let Some(page) = stream.next().await {
match page {
Ok(list_res) => {
if let Some(prefixes) = list_res.common_prefixes {
prefixes.iter().for_each(|p| f(p))
}
},
Err(sdk_err) => {
let mut err_str = format!("ListObjectV2 s3://{}/{} error: ", bucket, prefix);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
},
}
}
Ok(())
}
pub(crate) async fn do_list_objects(client: &Client, bucket: &str, prefix: &str, mut f: impl FnMut(&Object)) -> Result<()> {
let mut stream = client
.list_objects_v2()
.bucket(bucket)
.prefix(prefix)
.into_paginator()
.send();
while let Some(page) = stream.next().await {
match page {
Ok(list_res) => {
if let Some(objects) = list_res.contents {
objects.iter().for_each(|obj| f(obj))
}
},
Err(sdk_err) => {
let mut err_str = format!("ListObjectV2 s3://{}/{} error: ", bucket, prefix);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
},
}
}
Ok(())
}
pub(crate) async fn do_head_object(client: &Client, bucket: &str, key: &str) -> Result<HeadObjectOutput> {
let res = client
.head_object()
.bucket(bucket)
.key(key)
.send()
.await;
match res {
Ok(output) => {
return Ok(output);
},
Err(sdk_err) => {
let mut err_str = format!("HeadObject s3://{}/{} error: ", bucket, key);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
if sdk_err.as_service_error().map(|e| e.is_not_found()) == Some(true) {
warn!("{}", err_str);
return Err(Error::new(ErrorKind::NotFound, err_str));
}
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
}
}
}
pub(crate) async fn do_delete_object(client: &Client, bucket: &str, key: &str,
inode_state: &Option<OnDiskState>) -> Result<()>
{
let builder = client
.delete_object()
.bucket(bucket)
.key(key);
let op = if let Some(state) = inode_state {
builder
.if_match(state.checksum.as_str())
.if_match_last_modified_time(aws_sdk_s3::primitives::DateTime::from_secs(state.timestamp))
} else {
builder
};
match op.send().await {
Ok(_) => {
},
Err(sdk_err) => {
if let Some(resp) = sdk_err.raw_response() {
if resp.status().as_u16() == 412 {
let err_str = format!("Conditional DeleteObject failed on s3://{}/{}, status: 412 (concurrent modification)", bucket, key);
warn!("{}", err_str);
return Err(Error::new(ErrorKind::AlreadyExists, err_str));
}
}
let mut err_str = format!("DeleteObject s3://{}/{} error: ", bucket, key);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
}
}
Ok(())
}
pub(crate) async fn do_get_object_speculative(client: &Client, bucket: &str, key: &str,
range: Option<&str>, with_ods: bool) -> Result<(bytes::Bytes, Option<OnDiskState>)>
{
let builder = client
.get_object()
.bucket(bucket)
.key(key);
let op = if let Some(r) = range {
builder.range(r)
} else {
builder
};
match op.send().await {
Ok(output) => {
let bytes = output.body.collect().await.map_err(|e| body_error(e, bucket, key, range))?;
if with_ods {
let od_state = OnDiskState {
checksum: output.e_tag.unwrap().replace("\"", ""),
timestamp: output.last_modified.unwrap().secs(),
};
return Ok((bytes.into_bytes(), Some(od_state)));
}
return Ok((bytes.into_bytes(), None));
},
Err(sdk_err) => {
let mut err_str = if let Some(r) = range {
format!("GetObject s3://{}/{} by range {} error: ", bucket, key, r)
} else {
format!("GetObject s3://{}/{} error: ", bucket, key)
};
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
if sdk_err.as_service_error().map(|e| e.is_no_such_key()) == Some(true) {
warn!("{}", err_str);
return Err(Error::new(ErrorKind::NotFound, err_str));
}
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
},
}
}
pub(crate) async fn do_get_object(client: &Client, bucket: &str, key: &str,
buf: &mut [u8], range: Option<&str>, with_ods: bool) -> Result<Option<OnDiskState>>
{
let builder = client
.get_object()
.bucket(bucket)
.key(key);
let op = if let Some(r) = range {
builder.range(r)
} else {
builder
};
match op.send().await {
Ok(output) => {
let mut bytes = output.body.collect().await.map_err(|e| body_error(e, bucket, key, range))?;
if bytes.remaining() < buf.len() {
let err_str = format!("GetObject s3://{}/{} feched size {} less than input buffer size {}",
bucket, key, bytes.remaining(), buf.len());
error!("{}", err_str);
return Err(Error::new(ErrorKind::InvalidData, err_str));
}
bytes.copy_to_slice(buf);
if with_ods {
let od_state = OnDiskState {
checksum: output.e_tag.unwrap().replace("\"", ""),
timestamp: output.last_modified.unwrap().secs(),
};
return Ok(Some(od_state));
}
return Ok(None);
},
Err(sdk_err) => {
let mut err_str = if let Some(r) = range {
format!("GetObject s3://{}/{} by range {} error: ", bucket, key, r)
} else {
format!("GetObject s3://{}/{} error: ", bucket, key)
};
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
if sdk_err.as_service_error().map(|e| e.is_no_such_key()) == Some(true) {
warn!("{}", err_str);
return Err(Error::new(ErrorKind::NotFound, err_str));
}
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
},
}
}
pub(crate) async fn do_put_object(client: &Client, bucket: &str, key: &str,
buf: &[u8], inode_state: &Option<OnDiskState>) -> Result<Option<OnDiskState>>
{
let body = SdkBody::from(buf);
let mark = write_mark();
let builder = client
.put_object()
.bucket(bucket)
.key(key)
.metadata(WRITE_MARK, &mark)
.body(body.into());
let op = if let Some(state) = inode_state {
builder
.if_match(state.checksum.as_str())
} else {
builder
.if_none_match('*')
};
match op.send().await {
Ok(output) => {
let od_state = OnDiskState {
checksum: output.e_tag.unwrap().replace("\"", ""),
timestamp: 0,
};
return Ok(Some(od_state));
},
Err(sdk_err) => {
if let Some(resp) = sdk_err.raw_response() {
match resp.status().as_u16() {
412 | 409 => {
let err_str = format!("Conditional PutObject failed on s3://{}/{}, status: {} (concurrent modification)", bucket, key, resp.status().as_u16());
let conflict = Error::new(ErrorKind::AlreadyExists, err_str);
let res = Self::own_write_or(client, bucket, key, &mark, conflict).await;
if let Err(e) = &res { warn!("{}", e); }
return res;
},
_ => {},
}
}
let mut err_str = format!("PutObject s3://{}/{} error: ", bucket, key);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
}
}
}
async fn create_mp_upload(client: &Client, bucket: &str, key: &str, mark: &str) -> Result<String> {
match client.create_multipart_upload().bucket(bucket).key(key).metadata(WRITE_MARK, mark).send().await {
Ok(output) => output.upload_id.ok_or_else(|| {
let err_str = format!("CreateMultipartUpload s3://{}/{} error: unable to get a valid upload id", bucket, key);
error!("{}", err_str);
Error::new(ErrorKind::InvalidData, err_str)
}),
Err(sdk_err) => {
let mut err_str = format!("CreateMultipartUpload s3://{}/{} error: ", bucket, key);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
error!("{}", err_str);
Err(Error::new(s3_error_kind(&sdk_err), err_str))
}
}
}
async fn abort_mp_upload(client: &Client, bucket: &str, key: &str, upload_id: &str, e: Error) -> Error {
let note = match client.abort_multipart_upload().bucket(bucket).key(key).upload_id(upload_id).send().await {
Ok(_) => format!("multipart upload {} aborted", upload_id),
Err(sdk_err) if sdk_err.as_service_error().is_some_and(|s| s.is_no_such_upload()) =>
format!("multipart upload {} was already gone", upload_id),
Err(sdk_err) => {
let msg = match sdk_err.as_service_error() {
Some(serv_err) => format!("{}", serv_err),
None => format!("{}", sdk_err),
};
warn!("AbortMultipartUpload s3://{}/{} upload_id: {}, error: {}; its parts stay \
until a lifecycle rule for incomplete uploads removes them", bucket, key, upload_id, msg);
format!("aborting multipart upload {} failed as well ({}), so its parts stay until \
a lifecycle rule for incomplete uploads removes them", upload_id, msg)
},
};
Error::new(e.kind(), format!("{}; {}", e, note))
}
pub(crate) async fn do_mp_upload(client: &Client, bucket: &str, key: &str,
buf: &[u8], inode_state: &Option<OnDiskState>, mpu_chunk_size: usize, concurrency: usize) -> Result<Option<OnDiskState>>
{
let mark = write_mark();
let upload_id = Self::create_mp_upload(client, bucket, key, &mark).await?;
match Self::mp_upload_and_complete(client, bucket, key, &upload_id, &mark, buf, inode_state, mpu_chunk_size, concurrency).await {
Ok(state) => Ok(state),
Err(e) => Err(Self::abort_mp_upload(client, bucket, key, &upload_id, e).await),
}
}
async fn mp_upload_and_complete(client: &Client, bucket: &str, key: &str, upload_id: &str, mark: &str,
buf: &[u8], inode_state: &Option<OnDiskState>, mpu_chunk_size: usize, concurrency: usize) -> Result<Option<OnDiskState>>
{
let upload_id = upload_id.to_owned();
let parts = (buf.len() / mpu_chunk_size) + 1;
let concurrency = concurrency.max(1);
let mut complete_parts = Vec::with_capacity(parts);
let mut set: tokio::task::JoinSet<Result<(usize, String)>> = tokio::task::JoinSet::new();
let mut next = 1;
loop {
while next <= parts && set.len() < concurrency {
let part_id = next;
next += 1;
let part_data = if part_id == parts {
&buf[(part_id - 1) * mpu_chunk_size..]
} else {
&buf[(part_id - 1) * mpu_chunk_size..part_id * mpu_chunk_size]
};
let part_data = unsafe {
std::slice::from_raw_parts(part_data.as_ptr() as *const u8, part_data.len())
};
let uid = upload_id.clone();
let k = key.to_owned();
let c = client.clone();
let b = bucket.to_owned();
set.spawn(async move {
let etag = Self::do_upload_parts(&c, &b, &k, part_data, part_id, &uid).await?;
Ok((part_id, etag))
});
}
match set.join_next().await {
None => break,
Some(Ok(Ok((part_id, etag)))) => complete_parts.push((part_id as i32, etag)),
Some(Ok(Err(e))) => { set.shutdown().await; return Err(e); },
Some(Err(join_err)) => { set.shutdown().await; return Err(Error::from(join_err)); },
}
}
complete_parts.sort_by_key(|t| t.0);
Self::complete_mp_upload(client, bucket, key, &upload_id, mark, complete_parts, inode_state, "").await
}
async fn complete_mp_upload(client: &Client, bucket: &str, key: &str, upload_id: &str, mark: &str,
parts: Vec<(i32, String)>, inode_state: &Option<OnDiskState>, what: &str) -> Result<Option<OnDiskState>>
{
let parts = parts.into_iter().map(|p|
aws_sdk_s3::types::builders::CompletedPartBuilder::default()
.part_number(p.0)
.e_tag(p.1)
.build()
).collect();
let completed = aws_sdk_s3::types::builders::CompletedMultipartUploadBuilder::default()
.set_parts(Some(parts))
.build();
let builder = client
.complete_multipart_upload()
.bucket(bucket)
.key(key)
.multipart_upload(completed)
.upload_id(upload_id);
let op = if let Some(state) = inode_state {
builder
.if_match(state.checksum.as_str())
} else {
builder
.if_none_match('*')
};
match op.send().await {
Ok(output) => {
let od_state = output.e_tag.map(|t| OnDiskState {
checksum: t.replace("\"", ""),
timestamp: 0,
});
Ok(od_state)
},
Err(sdk_err) => {
if is_condition_failure(&sdk_err) {
let err_str = format!("Conditional CompleteMultipartUpload{} failed on s3://{}/{}, {} (concurrent modification)",
what, bucket, key, sdk_err.code().unwrap_or("conflict"));
let conflict = Error::new(ErrorKind::AlreadyExists, err_str);
let res = Self::own_write_or(client, bucket, key, mark, conflict).await;
if let Err(e) = &res { warn!("{}", e); }
return res;
}
let mut err_str = format!("CompleteMultipartUpload{} s3://{}/{} upload_id: {}, error: ",
what, bucket, key, upload_id);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
error!("{}", err_str);
Err(Error::new(s3_error_kind(&sdk_err), err_str))
}
}
}
async fn do_upload_parts(client: &Client, bucket: &str, key: &str,
buf: &[u8], part_id: usize, upload_id: &str) -> Result<String>
{
let body = SdkBody::from(buf);
let res = client
.upload_part()
.bucket(bucket)
.key(key)
.part_number(part_id as i32)
.upload_id(upload_id)
.body(body.into())
.send()
.await;
match res {
Ok(output) => {
return Ok(output.e_tag.unwrap());
},
Err(sdk_err) => {
let mut err_str = format!("UploadPart s3://{}/{} upload_id: {}, part_id: {}, error: ",
bucket, key, upload_id, part_id);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
};
error!("{}", err_str);
return Err(Error::new(s3_error_kind(&sdk_err), err_str));
},
}
}
pub(crate) async fn do_delete_objects(client: &Client, bucket: &str, delete_keys: Vec<String>) -> Result<()> {
let mut first: Option<(ErrorKind, String)> = None;
let mut failed = 0usize;
for keys in delete_keys.chunks(1000) {
let obj_ids = keys.iter().map(|k| {
ObjectIdentifier::builder()
.key(k.to_string())
.build()
.unwrap()
}).collect::<Vec<ObjectIdentifier>>();
let delete = aws_sdk_s3::types::Delete::builder()
.set_objects(Some(obj_ids))
.quiet(true)
.build()
.unwrap();
match client.delete_objects()
.bucket(bucket)
.delete(delete)
.send()
.await
{
Ok(output) => {
for e in output.errors() {
failed += 1;
let what = format!("DeleteObjects s3://{}/{}: {} {}", bucket,
e.key().unwrap_or("<no key>"), e.code().unwrap_or("<no code>"), e.message().unwrap_or(""));
error!("{}", what);
first.get_or_insert((s3_code_kind(e.code()), what));
}
},
Err(sdk_err) => {
failed += keys.len();
let what = format!("DeleteObjects in s3://{} of {} keys: {}", bucket, keys.len(),
sdk_err.as_service_error().map(|e| e.to_string()).unwrap_or_else(|| sdk_err.to_string()));
error!("{}", what);
first.get_or_insert((s3_error_kind(&sdk_err), what));
}
}
}
if let Some((kind, what)) = first {
return Err(Error::new(kind, format!("{} of {} keys were not deleted; the first: {}",
failed, delete_keys.len(), what)));
}
Ok(())
}
pub(crate) async fn do_put_object_pieces(
client: &Client,
bucket: &str,
key: &str,
body: crate::segment_body::SegmentBody,
inode_state: &Option<OnDiskState>,
) -> Result<Option<OnDiskState>> {
let sdk_body = body.materialize().await?.into_sdk_body();
let mark = write_mark();
let builder = client
.put_object()
.bucket(bucket)
.key(key)
.metadata(WRITE_MARK, &mark)
.body(sdk_body.into());
let op = if let Some(state) = inode_state {
builder.if_match(state.checksum.as_str())
} else {
builder.if_none_match('*')
};
match op.send().await {
Ok(output) => {
let od_state = OnDiskState {
checksum: output.e_tag.unwrap().replace("\"", ""),
timestamp: 0,
};
Ok(Some(od_state))
}
Err(sdk_err) => {
if let Some(resp) = sdk_err.raw_response() {
match resp.status().as_u16() {
412 | 409 => {
let err_str = format!(
"Conditional PutObject (pieces) failed on s3://{}/{}, status: {} (concurrent modification)",
bucket, key, resp.status().as_u16(),
);
let conflict = Error::new(ErrorKind::AlreadyExists, err_str);
let res = Self::own_write_or(client, bucket, key, &mark, conflict).await;
if let Err(e) = &res { warn!("{}", e); }
return res;
}
_ => {}
}
}
let mut err_str = format!("PutObject (pieces) s3://{}/{} error: ", bucket, key);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
}
error!("{}", err_str);
Err(Error::new(s3_error_kind(&sdk_err), err_str))
}
}
}
pub(crate) async fn do_mp_upload_pieces(
client: &Client,
bucket: &str,
key: &str,
body: crate::segment_body::SegmentBody,
inode_state: &Option<OnDiskState>,
mpu_chunk_size: usize,
concurrency: usize,
) -> Result<Option<OnDiskState>> {
let mark = write_mark();
let upload_id = Self::create_mp_upload(client, bucket, key, &mark).await?;
match Self::mp_upload_and_complete_pieces(client, bucket, key, &upload_id, &mark, body, inode_state, mpu_chunk_size, concurrency).await {
Ok(state) => Ok(state),
Err(e) => Err(Self::abort_mp_upload(client, bucket, key, &upload_id, e).await),
}
}
async fn mp_upload_and_complete_pieces(
client: &Client,
bucket: &str,
key: &str,
upload_id: &str,
mark: &str,
body: crate::segment_body::SegmentBody,
inode_state: &Option<OnDiskState>,
mpu_chunk_size: usize,
concurrency: usize,
) -> Result<Option<OnDiskState>> {
let total_len = body.len();
let upload_id = upload_id.to_owned();
let parts = total_len.div_ceil(mpu_chunk_size).max(1);
let concurrency = concurrency.max(1);
let mut complete_parts: Vec<(i32, String)> = Vec::with_capacity(parts);
let mut set: tokio::task::JoinSet<Result<(usize, String)>> = tokio::task::JoinSet::new();
let mut next = 1;
loop {
while next <= parts && set.len() < concurrency {
let part_id = next;
next += 1;
let start = (part_id - 1) * mpu_chunk_size;
let part_len = if part_id == parts {
total_len - start
} else {
mpu_chunk_size
};
let sub = body.slice(start, part_len);
let uid = upload_id.clone();
let k = key.to_owned();
let c = client.clone();
let b = bucket.to_owned();
set.spawn(async move {
let etag = Self::do_upload_parts_pieces(&c, &b, &k, sub, part_id, &uid).await?;
Ok((part_id, etag))
});
}
match set.join_next().await {
None => break,
Some(Ok(Ok((pid, etag)))) => complete_parts.push((pid as i32, etag)),
Some(Ok(Err(e))) => { set.shutdown().await; return Err(e); },
Some(Err(join_err)) => { set.shutdown().await; return Err(Error::from(join_err)); },
}
}
complete_parts.sort_by_key(|t| t.0);
Self::complete_mp_upload(client, bucket, key, &upload_id, mark, complete_parts, inode_state, " (pieces)").await
}
async fn do_upload_parts_pieces(
client: &Client,
bucket: &str,
key: &str,
body: crate::segment_body::SegmentBody,
part_id: usize,
upload_id: &str,
) -> Result<String> {
let sdk_body = body.materialize().await?.into_sdk_body();
let res = client
.upload_part()
.bucket(bucket)
.key(key)
.part_number(part_id as i32)
.upload_id(upload_id)
.body(sdk_body.into())
.send()
.await;
match res {
Ok(output) => Ok(output.e_tag.unwrap()),
Err(sdk_err) => {
let mut err_str = format!(
"UploadPart (pieces) s3://{}/{} upload_id: {}, part_id: {}, error: ",
bucket, key, upload_id, part_id,
);
if let Some(serv_err) = sdk_err.as_service_error() {
err_str.push_str(&format!("{}", serv_err));
} else {
err_str.push_str(&format!("{}", sdk_err));
}
error!("{}", err_str);
Err(Error::new(s3_error_kind(&sdk_err), err_str))
}
}
}
}