You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@opendal.apache.org by xu...@apache.org on 2024/05/21 06:19:02 UTC
(opendal) branch main updated: fix: Fetch the content length while end_bound is unknown (#4631)
This is an automated email from the ASF dual-hosted git repository.
xuanwo pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/opendal.git
The following commit(s) were added to refs/heads/main by this push:
new 227209f42c fix: Fetch the content length while end_bound is unknown (#4631)
227209f42c is described below
commit 227209f42c333dedd1fe51ab628e6d4f266e83d8
Author: Xuanwo <gi...@xuanwo.io>
AuthorDate: Tue May 21 14:18:52 2024 +0800
fix: Fetch the content length while end_bound is unknown (#4631)
* fix: Fetch the content length while end_bound is unknown
Signed-off-by: Xuanwo <gi...@xuanwo.io>
* Fix chunks
Signed-off-by: Xuanwo <gi...@xuanwo.io>
* Fix unit tests
Signed-off-by: Xuanwo <gi...@xuanwo.io>
---------
Signed-off-by: Xuanwo <gi...@xuanwo.io>
---
bin/oli/src/commands/cat.rs | 4 +-
bin/oli/src/commands/cp.rs | 8 ++-
bindings/nodejs/src/lib.rs | 6 +-
bindings/python/src/file.rs | 6 +-
bindings/python/src/operator.rs | 11 +++-
core/benches/ops/read.rs | 5 +-
core/src/layers/retry.rs | 7 +++
core/src/raw/atomic_util.rs | 56 ++++++++++++++++++
core/src/raw/mod.rs | 3 +
core/src/types/operator/operator.rs | 7 ++-
core/src/types/read/buffer_stream.rs | 52 +++++------------
core/src/types/read/futures_bytes_stream.rs | 4 +-
core/src/types/read/reader.rs | 91 ++++++++++++++++++++++++-----
core/tests/behavior/async_read.rs | 4 +-
integrations/object_store/src/lib.rs | 5 ++
15 files changed, 196 insertions(+), 73 deletions(-)
diff --git a/bin/oli/src/commands/cat.rs b/bin/oli/src/commands/cat.rs
index db3b2ef850..4bfe8190f7 100644
--- a/bin/oli/src/commands/cat.rs
+++ b/bin/oli/src/commands/cat.rs
@@ -39,7 +39,9 @@ pub async fn main(args: &ArgMatches) -> Result<()> {
let reader = op.reader(&path).await?;
let meta = op.stat(&path).await?;
- let mut buf_reader = reader.into_futures_async_read(0..meta.content_length());
+ let mut buf_reader = reader
+ .into_futures_async_read(0..meta.content_length())
+ .await?;
let mut stdout = io::AllowStdIo::new(std::io::stdout());
io::copy_buf(&mut buf_reader, &mut stdout).await?;
Ok(())
diff --git a/bin/oli/src/commands/cp.rs b/bin/oli/src/commands/cp.rs
index 7dd9c640f7..1ac27948a2 100644
--- a/bin/oli/src/commands/cp.rs
+++ b/bin/oli/src/commands/cp.rs
@@ -49,7 +49,9 @@ pub async fn main(args: &ArgMatches) -> Result<()> {
let mut dst_w = dst_op.writer(&dst_path).await?.into_futures_async_write();
let src_meta = src_op.stat(&src_path).await?;
let reader = src_op.reader_with(&src_path).chunk(8 * 1024 * 1024).await?;
- let buf_reader = reader.into_futures_async_read(0..src_meta.content_length());
+ let buf_reader = reader
+ .into_futures_async_read(0..src_meta.content_length())
+ .await?;
futures::io::copy_buf(buf_reader, &mut dst_w).await?;
// flush data
dst_w.close().await?;
@@ -71,7 +73,9 @@ pub async fn main(args: &ArgMatches) -> Result<()> {
.strip_prefix(prefix)
.expect("invalid path");
let reader = src_op.reader_with(de.path()).chunk(8 * 1024 * 1024).await?;
- let buf_reader = reader.into_futures_async_read(0..meta.content_length());
+ let buf_reader = reader
+ .into_futures_async_read(0..meta.content_length())
+ .await?;
let mut writer = dst_op
.writer(&dst_root.join(fp).to_string_lossy())
diff --git a/bindings/nodejs/src/lib.rs b/bindings/nodejs/src/lib.rs
index f8c44c7ebc..966d47100d 100644
--- a/bindings/nodejs/src/lib.rs
+++ b/bindings/nodejs/src/lib.rs
@@ -188,10 +188,12 @@ impl Operator {
/// It could be used to read large file in a streaming way.
#[napi]
pub async fn reader(&self, path: String) -> Result<Reader> {
- let meta = self.0.stat(&path).await.map_err(format_napi_error)?;
let r = self.0.reader(&path).await.map_err(format_napi_error)?;
Ok(Reader {
- inner: r.into_futures_async_read(0..meta.content_length()),
+ inner: r
+ .into_futures_async_read(..)
+ .await
+ .map_err(format_napi_error)?,
})
}
diff --git a/bindings/python/src/file.rs b/bindings/python/src/file.rs
index b71b903fcb..849a058a06 100644
--- a/bindings/python/src/file.rs
+++ b/bindings/python/src/file.rs
@@ -292,11 +292,9 @@ enum AsyncFileState {
}
impl AsyncFile {
- pub fn new_reader(reader: ocore::Reader, size: u64, capability: Capability) -> Self {
+ pub fn new_reader(reader: ocore::FuturesAsyncReader, capability: Capability) -> Self {
Self(
- Arc::new(Mutex::new(AsyncFileState::Reader(
- reader.into_futures_async_read(0..size),
- ))),
+ Arc::new(Mutex::new(AsyncFileState::Reader(reader))),
capability,
)
}
diff --git a/bindings/python/src/operator.rs b/bindings/python/src/operator.rs
index 8a3425ae7a..bf1cc90bc7 100644
--- a/bindings/python/src/operator.rs
+++ b/bindings/python/src/operator.rs
@@ -245,9 +245,14 @@ impl AsyncOperator {
future_into_py(py, async move {
if mode == "rb" {
- let meta = this.stat(&path).await.map_err(format_pyerr)?;
- let r = this.reader(&path).await.map_err(format_pyerr)?;
- Ok(AsyncFile::new_reader(r, meta.content_length(), capability))
+ let r = this
+ .reader(&path)
+ .await
+ .map_err(format_pyerr)?
+ .into_futures_async_read(..)
+ .await
+ .map_err(format_pyerr)?;
+ Ok(AsyncFile::new_reader(r, capability))
} else if mode == "wb" {
let w = this.writer(&path).await.map_err(format_pyerr)?;
Ok(AsyncFile::new_writer(w, capability))
diff --git a/core/benches/ops/read.rs b/core/benches/ops/read.rs
index b336c05efa..9520471835 100644
--- a/core/benches/ops/read.rs
+++ b/core/benches/ops/read.rs
@@ -52,7 +52,10 @@ fn bench_read_full(c: &mut Criterion, name: &str, op: Operator) {
b.to_async(&*TEST_RUNTIME).iter(|| async {
let r = op.reader_with(path).await.unwrap();
- let r = r.into_futures_async_read(0..size.bytes() as u64);
+ let r = r
+ .into_futures_async_read(0..size.bytes() as u64)
+ .await
+ .unwrap();
io::copy(r, &mut io::sink()).await.unwrap();
})
});
diff --git a/core/src/layers/retry.rs b/core/src/layers/retry.rs
index 0dca24193f..1756a6fbd2 100644
--- a/core/src/layers/retry.rs
+++ b/core/src/layers/retry.rs
@@ -940,6 +940,7 @@ mod tests {
let mut am = AccessorInfo::default();
am.set_native_capability(Capability {
read: true,
+ stat: true,
list: true,
list_with_recursive: true,
batch: true,
@@ -949,6 +950,12 @@ mod tests {
am
}
+ async fn stat(&self, _: &str, _: OpStat) -> Result<RpStat> {
+ Ok(RpStat::new(
+ Metadata::new(EntryMode::FILE).with_content_length(13),
+ ))
+ }
+
async fn read(&self, _: &str, _: OpRead) -> Result<(RpRead, Self::Reader)> {
Ok((
RpRead::new(),
diff --git a/core/src/raw/atomic_util.rs b/core/src/raw/atomic_util.rs
new file mode 100644
index 0000000000..4876dc7651
--- /dev/null
+++ b/core/src/raw/atomic_util.rs
@@ -0,0 +1,56 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::atomic::{AtomicU64, Ordering};
+
+/// AtomicContentLength is a wrapper of AtomicU64 that used to store content length.
+///
+/// It provides a way to store and load content length in atomic way, so caller don't need to
+/// use `Mutex` or `RwLock` to protect the content length.
+///
+/// We use value `u64::MAX` to represent unknown size, it's impossible for us to
+/// handle a file that has `u64::MAX` bytes.
+#[derive(Debug)]
+pub struct AtomicContentLength(AtomicU64);
+
+impl Default for AtomicContentLength {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl AtomicContentLength {
+ /// Create a new AtomicContentLength.
+ pub const fn new() -> Self {
+ Self(AtomicU64::new(u64::MAX))
+ }
+
+ /// Load content length from AtomicU64.
+ #[inline]
+ pub fn load(&self) -> Option<u64> {
+ match self.0.load(Ordering::Relaxed) {
+ u64::MAX => None,
+ v => Some(v),
+ }
+ }
+
+ /// Store content length to AtomicU64.
+ #[inline]
+ pub fn store(&self, v: u64) {
+ self.0.store(v, Ordering::Relaxed)
+ }
+}
diff --git a/core/src/raw/mod.rs b/core/src/raw/mod.rs
index 8065c86d27..fc2965e377 100644
--- a/core/src/raw/mod.rs
+++ b/core/src/raw/mod.rs
@@ -78,6 +78,9 @@ pub use futures_util::MaybeSend;
mod enum_utils;
pub use enum_utils::*;
+mod atomic_util;
+pub use atomic_util::*;
+
// Expose as a pub mod to avoid confusing.
pub mod adapters;
pub mod oio;
diff --git a/core/src/types/operator/operator.rs b/core/src/types/operator/operator.rs
index 2db831ba7c..ece483ab33 100644
--- a/core/src/types/operator/operator.rs
+++ b/core/src/types/operator/operator.rs
@@ -16,6 +16,7 @@
// under the License.
use std::future::Future;
+use std::sync::Arc;
use std::time::Duration;
use futures::stream;
@@ -547,8 +548,9 @@ impl Operator {
);
}
+ let path = Arc::new(path);
let range = options.range();
- let r = Reader::create(inner, &path, args, options).await?;
+ let r = Reader::create(inner, path, args, options).await?;
let buf = r.read(range.to_range()).await?;
Ok(buf)
},
@@ -656,7 +658,8 @@ impl Operator {
);
}
- Reader::create(inner.clone(), &path, args, options).await
+ let path = Arc::new(path);
+ Reader::create(inner.clone(), path, args, options).await
},
)
}
diff --git a/core/src/types/read/buffer_stream.rs b/core/src/types/read/buffer_stream.rs
index 62331f6eab..21297367aa 100644
--- a/core/src/types/read/buffer_stream.rs
+++ b/core/src/types/read/buffer_stream.rs
@@ -15,12 +15,8 @@
// specific language governing permissions and limitations
// under the License.
-use std::ops::Bound;
-use std::ops::RangeBounds;
+use std::ops::Range;
use std::pin::Pin;
-use std::sync::atomic::AtomicBool;
-use std::sync::atomic::Ordering;
-use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
@@ -41,30 +37,17 @@ struct FutureIterator {
chunk: Option<usize>,
offset: u64,
- end: Option<u64>,
- finished: Arc<AtomicBool>,
+ end: u64,
}
impl FutureIterator {
#[inline]
- fn new(r: oio::Reader, chunk: Option<usize>, range: impl RangeBounds<u64>) -> Self {
- let start = match range.start_bound().cloned() {
- Bound::Included(start) => start,
- Bound::Excluded(start) => start + 1,
- Bound::Unbounded => 0,
- };
- let end = match range.end_bound().cloned() {
- Bound::Included(end) => Some(end + 1),
- Bound::Excluded(end) => Some(end),
- Bound::Unbounded => None,
- };
-
+ fn new(r: oio::Reader, chunk: Option<usize>, range: Range<u64>) -> Self {
FutureIterator {
r,
chunk,
- offset: start,
- end,
- finished: Arc::default(),
+ offset: range.start,
+ end: range.end,
}
}
}
@@ -73,31 +56,21 @@ impl Iterator for FutureIterator {
type Item = BoxedFuture<'static, Result<Buffer>>;
fn next(&mut self) -> Option<Self::Item> {
- if self.offset >= self.end.unwrap_or(u64::MAX) {
- return None;
- }
- if self.finished.load(Ordering::Relaxed) {
+ if self.offset >= self.end {
return None;
}
let offset = self.offset;
- // TODO: replace with services preferred chunk size.
- let chunk = self.chunk.unwrap_or(4 * 1024 * 1024);
- let limit = self
- .end
- .map(|end| ((end - self.offset) as usize).min(chunk))
- .unwrap_or(chunk);
- let finished = self.finished.clone();
- let r = self.r.clone();
+ let mut limit = (self.end - self.offset) as usize;
+ if let Some(chunk) = self.chunk {
+ limit = limit.min(chunk)
+ }
// Update self.offset before building future.
self.offset += limit as u64;
+ let r = self.r.clone();
let fut = async move {
let buf = r.read_at_dyn(offset, limit).await?;
- if buf.len() < limit || limit == 0 {
- // Update finished marked if buf is less than limit.
- finished.store(true, Ordering::Relaxed);
- }
Ok(buf)
};
@@ -118,7 +91,7 @@ pub struct BufferStream(Buffered<Iter<FutureIterator>>);
impl BufferStream {
/// Create a new buffer stream from given reader.
#[inline]
- pub(crate) fn new(r: oio::Reader, options: OpReader, range: impl RangeBounds<u64>) -> Self {
+ pub(crate) fn new(r: oio::Reader, options: OpReader, range: Range<u64>) -> Self {
let iter = FutureIterator::new(r, options.chunk(), range);
let stream = stream::iter(iter).buffered(options.concurrent());
@@ -147,6 +120,7 @@ mod tests {
use bytes::Bytes;
use futures::TryStreamExt;
use pretty_assertions::assert_eq;
+ use std::sync::Arc;
use super::*;
diff --git a/core/src/types/read/futures_bytes_stream.rs b/core/src/types/read/futures_bytes_stream.rs
index 4301a3aa6e..7fc98007ba 100644
--- a/core/src/types/read/futures_bytes_stream.rs
+++ b/core/src/types/read/futures_bytes_stream.rs
@@ -16,7 +16,7 @@
// under the License.
use std::io;
-use std::ops::RangeBounds;
+use std::ops::Range;
use std::pin::Pin;
use std::task::ready;
use std::task::Context;
@@ -46,7 +46,7 @@ unsafe impl Sync for FuturesBytesStream {}
impl FuturesBytesStream {
/// NOTE: don't allow users to create FuturesStream directly.
#[inline]
- pub(crate) fn new(r: oio::Reader, options: OpReader, range: impl RangeBounds<u64>) -> Self {
+ pub(crate) fn new(r: oio::Reader, options: OpReader, range: Range<u64>) -> Self {
let stream = BufferStream::new(r, options, range);
FuturesBytesStream {
diff --git a/core/src/types/read/reader.rs b/core/src/types/read/reader.rs
index 27f0657ee6..e779cbb465 100644
--- a/core/src/types/read/reader.rs
+++ b/core/src/types/read/reader.rs
@@ -15,8 +15,9 @@
// specific language governing permissions and limitations
// under the License.
-use std::ops::Range;
use std::ops::RangeBounds;
+use std::ops::{Bound, Range};
+use std::sync::Arc;
use bytes::BufMut;
use futures::stream;
@@ -39,8 +40,14 @@ use crate::*;
/// [`Reader`] provides public API including [`Reader::read`], [`Reader:read_range`], and [`Reader::read_to_end`]. You can use those APIs directly without extra copy.
#[derive(Clone)]
pub struct Reader {
+ acc: Accessor,
+ path: Arc<String>,
+
inner: oio::Reader,
options: OpReader,
+
+ /// Total size of the reader.
+ size: Arc<AtomicContentLength>,
}
impl Reader {
@@ -53,13 +60,50 @@ impl Reader {
/// in crate only.
pub(crate) async fn create(
acc: Accessor,
- path: &str,
+ path: Arc<String>,
args: OpRead,
options: OpReader,
) -> Result<Self> {
- let (_, r) = acc.read(path, args).await?;
+ let (_, r) = acc.read(&path, args).await?;
+
+ Ok(Reader {
+ acc,
+ path,
+ inner: r,
+ options,
+ size: Arc::new(AtomicContentLength::new()),
+ })
+ }
+
+ /// Parse users input range bounds into valid `Range<u64>`.
+ ///
+ /// To avoid duplicated stat call, we will cache the size of the reader.
+ async fn parse_range(&self, range: impl RangeBounds<u64>) -> Result<Range<u64>> {
+ let start = match range.start_bound() {
+ Bound::Included(v) => *v,
+ Bound::Excluded(v) => v + 1,
+ Bound::Unbounded => 0,
+ };
+
+ let end = match range.end_bound() {
+ Bound::Included(v) => v + 1,
+ Bound::Excluded(v) => *v,
+ Bound::Unbounded => match self.size.load() {
+ Some(v) => v,
+ None => {
+ let size = self
+ .acc
+ .stat(&self.path, OpStat::new())
+ .await?
+ .into_metadata()
+ .content_length();
+ self.size.store(size);
+ size
+ }
+ },
+ };
- Ok(Reader { inner: r, options })
+ Ok(start..end)
}
/// Read give range from reader into [`Buffer`].
@@ -71,7 +115,7 @@ impl Reader {
///
/// - Buffer length smaller than range means we have reached the end of file.
pub async fn read(&self, range: impl RangeBounds<u64>) -> Result<Buffer> {
- let bufs: Vec<_> = self.clone().into_stream(range).try_collect().await?;
+ let bufs: Vec<_> = self.clone().into_stream(range).await?.try_collect().await?;
Ok(bufs.into_iter().flatten().collect())
}
@@ -88,7 +132,7 @@ impl Reader {
buf: &mut impl BufMut,
range: impl RangeBounds<u64>,
) -> Result<usize> {
- let mut stream = self.clone().into_stream(range);
+ let mut stream = self.clone().into_stream(range).await?;
let mut read = 0;
loop {
@@ -163,8 +207,9 @@ impl Reader {
/// This API can be public but we are not sure if it's useful for users.
/// And the name `BufferStream` is not good enough to expose to users.
/// Let's keep it inside for now.
- fn into_stream(self, range: impl RangeBounds<u64>) -> BufferStream {
- BufferStream::new(self.inner, self.options, range)
+ async fn into_stream(self, range: impl RangeBounds<u64>) -> Result<BufferStream> {
+ let range = self.parse_range(range).await?;
+ Ok(BufferStream::new(self.inner, self.options, range))
}
/// Convert reader into [`FuturesAsyncReader`] which implements [`futures::AsyncRead`],
@@ -190,7 +235,8 @@ impl Reader {
/// let mut r = op
/// .reader("hello.txt")
/// .await?
- /// .into_futures_async_read(1024..2048);
+ /// .into_futures_async_read(1024..2048)
+ /// .await?;
/// let mut bs = Vec::new();
/// r.read_to_end(&mut bs).await?;
///
@@ -215,7 +261,8 @@ impl Reader {
/// .concurrent(8)
/// .chunk(256)
/// .await?
- /// .into_futures_async_read(1024..2048);
+ /// .into_futures_async_read(1024..2048)
+ /// .await?;
/// let mut bs = Vec::new();
/// r.read_to_end(&mut bs).await?;
///
@@ -223,8 +270,12 @@ impl Reader {
/// }
/// ```
#[inline]
- pub fn into_futures_async_read(self, range: Range<u64>) -> FuturesAsyncReader {
- FuturesAsyncReader::new(self.inner, self.options, range)
+ pub async fn into_futures_async_read(
+ self,
+ range: impl RangeBounds<u64>,
+ ) -> Result<FuturesAsyncReader> {
+ let range = self.parse_range(range).await?;
+ Ok(FuturesAsyncReader::new(self.inner, self.options, range))
}
/// Convert reader into [`FuturesBytesStream`] which implements [`futures::Stream`].
@@ -242,7 +293,7 @@ impl Reader {
/// use opendal::Result;
///
/// async fn test(op: Operator) -> io::Result<()> {
- /// let mut s = op.reader("hello.txt").await?.into_bytes_stream(1024..2048);
+ /// let mut s = op.reader("hello.txt").await?.into_bytes_stream(1024..2048).await?;;
/// let bs: Vec<Bytes> = s.try_collect().await?;
///
/// Ok(())
@@ -267,20 +318,25 @@ impl Reader {
/// .concurrent(8)
/// .chunk(256)
/// .await?
- /// .into_bytes_stream(1024..2048);
+ /// .into_bytes_stream(1024..2048).await?;
/// let bs: Vec<Bytes> = s.try_collect().await?;
///
/// Ok(())
/// }
/// ```
#[inline]
- pub fn into_bytes_stream(self, range: impl RangeBounds<u64>) -> FuturesBytesStream {
- FuturesBytesStream::new(self.inner, self.options, range)
+ pub async fn into_bytes_stream(
+ self,
+ range: impl RangeBounds<u64>,
+ ) -> Result<FuturesBytesStream> {
+ let range = self.parse_range(range).await?;
+ Ok(FuturesBytesStream::new(self.inner, self.options, range))
}
}
#[cfg(test)]
mod tests {
+ use crate::layers::TypeEraseLayer;
use rand::rngs::ThreadRng;
use rand::Rng;
use rand::RngCore;
@@ -294,8 +350,11 @@ mod tests {
#[test]
fn test_trait() {
let v = Reader {
+ acc: Arc::new(TypeEraseLayer.layer(())),
+ path: Arc::new("hello".to_string()),
inner: Arc::new(Buffer::new()),
options: OpReader::new(),
+ size: Arc::new(AtomicContentLength::new()),
};
let _: Box<dyn Unpin + MaybeSend + Sync + 'static> = Box::new(v);
diff --git a/core/tests/behavior/async_read.rs b/core/tests/behavior/async_read.rs
index b6868ff35f..75b9e9961f 100644
--- a/core/tests/behavior/async_read.rs
+++ b/core/tests/behavior/async_read.rs
@@ -159,6 +159,7 @@ pub async fn test_reader(op: Operator) -> anyhow::Result<()> {
.reader(&path)
.await?
.into_bytes_stream(..)
+ .await?
.try_fold(Vec::new(), |mut acc, chunk| {
acc.extend_from_slice(&chunk);
async { Ok(acc) }
@@ -175,7 +176,8 @@ pub async fn test_reader(op: Operator) -> anyhow::Result<()> {
let mut futures_reader = op
.reader(&path)
.await?
- .into_futures_async_read(0..size as u64);
+ .into_futures_async_read(0..size as u64)
+ .await?;
let mut bs = Vec::new();
futures_reader.read_to_end(&mut bs).await?;
assert_eq!(size, bs.len(), "read size");
diff --git a/integrations/object_store/src/lib.rs b/integrations/object_store/src/lib.rs
index c00c4e38a0..6c18102e9e 100644
--- a/integrations/object_store/src/lib.rs
+++ b/integrations/object_store/src/lib.rs
@@ -139,6 +139,11 @@ impl ObjectStore for OpendalStore {
let stream = r
.into_bytes_stream(0..meta.size as u64)
+ .await
+ .map_err(|err| object_store::Error::Generic {
+ store: "IoError",
+ source: Box::new(err),
+ })?
.into_send()
.map_err(|err| object_store::Error::Generic {
store: "IoError",