Skip to content

enhancement: Forward 5XX 4XX from source to client #2031

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
entrop-tankos wants to merge 1 commit into
base: main
Choose a base branch
from
Open

enhancement: Forward 5XX 4XX from source to client #2031

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions dragonfly-client-backend/src/hdfs.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ impl Backend for Hdfs {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?
.into_iter()
Expand Down Expand Up @@ -177,6 +178,7 @@ impl Backend for Hdfs {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -194,6 +196,7 @@ impl Backend for Hdfs {
http_status_code: None,
error_message: None,
entries,
body: None,
})
}

Expand Down Expand Up @@ -227,6 +230,7 @@ impl Backend for Hdfs {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -244,6 +248,7 @@ impl Backend for Hdfs {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?,
None => operator_reader.into_bytes_stream(..).await.map_err(|err| {
Expand All @@ -256,6 +261,7 @@ impl Backend for Hdfs {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?,
};
Expand Down
89 changes: 64 additions & 25 deletions dragonfly-client-backend/src/http.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ use async_trait::async_trait;
use dashmap::{mapref::entry::Entry, DashMap};
use dragonfly_api::common::v2::Range;
use dragonfly_client_core::{
error::{ErrorType, OrErr},
error::{BackendError, ErrorType, OrErr},
Error, Result,
};
use dragonfly_client_util::{http::validate_ranged_response, tls::NoVerifier};
Expand Down Expand Up @@ -373,6 +373,16 @@ impl HTTP {
}
}

fn parse_content_range_total(headers: &HeaderMap) -> Option<u64> {
let content_range = headers.get(CONTENT_RANGE)?.to_str().ok()?;
// Expected format: "bytes <start>-<end>/<total>"
let (_, total) = content_range.split_once('/')?;
if total == "*" {
return None;
}
total.parse::<u64>().ok()
}

/// Implements the Backend trait.
#[async_trait]
impl Backend for HTTP {
Expand Down Expand Up @@ -479,6 +489,7 @@ impl Backend for HTTP {
http_status_code: None,
entries: Vec::new(),
error_message: Some(err.to_string()),
body: None,
});
}
}
Expand All @@ -497,6 +508,7 @@ impl Backend for HTTP {
error_message: Some(
"got 307 Temporary Redirect without Location header".to_string(),
),
body: None,
});
}
}
Expand Down Expand Up @@ -533,6 +545,7 @@ impl Backend for HTTP {
http_status_code: None,
entries: Vec::new(),
error_message: Some(err.to_string()),
body: None,
});
}
}
Expand Down Expand Up @@ -570,6 +583,7 @@ impl Backend for HTTP {
http_status_code: None,
entries: Vec::new(),
error_message: Some(err.to_string()),
body: None,
});
}
}
Expand All @@ -588,28 +602,21 @@ impl Backend for HTTP {
http_status_code: None,
entries: Vec::new(),
error_message: None,
body: None,
});
}
};

let response_status_code = response.status();
let mut response_header = response.headers().clone();
let content_length = if response_status_code == reqwest::StatusCode::PARTIAL_CONTENT {
// The total length of a ranged response is in the Content-Range header,
// e.g. "bytes 0-0/1048576".
let content_length = response_header
.get(CONTENT_RANGE)
.and_then(|content_range| content_range.to_str().ok())
.and_then(|content_range| content_range.rsplit_once('/'))
.and_then(|(_, total)| total.parse::<u64>().ok());

if response_status_code == reqwest::StatusCode::PARTIAL_CONTENT {
let content_length = parse_content_range_total(&response_header);
if content_length.is_none() {
error!(
"stat request got 206 Partial Content without valid Content-Range {} {}",
request.task_id, request_url
);
}

// Read the one-byte body to completion, so the connection can be reused by
// the connection pool instead of being closed with an unread body.
if let Err(err) = response.bytes().await {
Expand All @@ -618,39 +625,71 @@ impl Backend for HTTP {
request.task_id, request_url, err
);
}

// Rewrite the 206-shaped headers to look like the full-object response, since
// the response header is persisted as the task response header and echoed to
// the clients.
response_header.remove(CONTENT_RANGE);
if let Some(content_length) = content_length {
response_header.insert(CONTENT_LENGTH, HeaderValue::from(content_length));
}
debug!(
"stat response {} {}: {:?} {:?} {:?}",
request.task_id, request_url, response_status_code, content_length, response_header
);
return Ok(StatResponse {
success: true,
content_length,
http_header: Some(response_header),
http_status_code: Some(response_status_code),
entries: Vec::new(),
body: None,
error_message: Some(response_status_code.to_string()),
});
}

content_length
} else {
let content_length = match response_header.get(CONTENT_LENGTH) {
Some(content_length) => content_length.to_str()?.parse::<u64>().ok(),
None => response.content_length(),
};

// Drop the response body to avoid reading it.
drop(response);
content_length
let content_length = match response_header.get(CONTENT_LENGTH) {
Some(content_length) => content_length.to_str()?.parse::<u64>().ok(),
None => response.content_length(),
};

debug!(
"stat response {} {}: {:?} {:?} {:?}",
request.task_id, request_url, response_status_code, content_length, response_header
);

if !response_status_code.is_success() {
let body = response.bytes().await.map_err(|err| {
error!(
"stat request failed to read response body {} {}: {}",
request.task_id, request_url, err
);
Error::BackendError(Box::new(BackendError {
message: err.to_string(),
status_code: Some(response_status_code),
header: Some(response_header.clone()),
body: None,
}))
})?;

return Ok(StatResponse {
success: false,
content_length,
http_header: Some(response_header),
http_status_code: Some(response_status_code),
entries: Vec::new(),
body: Some(body.to_vec()),
error_message: Some(response_status_code.to_string()),
});
}

drop(response);
Ok(StatResponse {
success: response_status_code.is_success(),
success: true,
content_length,
http_header: Some(response_header),
http_status_code: Some(response_status_code),
error_message: Some(response_status_code.to_string()),
entries: Vec::new(),
body: None,
error_message: Some(response_status_code.to_string()),
})
}

Expand Down
8 changes: 8 additions & 0 deletions dragonfly-client-backend/src/hugging_face.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,7 @@ impl Backend for HuggingFace {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -435,6 +436,7 @@ impl Backend for HuggingFace {
message: response_status_code.to_string(),
status_code: Some(response_status_code),
header: Some(response_header),
body: None,
})));
}

Expand All @@ -454,6 +456,7 @@ impl Backend for HuggingFace {
http_status_code: Some(response_status_code),
error_message: Some(response_status_code.to_string()),
entries: Vec::new(),
body: None,
})
}
None => {
Expand All @@ -480,6 +483,7 @@ impl Backend for HuggingFace {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -500,6 +504,7 @@ impl Backend for HuggingFace {
message: response_status_code.to_string(),
status_code: Some(response_status_code),
header: Some(response_header),
body: None,
})));
}

Expand All @@ -513,6 +518,7 @@ impl Backend for HuggingFace {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -526,6 +532,7 @@ impl Backend for HuggingFace {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand Down Expand Up @@ -568,6 +575,7 @@ impl Backend for HuggingFace {
http_status_code: Some(response_status_code),
error_message: Some(response_status_code.to_string()),
entries,
body: None,
})
}
}
Expand Down
3 changes: 3 additions & 0 deletions dragonfly-client-backend/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,9 @@ pub struct StatResponse {
/// The information of the entries in the directory.
pub entries: Vec<DirEntry>,

/// The body of the non-success response.
pub body: Option<Vec<u8>>,

/// The error message of the response.
pub error_message: Option<String>,
}
Expand Down
9 changes: 9 additions & 0 deletions dragonfly-client-backend/src/model_scope.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -388,6 +388,7 @@ impl Backend for ModelScope {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -408,6 +409,7 @@ impl Backend for ModelScope {
message: response_status_code.to_string(),
status_code: Some(response_status_code),
header: Some(response_header),
body: None,
})));
}

Expand All @@ -428,6 +430,7 @@ impl Backend for ModelScope {
http_status_code: Some(response_status_code),
error_message: Some(response_status_code.to_string()),
entries: Vec::new(),
body: None,
})
}
None => {
Expand All @@ -451,6 +454,7 @@ impl Backend for ModelScope {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -471,6 +475,7 @@ impl Backend for ModelScope {
message: response_status_code.to_string(),
status_code: Some(response_status_code),
header: Some(response_header),
body: None,
})));
}

Expand All @@ -484,6 +489,7 @@ impl Backend for ModelScope {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -497,6 +503,7 @@ impl Backend for ModelScope {
message: err.to_string(),
status_code: None,
header: None,
body: None,
}))
})?;

Expand All @@ -509,6 +516,7 @@ impl Backend for ModelScope {
response.code,
response.message.unwrap_or_default()
),
body: None,
})));
}

Expand Down Expand Up @@ -547,6 +555,7 @@ impl Backend for ModelScope {
http_status_code: Some(response_status_code),
error_message: Some(response_status_code.to_string()),
entries,
body: None,
})
}
}
Expand Down
Loading
Loading