Add exact resumable HTTP source streaming
This commit is contained in:
parent
3041257d78
commit
5b43b6c660
1 changed files with 1361 additions and 5 deletions
335
crates/remora-http-source/src/source.rs
Normal file
335
crates/remora-http-source/src/source.rs
Normal file
|
|
@ -0,0 +1,335 @@
|
|||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use http::{HeaderMap, HeaderValue, StatusCode, header};
|
||||
use remora_domain::MemberPath;
|
||||
use remora_outbound::{
|
||||
CanonicalUrl, OutboundClient, OutboundError, OutboundPolicy, OutboundRequest, Resolver,
|
||||
};
|
||||
use remora_stream::{IntegrityDigest, RetryTiming, capped_server_hint};
|
||||
use tokio::time::timeout;
|
||||
|
||||
use crate::{
|
||||
HttpResumeCheckpoint, HttpSourceBody, HttpSourceError, HttpSourceIdentity, StrongEtag,
|
||||
};
|
||||
|
||||
pub struct HttpSourceRequest {
|
||||
url: CanonicalUrl,
|
||||
result_path: MemberPath,
|
||||
credential_headers: HeaderMap,
|
||||
}
|
||||
|
||||
impl HttpSourceRequest {
|
||||
pub fn new(url: &str, result_path: &str) -> Result<Self, HttpSourceError> {
|
||||
Ok(Self {
|
||||
url: CanonicalUrl::parse(url).map_err(|_| HttpSourceError::InvalidRequest)?,
|
||||
result_path: MemberPath::parse(result_path)
|
||||
.map_err(|_| HttpSourceError::InvalidRequest)?,
|
||||
credential_headers: HeaderMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_credential_headers(mut self, headers: HeaderMap) -> Result<Self, HttpSourceError> {
|
||||
OutboundRequest::get(self.url.as_url().as_str())
|
||||
.map_err(|_| HttpSourceError::InvalidRequest)?
|
||||
.with_credential_headers(headers.clone())
|
||||
.map_err(|_| HttpSourceError::InvalidRequest)?;
|
||||
self.credential_headers = headers;
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for HttpSourceRequest {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("HttpSourceRequest")
|
||||
.field("result_path", &self.result_path.canonical())
|
||||
.field("credential_header_count", &self.credential_headers.len())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct OpenedHttpSource {
|
||||
result_path: MemberPath,
|
||||
identity: Option<HttpSourceIdentity>,
|
||||
body: HttpSourceBody,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for OpenedHttpSource {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("OpenedHttpSource")
|
||||
.field("result_path", &self.result_path.canonical())
|
||||
.field("has_resume_identity", &self.identity.is_some())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl OpenedHttpSource {
|
||||
#[must_use]
|
||||
pub fn result_path(&self) -> &MemberPath {
|
||||
&self.result_path
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn identity(&self) -> Option<&HttpSourceIdentity> {
|
||||
self.identity.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn into_body(self) -> HttpSourceBody {
|
||||
self.body
|
||||
}
|
||||
}
|
||||
|
||||
pub struct HttpSource<R> {
|
||||
client: OutboundClient<R>,
|
||||
timing: RetryTiming,
|
||||
}
|
||||
|
||||
impl<R: Resolver> HttpSource<R> {
|
||||
#[must_use]
|
||||
pub fn new(resolver: R, policy: OutboundPolicy, timing: RetryTiming) -> Self {
|
||||
let timing = timing.sanitized();
|
||||
Self {
|
||||
client: OutboundClient::builder(resolver, policy)
|
||||
.connect_timeout(timing.connect_timeout)
|
||||
.build(),
|
||||
timing,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn open(
|
||||
&self,
|
||||
request: HttpSourceRequest,
|
||||
) -> Result<OpenedHttpSource, HttpSourceError> {
|
||||
let requested_url = request.url.clone();
|
||||
let outbound = build_request(&request.url, request.credential_headers, None)?;
|
||||
let response = self.execute(outbound).await?;
|
||||
let status = response.response().status();
|
||||
if status != StatusCode::OK {
|
||||
return Err(status_error(
|
||||
status,
|
||||
response.response().headers(),
|
||||
&self.timing,
|
||||
));
|
||||
}
|
||||
ensure_identity_encoding(response.response().headers())?;
|
||||
let expected_length = content_length(response.response().headers())?;
|
||||
let identity = resumable_identity(
|
||||
requested_url,
|
||||
response.final_url().clone(),
|
||||
response.response().headers(),
|
||||
expected_length,
|
||||
);
|
||||
Ok(OpenedHttpSource {
|
||||
result_path: request.result_path,
|
||||
identity,
|
||||
body: HttpSourceBody::new(response.into_response(), expected_length),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn resume(
|
||||
&self,
|
||||
checkpoint: &HttpResumeCheckpoint,
|
||||
observed_prefix: &IntegrityDigest,
|
||||
credential_headers: HeaderMap,
|
||||
) -> Result<OpenedHttpSource, HttpSourceError> {
|
||||
if checkpoint.prefix_integrity() != observed_prefix {
|
||||
return Err(HttpSourceError::SourceChanged);
|
||||
}
|
||||
let outbound = build_request(
|
||||
checkpoint.identity().requested_url(),
|
||||
credential_headers,
|
||||
Some(checkpoint),
|
||||
)?;
|
||||
let response = self.execute(outbound).await?;
|
||||
if response.response().status() != StatusCode::PARTIAL_CONTENT {
|
||||
return Err(match response.response().status() {
|
||||
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => HttpSourceError::Authorization,
|
||||
StatusCode::NOT_FOUND | StatusCode::GONE => HttpSourceError::NotFound,
|
||||
status if status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS => {
|
||||
status_error(status, response.response().headers(), &self.timing)
|
||||
}
|
||||
_ => HttpSourceError::SourceChanged,
|
||||
});
|
||||
}
|
||||
validate_resume_response(checkpoint, &response)?;
|
||||
let remaining = checkpoint.identity().representation_length() - checkpoint.offset();
|
||||
Ok(OpenedHttpSource {
|
||||
result_path: checkpoint.result_path().clone(),
|
||||
identity: Some(checkpoint.identity().clone()),
|
||||
body: HttpSourceBody::new(response.into_response(), Some(remaining)),
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
request: OutboundRequest,
|
||||
) -> Result<remora_outbound::OutboundResponse, HttpSourceError> {
|
||||
match timeout(
|
||||
self.timing.response_start_timeout,
|
||||
self.client.execute(request),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(response)) => Ok(response),
|
||||
Ok(Err(error)) => Err(map_outbound_error(error)),
|
||||
Err(_) => Err(HttpSourceError::Transient { retry_after: None }),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_request(
|
||||
url: &CanonicalUrl,
|
||||
credentials: HeaderMap,
|
||||
checkpoint: Option<&HttpResumeCheckpoint>,
|
||||
) -> Result<OutboundRequest, HttpSourceError> {
|
||||
let mut public = HeaderMap::new();
|
||||
public.insert(
|
||||
header::ACCEPT_ENCODING,
|
||||
HeaderValue::from_static("identity"),
|
||||
);
|
||||
if let Some(checkpoint) = checkpoint {
|
||||
public.insert(
|
||||
header::RANGE,
|
||||
HeaderValue::from_str(&format!("bytes={}-", checkpoint.offset()))
|
||||
.map_err(|_| HttpSourceError::InvalidRequest)?,
|
||||
);
|
||||
public.insert(
|
||||
header::IF_MATCH,
|
||||
HeaderValue::from_str(checkpoint.identity().etag().as_str())
|
||||
.map_err(|_| HttpSourceError::InvalidRequest)?,
|
||||
);
|
||||
}
|
||||
OutboundRequest::get(url.as_url().as_str())
|
||||
.and_then(|request| request.with_public_headers(public))
|
||||
.and_then(|request| request.with_credential_headers(credentials))
|
||||
.map_err(|_| HttpSourceError::InvalidRequest)
|
||||
}
|
||||
|
||||
fn ensure_identity_encoding(headers: &HeaderMap) -> Result<(), HttpSourceError> {
|
||||
match headers.get(header::CONTENT_ENCODING) {
|
||||
None => Ok(()),
|
||||
Some(value) if value.as_bytes().eq_ignore_ascii_case(b"identity") => Ok(()),
|
||||
Some(_) => Err(HttpSourceError::InvalidResponse),
|
||||
}
|
||||
}
|
||||
|
||||
fn content_length(headers: &HeaderMap) -> Result<Option<u64>, HttpSourceError> {
|
||||
headers
|
||||
.get(header::CONTENT_LENGTH)
|
||||
.map(|value| {
|
||||
value
|
||||
.to_str()
|
||||
.map_err(|_| HttpSourceError::InvalidResponse)?
|
||||
.parse::<u64>()
|
||||
.map_err(|_| HttpSourceError::InvalidResponse)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn resumable_identity(
|
||||
requested_url: CanonicalUrl,
|
||||
final_url: CanonicalUrl,
|
||||
headers: &HeaderMap,
|
||||
length: Option<u64>,
|
||||
) -> Option<HttpSourceIdentity> {
|
||||
let etag = headers
|
||||
.get(header::ETAG)?
|
||||
.to_str()
|
||||
.ok()
|
||||
.and_then(|value| StrongEtag::parse(value).ok())?;
|
||||
Some(HttpSourceIdentity::from_parts(
|
||||
requested_url,
|
||||
final_url,
|
||||
etag,
|
||||
length?,
|
||||
))
|
||||
}
|
||||
|
||||
fn validate_resume_response(
|
||||
checkpoint: &HttpResumeCheckpoint,
|
||||
response: &remora_outbound::OutboundResponse,
|
||||
) -> Result<(), HttpSourceError> {
|
||||
ensure_identity_encoding(response.response().headers())
|
||||
.map_err(|_| HttpSourceError::SourceChanged)?;
|
||||
if response.final_url() != checkpoint.identity().final_url() {
|
||||
return Err(HttpSourceError::SourceChanged);
|
||||
}
|
||||
let etag = response
|
||||
.response()
|
||||
.headers()
|
||||
.get(header::ETAG)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| StrongEtag::parse(value).ok())
|
||||
.ok_or(HttpSourceError::SourceChanged)?;
|
||||
if &etag != checkpoint.identity().etag() {
|
||||
return Err(HttpSourceError::SourceChanged);
|
||||
}
|
||||
let total = checkpoint.identity().representation_length();
|
||||
let remaining = total - checkpoint.offset();
|
||||
if content_length(response.response().headers())? != Some(remaining) {
|
||||
return Err(HttpSourceError::SourceChanged);
|
||||
}
|
||||
let range = response
|
||||
.response()
|
||||
.headers()
|
||||
.get(header::CONTENT_RANGE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.ok_or(HttpSourceError::SourceChanged)?;
|
||||
let expected = format!("bytes {}-{}/{total}", checkpoint.offset(), total - 1);
|
||||
if range != expected {
|
||||
return Err(HttpSourceError::SourceChanged);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn status_error(status: StatusCode, headers: &HeaderMap, timing: &RetryTiming) -> HttpSourceError {
|
||||
match status {
|
||||
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => HttpSourceError::Authorization,
|
||||
StatusCode::NOT_FOUND | StatusCode::GONE => HttpSourceError::NotFound,
|
||||
status if status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error() => {
|
||||
HttpSourceError::Transient {
|
||||
retry_after: parse_retry_after(headers, timing, SystemTime::now()),
|
||||
}
|
||||
}
|
||||
_ => HttpSourceError::InvalidResponse,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_retry_after(
|
||||
headers: &HeaderMap,
|
||||
timing: &RetryTiming,
|
||||
now: SystemTime,
|
||||
) -> Option<Duration> {
|
||||
let value = headers.get(header::RETRY_AFTER)?.to_str().ok()?;
|
||||
let hint = value
|
||||
.parse::<u64>()
|
||||
.ok()
|
||||
.map(Duration::from_secs)
|
||||
.or_else(|| {
|
||||
httpdate::parse_http_date(value)
|
||||
.ok()?
|
||||
.duration_since(now)
|
||||
.ok()
|
||||
});
|
||||
capped_server_hint(hint, timing)
|
||||
}
|
||||
|
||||
fn map_outbound_error(error: OutboundError) -> HttpSourceError {
|
||||
match error {
|
||||
OutboundError::ResolutionFailed | OutboundError::Transport => {
|
||||
HttpSourceError::Transient { retry_after: None }
|
||||
}
|
||||
OutboundError::InvalidUrl | OutboundError::InvalidHeaders => {
|
||||
HttpSourceError::InvalidRequest
|
||||
}
|
||||
OutboundError::EmptyResolution
|
||||
| OutboundError::AddressDenied
|
||||
| OutboundError::SchemePortDenied
|
||||
| OutboundError::HttpsDowngrade => HttpSourceError::Authorization,
|
||||
OutboundError::InvalidRedirect | OutboundError::RedirectLimit => {
|
||||
HttpSourceError::InvalidResponse
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue