84 lines
2.6 KiB
Rust
84 lines
2.6 KiB
Rust
use bytes::Bytes;
|
|
use futures_util::StreamExt;
|
|
use opendal::FuturesBytesStream;
|
|
use remora_stream::{SourceBody, StreamError, StreamSide};
|
|
|
|
pub struct ObjectSourceBody {
|
|
stream: FuturesBytesStream,
|
|
pending: Bytes,
|
|
pending_offset: usize,
|
|
remaining: u64,
|
|
finished: bool,
|
|
}
|
|
|
|
impl ObjectSourceBody {
|
|
pub(crate) fn new(stream: FuturesBytesStream, expected_length: u64) -> Self {
|
|
Self {
|
|
stream,
|
|
pending: Bytes::new(),
|
|
pending_offset: 0,
|
|
remaining: expected_length,
|
|
finished: false,
|
|
}
|
|
}
|
|
|
|
fn copy_pending(&mut self, output: &mut [u8]) -> usize {
|
|
let available = &self.pending[self.pending_offset..];
|
|
let count = available.len().min(output.len());
|
|
output[..count].copy_from_slice(&available[..count]);
|
|
self.pending_offset += count;
|
|
self.remaining -= count as u64;
|
|
if self.pending_offset == self.pending.len() {
|
|
self.pending = Bytes::new();
|
|
self.pending_offset = 0;
|
|
}
|
|
count
|
|
}
|
|
}
|
|
|
|
impl SourceBody for ObjectSourceBody {
|
|
async fn read(&mut self, output: &mut [u8]) -> Result<usize, StreamError> {
|
|
if output.is_empty() {
|
|
return Ok(0);
|
|
}
|
|
if self.pending_offset < self.pending.len() {
|
|
return Ok(self.copy_pending(output));
|
|
}
|
|
if self.finished {
|
|
return Ok(0);
|
|
}
|
|
match self.stream.next().await {
|
|
Some(Ok(chunk)) if chunk.len() as u64 <= self.remaining => {
|
|
self.pending = chunk;
|
|
Ok(self.copy_pending(output))
|
|
}
|
|
Some(Ok(_)) => Err(StreamError::SourceChanged),
|
|
Some(Err(error)) => Err(map_stream_error(&error)),
|
|
None if self.remaining != 0 => Err(StreamError::transient_source(
|
|
"object source ended before its declared length",
|
|
)),
|
|
None => {
|
|
self.finished = true;
|
|
Ok(0)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn map_stream_error(error: &std::io::Error) -> StreamError {
|
|
let kind = error
|
|
.get_ref()
|
|
.and_then(|source| source.downcast_ref::<opendal::Error>())
|
|
.map(opendal::Error::kind);
|
|
match kind {
|
|
Some(opendal::ErrorKind::PermissionDenied) => StreamError::Authorization {
|
|
side: StreamSide::Source,
|
|
},
|
|
Some(
|
|
opendal::ErrorKind::ConditionNotMatch
|
|
| opendal::ErrorKind::NotFound
|
|
| opendal::ErrorKind::RangeNotSatisfied,
|
|
) => StreamError::SourceChanged,
|
|
_ => StreamError::transient_source("object source read failed"),
|
|
}
|
|
}
|