Add exact resumable object source streaming

crates/remora-object-source/src/body.rs

@@ -0,0 +1,84 @@
+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"),
+    }
+}