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/06/30 08:09:27 UTC
(opendal) 02/02: refactor: Move ChunkedWrite logic into WriteContext
This is an automated email from the ASF dual-hosted git repository.
xuanwo pushed a commit to branch move-chunked-write
in repository https://gitbox.apache.org/repos/asf/opendal.git
commit 63dfe3606beb2fd933cbaad7ad66a6f8c1d3f711
Author: Xuanwo <gi...@xuanwo.io>
AuthorDate: Sun Jun 30 16:08:58 2024 +0800
refactor: Move ChunkedWrite logic into WriteContext
Signed-off-by: Xuanwo <gi...@xuanwo.io>
---
core/Cargo.toml | 117 ++++----
core/benches/oio/main.rs | 25 --
core/benches/oio/utils.rs | 46 ----
core/benches/oio/write.rs | 60 ----
core/src/layers/complete.rs | 33 +--
core/src/lib.rs | 2 +-
core/src/raw/oio/buf/queue_buf.rs | 14 +-
core/src/raw/oio/write/mod.rs | 3 -
core/src/raw/ops.rs | 22 --
core/src/types/context/mod.rs | 3 +
.../chunked_write.rs => types/context/write.rs} | 302 ++++++++++++++-------
core/src/types/operator/operator.rs | 22 +-
core/src/types/operator/operator_functions.rs | 26 --
core/src/types/operator/operator_futures.rs | 32 +--
core/src/types/write/buffer_sink.rs | 39 ++-
core/src/types/write/futures_async_writer.rs | 23 +-
core/src/types/write/futures_bytes_sink.rs | 21 +-
core/src/types/write/writer.rs | 28 +-
18 files changed, 372 insertions(+), 446 deletions(-)
diff --git a/core/Cargo.toml b/core/Cargo.toml
index abdb7bae27..d3e2dc5df8 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -54,16 +54,16 @@ default = ["reqwest/rustls-tls", "executors-tokio", "services-memory"]
#
# You should never enable this feature unless you are developing opendal.
tests = [
- "dep:rand",
- "dep:sha2",
- "dep:dotenvy",
- "layers-blocking",
- "services-azblob",
- "services-fs",
- "services-http",
- "services-memory",
- "internal-tokio-rt",
- "services-s3",
+ "dep:rand",
+ "dep:sha2",
+ "dep:dotenvy",
+ "layers-blocking",
+ "services-azblob",
+ "services-fs",
+ "services-http",
+ "services-memory",
+ "internal-tokio-rt",
+ "services-s3",
]
# Enable path cache.
@@ -103,20 +103,20 @@ services-aliyun-drive = []
services-alluxio = []
services-atomicserver = ["dep:atomic_lib"]
services-azblob = [
- "dep:sha2",
- "dep:reqsign",
- "reqsign?/services-azblob",
- "reqsign?/reqwest_request",
+ "dep:sha2",
+ "dep:reqsign",
+ "reqsign?/services-azblob",
+ "reqsign?/reqwest_request",
]
services-azdls = [
- "dep:reqsign",
- "reqsign?/services-azblob",
- "reqsign?/reqwest_request",
+ "dep:reqsign",
+ "reqsign?/services-azblob",
+ "reqsign?/reqwest_request",
]
services-azfile = [
- "dep:reqsign",
- "reqsign?/services-azblob",
- "reqsign?/reqwest_request",
+ "dep:reqsign",
+ "reqsign?/services-azblob",
+ "reqsign?/reqwest_request",
]
services-b2 = []
services-cacache = ["dep:cacache"]
@@ -124,9 +124,9 @@ services-chainsafe = []
services-cloudflare-kv = []
services-compfs = ["dep:compio"]
services-cos = [
- "dep:reqsign",
- "reqsign?/services-tencent",
- "reqsign?/reqwest_request",
+ "dep:reqsign",
+ "reqsign?/services-tencent",
+ "reqsign?/reqwest_request",
]
services-d1 = []
services-dashmap = ["dep:dashmap"]
@@ -137,9 +137,9 @@ services-foundationdb = ["dep:foundationdb"]
services-fs = ["tokio/fs"]
services-ftp = ["dep:suppaftp", "dep:bb8", "dep:async-tls"]
services-gcs = [
- "dep:reqsign",
- "reqsign?/services-google",
- "reqsign?/reqwest_request",
+ "dep:reqsign",
+ "reqsign?/services-google",
+ "reqsign?/reqwest_request",
]
services-gdrive = ["internal-path-cache"]
services-ghac = []
@@ -161,15 +161,15 @@ services-moka = ["dep:moka"]
services-mongodb = ["dep:mongodb"]
services-mysql = ["dep:mysql_async"]
services-obs = [
- "dep:reqsign",
- "reqsign?/services-huaweicloud",
- "reqsign?/reqwest_request",
+ "dep:reqsign",
+ "reqsign?/services-huaweicloud",
+ "reqsign?/reqwest_request",
]
services-onedrive = []
services-oss = [
- "dep:reqsign",
- "reqsign?/services-aliyun",
- "reqsign?/reqwest_request",
+ "dep:reqsign",
+ "reqsign?/services-aliyun",
+ "reqsign?/reqwest_request",
]
services-pcloud = []
services-persy = ["dep:persy", "internal-tokio-rt"]
@@ -179,10 +179,10 @@ services-redis = ["dep:redis", "redis?/tokio-rustls-comp"]
services-redis-native-tls = ["services-redis", "redis?/tokio-native-tls-comp"]
services-rocksdb = ["dep:rocksdb", "internal-tokio-rt"]
services-s3 = [
- "dep:reqsign",
- "reqsign?/services-aws",
- "reqsign?/reqwest_request",
- "dep:crc32c",
+ "dep:reqsign",
+ "reqsign?/services-aws",
+ "reqsign?/reqwest_request",
+ "dep:crc32c",
]
services-seafile = []
services-sftp = ["dep:openssh", "dep:openssh-sftp-client", "dep:bb8"]
@@ -207,11 +207,6 @@ harness = false
name = "ops"
required-features = ["tests"]
-[[bench]]
-harness = false
-name = "oio"
-required-features = ["tests"]
-
[[bench]]
harness = false
name = "types"
@@ -233,13 +228,13 @@ backon = "0.4.4"
base64 = "0.22"
bytes = "1.6"
chrono = { version = "0.4.28", default-features = false, features = [
- "clock",
- "std",
+ "clock",
+ "std",
] }
flagset = "0.4"
futures = { version = "0.3", default-features = false, features = [
- "std",
- "async-await",
+ "std",
+ "async-await",
] }
http = "1.1"
log = "0.4"
@@ -249,7 +244,7 @@ once_cell = "1"
percent-encoding = "2"
quick-xml = { version = "0.31", features = ["serialize", "overlapped-lists"] }
reqwest = { version = "0.12.2", features = [
- "stream",
+ "stream",
], default-features = false }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
@@ -280,8 +275,8 @@ bb8-postgres = { version = "0.8.1", optional = true }
tokio-postgres = { version = "0.7.8", optional = true }
# for services-cacache
cacache = { version = "13.0", default-features = false, features = [
- "tokio-runtime",
- "mmap",
+ "tokio-runtime",
+ "mmap",
], optional = true }
# for services-dashmap
dashmap = { version = "5.4", optional = true }
@@ -289,7 +284,7 @@ dashmap = { version = "5.4", optional = true }
etcd-client = { version = "0.12", optional = true, features = ["tls"] }
# for services-foundationdb
foundationdb = { version = "0.8.0", features = [
- "embedded-fdb-include",
+ "embedded-fdb-include",
], optional = true }
# for services-hdfs
hdrs = { version = "0.3.2", optional = true, features = ["async_file"] }
@@ -305,13 +300,13 @@ moka = { version = "0.12", optional = true, features = ["future", "sync"] }
mongodb = { version = "2.8.1", optional = true, features = ["tokio-runtime"] }
# for services-mysql
mysql_async = { version = "0.32.2", default-features = false, features = [
- "default-rustls",
+ "default-rustls",
], optional = true }
# for services-sftp
openssh = { version = "0.10.0", optional = true }
openssh-sftp-client = { version = "0.14.0", optional = true, features = [
- "openssh",
- "tracing",
+ "openssh",
+ "tracing",
] }
# for services-persy
persy = { version = "1.4.6", optional = true }
@@ -319,9 +314,9 @@ persy = { version = "1.4.6", optional = true }
redb = { version = "1.1.0", optional = true }
# for services-redis
redis = { version = "0.23.1", features = [
- "cluster-async",
- "tokio-comp",
- "connection-manager",
+ "cluster-async",
+ "tokio-comp",
+ "connection-manager",
], optional = true }
# for services-rocksdb
rocksdb = { version = "0.21.0", default-features = false, optional = true }
@@ -331,9 +326,9 @@ rusqlite = { version = "0.31.0", optional = true, features = ["bundled"] }
sled = { version = "0.34.7", optional = true }
# for services-ftp
suppaftp = { version = "5.3.1", default-features = false, features = [
- "async-secure",
- "rustls",
- "async-rustls",
+ "async-secure",
+ "rustls",
+ "async-rustls",
], optional = true }
# for services-tikv
tikv-client = { version = "0.3.0", optional = true, default-features = false }
@@ -377,7 +372,7 @@ dotenvy = "0.15"
libtest-mimic = "0.6"
minitrace = { version = "0.6", features = ["enable"] }
opentelemetry = { version = "0.21", default-features = false, features = [
- "trace",
+ "trace",
] }
pretty_assertions = "1"
rand = "0.8"
@@ -385,6 +380,6 @@ sha2 = "0.10"
size = "0.4"
tokio = { version = "1.27", features = ["fs", "macros", "rt-multi-thread"] }
tracing-subscriber = { version = "0.3", features = [
- "env-filter",
- "tracing-log",
+ "env-filter",
+ "tracing-log",
] }
diff --git a/core/benches/oio/main.rs b/core/benches/oio/main.rs
deleted file mode 100644
index 41ae52f5d1..0000000000
--- a/core/benches/oio/main.rs
+++ /dev/null
@@ -1,25 +0,0 @@
-// 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.
-
-mod utils;
-mod write;
-
-use criterion::criterion_group;
-use criterion::criterion_main;
-
-criterion_group!(benches, write::bench_exact_buf_write);
-criterion_main!(benches);
diff --git a/core/benches/oio/utils.rs b/core/benches/oio/utils.rs
deleted file mode 100644
index fd46fef6c0..0000000000
--- a/core/benches/oio/utils.rs
+++ /dev/null
@@ -1,46 +0,0 @@
-// 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 bytes::Bytes;
-use opendal::raw::oio;
-use opendal::*;
-use rand::prelude::ThreadRng;
-use rand::RngCore;
-
-/// BlackHoleWriter will discard all data written to it so we can measure the buffer's cost.
-pub struct BlackHoleWriter;
-
-impl oio::Write for BlackHoleWriter {
- async fn write(&mut self, bs: Buffer) -> opendal::Result<usize> {
- Ok(bs.len())
- }
-
- async fn abort(&mut self) -> opendal::Result<()> {
- Ok(())
- }
-
- async fn close(&mut self) -> opendal::Result<()> {
- Ok(())
- }
-}
-
-pub fn gen_bytes(rng: &mut ThreadRng, size: usize) -> Bytes {
- let mut content = vec![0; size];
- rng.fill_bytes(&mut content);
-
- content.into()
-}
diff --git a/core/benches/oio/write.rs b/core/benches/oio/write.rs
deleted file mode 100644
index d9d8dd8b93..0000000000
--- a/core/benches/oio/write.rs
+++ /dev/null
@@ -1,60 +0,0 @@
-// 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 bytes::Buf;
-use criterion::Criterion;
-use once_cell::sync::Lazy;
-use opendal::raw::oio::ChunkedWriter;
-use opendal::raw::oio::Write;
-use rand::thread_rng;
-use size::Size;
-
-use super::utils::*;
-
-pub static TOKIO: Lazy<tokio::runtime::Runtime> =
- Lazy::new(|| tokio::runtime::Runtime::new().expect("build tokio runtime"));
-
-pub fn bench_exact_buf_write(c: &mut Criterion) {
- let mut group = c.benchmark_group("exact_buf_write");
-
- let mut rng = thread_rng();
-
- for size in [
- Size::from_kibibytes(4),
- Size::from_kibibytes(256),
- Size::from_mebibytes(4),
- Size::from_mebibytes(16),
- ] {
- let content = gen_bytes(&mut rng, size.bytes() as usize);
-
- group.throughput(criterion::Throughput::Bytes(size.bytes() as u64));
- group.bench_with_input(size.to_string(), &content, |b, content| {
- b.to_async(&*TOKIO).iter(|| async {
- let mut w = ChunkedWriter::new(BlackHoleWriter, 256 * 1024, true);
-
- let mut bs = content.clone();
- while !bs.is_empty() {
- let n = w.write(bs.clone().into()).await.unwrap();
- bs.advance(n);
- }
- w.close().await.unwrap();
- })
- });
- }
-
- group.finish()
-}
diff --git a/core/src/layers/complete.rs b/core/src/layers/complete.rs
index 224a63de03..c2156a0418 100644
--- a/core/src/layers/complete.rs
+++ b/core/src/layers/complete.rs
@@ -22,7 +22,6 @@ use std::sync::Arc;
use crate::raw::oio::FlatLister;
use crate::raw::oio::PrefixLister;
-use crate::raw::TwoWays;
use crate::raw::*;
use crate::*;
@@ -372,7 +371,7 @@ impl<A: Access> LayeredAccess for CompleteAccessor<A> {
type Inner = A;
type Reader = CompleteReader<A::Reader>;
type BlockingReader = CompleteReader<A::BlockingReader>;
- type Writer = TwoWays<CompleteWriter<A::Writer>, oio::ChunkedWriter<CompleteWriter<A::Writer>>>;
+ type Writer = CompleteWriter<A::Writer>;
type BlockingWriter = CompleteWriter<A::BlockingWriter>;
type Lister = CompleteLister<A, A::Lister>;
type BlockingLister = CompleteLister<A, A::BlockingLister>;
@@ -422,38 +421,8 @@ impl<A: Access> LayeredAccess for CompleteAccessor<A> {
));
}
- // Calculate buffer size.
- // If `chunk` is not set, we use `write_multi_min_size` or `write_multi_align_size`
- // as the default size.
- let chunk_size = args
- .chunk()
- .or(capability.write_multi_min_size)
- .or(capability.write_multi_align_size)
- .map(|mut size| {
- if let Some(v) = capability.write_multi_max_size {
- size = size.min(v);
- }
- if let Some(v) = capability.write_multi_min_size {
- size = size.max(v);
- }
- if let Some(v) = capability.write_multi_align_size {
- // Make sure size >= size first.
- size = size.max(v);
- size -= size % v;
- }
-
- size
- });
- let exact = args.chunk().is_some() || capability.write_multi_align_size.is_some();
-
let (rp, w) = self.inner.write(path, args.clone()).await?;
let w = CompleteWriter::new(w);
-
- let w = match chunk_size {
- None => TwoWays::One(w),
- Some(size) => TwoWays::Two(oio::ChunkedWriter::new(w, size, exact)),
- };
-
Ok((rp, w))
}
diff --git a/core/src/lib.rs b/core/src/lib.rs
index 2705f042ee..6018121051 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -176,6 +176,6 @@ mod tests {
/// This is used to make sure our public API implement Send + Sync
#[test]
fn test_trait() {
- let _: Box<dyn AssertSendSync> = Box::new(Writer::new(Box::new(())));
+ let _: Box<dyn AssertSendSync> = Box::new(Capability::default());
}
}
diff --git a/core/src/raw/oio/buf/queue_buf.rs b/core/src/raw/oio/buf/queue_buf.rs
index 168f2ee8f5..6e4fcee22c 100644
--- a/core/src/raw/oio/buf/queue_buf.rs
+++ b/core/src/raw/oio/buf/queue_buf.rs
@@ -15,9 +15,9 @@
// specific language governing permissions and limitations
// under the License.
-use std::collections::VecDeque;
-
use bytes::Buf;
+use std::collections::VecDeque;
+use std::mem;
use crate::*;
@@ -44,6 +44,10 @@ impl QueueBuf {
/// Push new [`Buffer`] into the queue.
#[inline]
pub fn push(&mut self, buf: Buffer) {
+ if buf.is_empty() {
+ return;
+ }
+
self.0.push_back(buf);
}
@@ -59,6 +63,12 @@ impl QueueBuf {
self.len() == 0
}
+ /// Take the entire buffer queue and leave `self` in empty states.
+ #[inline]
+ pub fn take(&mut self) -> QueueBuf {
+ mem::take(self)
+ }
+
/// Build a new [`Buffer`] from the queue.
///
/// If the queue is empty, it will return an empty buffer. Otherwise, it will iterate over all
diff --git a/core/src/raw/oio/write/mod.rs b/core/src/raw/oio/write/mod.rs
index 0649276f1c..48a1334407 100644
--- a/core/src/raw/oio/write/mod.rs
+++ b/core/src/raw/oio/write/mod.rs
@@ -35,9 +35,6 @@ mod one_shot_write;
pub use one_shot_write::OneShotWrite;
pub use one_shot_write::OneShotWriter;
-mod chunked_write;
-pub use chunked_write::ChunkedWriter;
-
mod range_write;
pub use range_write::RangeWrite;
pub use range_write::RangeWriter;
diff --git a/core/src/raw/ops.rs b/core/src/raw/ops.rs
index 8b9d3ea94f..aff6021ca5 100644
--- a/core/src/raw/ops.rs
+++ b/core/src/raw/ops.rs
@@ -575,7 +575,6 @@ impl OpStat {
#[derive(Debug, Clone, Default)]
pub struct OpWrite {
append: bool,
- chunk: Option<usize>,
concurrent: usize,
content_type: Option<String>,
content_disposition: Option<String>,
@@ -610,27 +609,6 @@ impl OpWrite {
self
}
- /// Get the chunk from op.
- ///
- /// The chunk is used by service to decide the chunk size of the underlying writer.
- pub fn chunk(&self) -> Option<usize> {
- self.chunk
- }
-
- /// Set the chunk of op.
- ///
- /// If chunk is set, the data will be chunked by the underlying writer.
- ///
- /// ## NOTE
- ///
- /// Service could have their own minimum chunk size while perform write
- /// operations like multipart uploads. So the chunk size may be larger than
- /// the given buffer size.
- pub fn with_chunk(mut self, chunk: usize) -> Self {
- self.chunk = Some(chunk);
- self
- }
-
/// Get the content type from option
pub fn content_type(&self) -> Option<&str> {
self.content_type.as_deref()
diff --git a/core/src/types/context/mod.rs b/core/src/types/context/mod.rs
index d8bcda96f5..9b2514b768 100644
--- a/core/src/types/context/mod.rs
+++ b/core/src/types/context/mod.rs
@@ -17,3 +17,6 @@
mod read;
pub use read::*;
+
+mod write;
+pub use write::*;
diff --git a/core/src/raw/oio/write/chunked_write.rs b/core/src/types/context/write.rs
similarity index 56%
rename from core/src/raw/oio/write/chunked_write.rs
rename to core/src/types/context/write.rs
index c98ad41c97..7cba9a0ed2 100644
--- a/core/src/raw/oio/write/chunked_write.rs
+++ b/core/src/types/context/write.rs
@@ -15,34 +15,116 @@
// specific language governing permissions and limitations
// under the License.
-use bytes::Buf;
-
+use crate::raw::oio::Write;
use crate::raw::*;
use crate::*;
+use bytes::Buf;
+use std::sync::Arc;
+
+/// WriteContext holds the immutable context for give write operation.
+pub struct WriteContext {
+ /// The accessor to the storage services.
+ acc: Accessor,
+ /// Path to the file.
+ path: String,
+ /// Arguments for the write operation.
+ args: OpWrite,
+ /// Options for the writer.
+ options: OpWriter,
+}
-/// ChunkedWriter is used to implement [`oio::Write`] based on chunk:
-/// flush the underlying storage at the `chunk`` size.
-///
-/// ChunkedWriter makes sure that the size of the data written to the
-/// underlying storage is
-/// - exactly `chunk` bytes if `exact` is true
-/// - at least `chunk` bytes if `exact` is false
-pub struct ChunkedWriter<W: oio::Write> {
- inner: W,
+impl WriteContext {
+ /// Create a new WriteContext.
+ #[inline]
+ pub fn new(acc: Accessor, path: String, args: OpWrite, options: OpWriter) -> Self {
+ Self {
+ acc,
+ path,
+ args,
+ options,
+ }
+ }
+
+ /// Get the accessor.
+ #[inline]
+ pub fn accessor(&self) -> &Accessor {
+ &self.acc
+ }
+
+ /// Get the path.
+ #[inline]
+ pub fn path(&self) -> &str {
+ &self.path
+ }
+
+ /// Get the arguments.
+ #[inline]
+ pub fn args(&self) -> &OpWrite {
+ &self.args
+ }
+
+ /// Get the options.
+ #[inline]
+ pub fn options(&self) -> &OpWriter {
+ &self.options
+ }
+
+ /// Calculate the chunk size for this write process.
+ ///
+ /// Returns the chunk size and if the chunk size is exact.
+ fn calculate_chunk_size(&self) -> (Option<usize>, bool) {
+ let cap = self.accessor().info().full_capability();
+
+ let exact = self.options().chunk().is_some();
+ let chunk_size = self
+ .options()
+ .chunk()
+ .or(cap.write_multi_min_size)
+ .map(|mut size| {
+ if let Some(v) = cap.write_multi_max_size {
+ size = size.min(v);
+ }
+ if let Some(v) = cap.write_multi_min_size {
+ size = size.max(v);
+ }
+
+ size
+ });
+
+ (chunk_size, exact)
+ }
+}
+
+pub struct WriteGenerator<W> {
+ w: W,
/// The size for buffer, we will flush the underlying storage at the size of this buffer.
- chunk_size: usize,
+ chunk_size: Option<usize>,
/// If `exact` is true, the size of the data written to the underlying storage is
/// exactly `chunk_size` bytes.
exact: bool,
buffer: oio::QueueBuf,
}
-impl<W: oio::Write> ChunkedWriter<W> {
+impl WriteGenerator<oio::Writer> {
/// Create a new exact buf writer.
- pub fn new(inner: W, chunk_size: usize, exact: bool) -> Self {
+ pub async fn create(ctx: Arc<WriteContext>) -> Result<Self> {
+ let (chunk_size, exact) = ctx.calculate_chunk_size();
+ let (_, w) = ctx.acc.write(ctx.path(), ctx.args().clone()).await?;
+
+ Ok(Self {
+ w,
+ chunk_size,
+ exact,
+ buffer: oio::QueueBuf::new(),
+ })
+ }
+
+ /// Allow building from existing oio::Writer for easier testing.
+ #[cfg(test)]
+ fn new(w: oio::Writer, chunk_size: Option<usize>, exact: bool) -> Self {
Self {
- inner,
+ w,
chunk_size,
exact,
buffer: oio::QueueBuf::new(),
@@ -50,75 +132,84 @@ impl<W: oio::Write> ChunkedWriter<W> {
}
}
-impl<W: oio::Write> oio::Write for ChunkedWriter<W> {
- async fn write(&mut self, mut bs: Buffer) -> Result<usize> {
- if self.exact {
- if self.buffer.len() >= self.chunk_size {
- let written = self.inner.write(self.buffer.clone().collect()).await?;
- self.buffer.advance(written);
- }
+impl WriteGenerator<oio::Writer> {
+ /// Write the entire buffer into writer.
+ pub async fn write(&mut self, mut bs: Buffer) -> Result<usize> {
+ let Some(chunk_size) = self.chunk_size else {
+ return self.w.write_dyn(bs).await;
+ };
- let remaining = self.chunk_size - self.buffer.len();
- bs.truncate(remaining);
- let n = bs.len();
+ if self.buffer.len() + bs.len() < chunk_size {
+ let size = bs.len();
self.buffer.push(bs);
- return Ok(n);
+ return Ok(size);
}
- // We are in inexact mode.
- if self.buffer.len() + bs.len() < self.chunk_size {
- // We haven't buffered enough data.
- let n = bs.len();
+ // Condition:
+ // - exact is false
+ // - buffer + bs is larger than chunk_size.
+ // Action:
+ // - write buffer + bs directly.
+ if !self.exact {
+ let fill_size = bs.len();
self.buffer.push(bs);
- return Ok(n);
+ let mut buf = self.buffer.take().collect();
+ let written = self.w.write_dyn(buf.clone()).await?;
+ buf.advance(written);
+ self.buffer.push(buf);
+ return Ok(fill_size);
}
- // We have enough data to send.
- if self.buffer.is_empty() {
- // Fast path: Sends the buffer directly if the buffer queue is empty.
- return self.inner.write(bs).await;
+ // Condition:
+ // - exact is true: we need write buffer in exact chunk size.
+ // - buffer is larger than chunk_size
+ // - in exact mode, the size must be chunk_size, use `>=` just for safe coding.
+ // Action:
+ // - write existing buffer in chunk_size to make more rooms for writing data.
+ if self.buffer.len() >= chunk_size {
+ let mut buf = self.buffer.take().collect();
+ let written = self.w.write_dyn(buf.clone()).await?;
+ buf.advance(written);
+ self.buffer.push(buf);
}
- // If we always push `bs` to the buffer queue, the buffer queue may grow infinitely if inner
- // doesn't fully consume the queue. So we clone the buffer queue and send it with `bs` first.
- let mut buffer = self.buffer.clone();
- buffer.push(bs.clone());
- let written = self.inner.write(buffer.collect()).await?;
- // The number of bytes in `self.buffer` that already written.
- let queue_written = written.min(self.buffer.len());
- self.buffer.advance(queue_written);
- // The number of bytes in `bs` that already written.
- let bs_written = written - queue_written;
- // Skip bytes that already written.
- bs.advance(bs_written);
- // We already sent `written` bytes so we put more `written` bytes into the buffer queue.
- bs.truncate(written);
- let n = bs_written + bs.len();
+ // Condition
+ // - exact is ture.
+ // - buffer size must lower than chunk_size.
+ // Action:
+ // - write bs to buffer with remaining size.
+ let remaining = chunk_size - self.buffer.len();
+ bs.truncate(remaining);
+ let n = bs.len();
self.buffer.push(bs);
Ok(n)
}
- async fn close(&mut self) -> Result<()> {
+ /// Finish the write process.
+ pub async fn close(&mut self) -> Result<()> {
loop {
if self.buffer.is_empty() {
break;
}
- let written = self.inner.write(self.buffer.clone().collect()).await?;
+ let written = self.w.write_dyn(self.buffer.clone().collect()).await?;
self.buffer.advance(written);
}
- self.inner.close().await
+ self.w.close().await
}
- async fn abort(&mut self) -> Result<()> {
+ /// Abort the write process.
+ pub async fn abort(&mut self) -> Result<()> {
self.buffer.clear();
- self.inner.abort().await
+ self.w.abort().await
}
}
#[cfg(test)]
mod tests {
+ use super::*;
+ use crate::raw::oio::Write;
use bytes::Buf;
use bytes::Bytes;
use log::debug;
@@ -128,12 +219,10 @@ mod tests {
use rand::RngCore;
use sha2::Digest;
use sha2::Sha256;
-
- use super::*;
- use crate::raw::oio::Write;
+ use tokio::sync::Mutex;
struct MockWriter {
- buf: Vec<u8>,
+ buf: Arc<Mutex<Vec<u8>>>,
}
impl Write for MockWriter {
@@ -141,7 +230,8 @@ mod tests {
debug!("test_fuzz_exact_buf_writer: flush size: {}", &bs.len());
let chunk = bs.chunk();
- self.buf.extend_from_slice(chunk);
+ let mut buf = self.buf.lock().await;
+ buf.extend_from_slice(chunk);
Ok(chunk.len())
}
@@ -166,7 +256,8 @@ mod tests {
let mut expected = vec![0; 5];
rng.fill_bytes(&mut expected);
- let mut w = ChunkedWriter::new(MockWriter { buf: vec![] }, 10, true);
+ let buf = Arc::new(Mutex::new(vec![]));
+ let mut w = WriteGenerator::new(Box::new(MockWriter { buf: buf.clone() }), Some(10), true);
let mut bs = Bytes::from(expected.clone());
while !bs.is_empty() {
@@ -176,9 +267,10 @@ mod tests {
w.close().await?;
- assert_eq!(w.inner.buf.len(), expected.len());
+ let buf = buf.lock().await;
+ assert_eq!(buf.len(), expected.len());
assert_eq!(
- format!("{:x}", Sha256::digest(&w.inner.buf)),
+ format!("{:x}", Sha256::digest(&*buf)),
format!("{:x}", Sha256::digest(&expected))
);
Ok(())
@@ -192,7 +284,8 @@ mod tests {
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.try_init();
- let mut w = ChunkedWriter::new(MockWriter { buf: vec![] }, 10, false);
+ let buf = Arc::new(Mutex::new(vec![]));
+ let mut w = WriteGenerator::new(Box::new(MockWriter { buf: buf.clone() }), Some(10), false);
let mut rng = thread_rng();
let mut expected = vec![0; 15];
@@ -205,9 +298,10 @@ mod tests {
w.close().await?;
- assert_eq!(w.inner.buf.len(), expected.len());
+ let buf = buf.lock().await;
+ assert_eq!(buf.len(), expected.len());
assert_eq!(
- format!("{:x}", Sha256::digest(&w.inner.buf)),
+ format!("{:x}", Sha256::digest(&*buf)),
format!("{:x}", Sha256::digest(&expected))
);
Ok(())
@@ -221,7 +315,8 @@ mod tests {
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.try_init();
- let mut w = ChunkedWriter::new(MockWriter { buf: vec![] }, 10, false);
+ let buf = Arc::new(Mutex::new(vec![]));
+ let mut w = WriteGenerator::new(Box::new(MockWriter { buf: buf.clone() }), Some(10), false);
let mut rng = thread_rng();
let mut expected = vec![];
@@ -240,19 +335,16 @@ mod tests {
let content = new_content(5);
assert_eq!(5, w.write(content.into()).await?);
// content > chunk size, but 5 bytes in queue.
- let mut content = new_content(15);
- // The MockWriter can only send 5 bytes each time, so we can only advance 5 bytes.
- assert_eq!(5, w.write(content.clone().into()).await?);
- content.advance(5);
- assert_eq!(5, w.write(content.clone().into()).await?);
- content.advance(5);
- assert_eq!(5, w.write(content.clone().into()).await?);
+ let content = new_content(15);
+ // The MockWriter can send all 15 bytes together, so we can only advance 5 bytes.
+ assert_eq!(15, w.write(content.clone().into()).await?);
w.close().await?;
- assert_eq!(w.inner.buf.len(), expected.len());
+ let buf = buf.lock().await;
+ assert_eq!(buf.len(), expected.len());
assert_eq!(
- format!("{:x}", Sha256::digest(&w.inner.buf)),
+ format!("{:x}", Sha256::digest(&*buf)),
format!("{:x}", Sha256::digest(&expected))
);
Ok(())
@@ -266,7 +358,8 @@ mod tests {
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.try_init();
- let mut w = ChunkedWriter::new(MockWriter { buf: vec![] }, 10, false);
+ let buf = Arc::new(Mutex::new(vec![]));
+ let mut w = WriteGenerator::new(Box::new(MockWriter { buf: buf.clone() }), Some(10), false);
let mut rng = thread_rng();
let mut expected = vec![];
@@ -287,51 +380,42 @@ mod tests {
// content < chunk size.
let content = new_content(3);
assert_eq!(3, w.write(content.into()).await?);
- // content > chunk size, but only sends the first chunk in the queue.
- let mut content = new_content(15);
- assert_eq!(5, w.write(content.clone().into()).await?);
- // queue: 3, 5, bs: 10
- content.advance(5);
- assert_eq!(3, w.write(content.clone().into()).await?);
- // queue: 5, 3, bs: 7
- content.advance(3);
- assert_eq!(5, w.write(content.clone().into()).await?);
- // queue: 3, 5, bs: 2
- content.advance(5);
- assert_eq!(2, w.write(content.clone().into()).await?);
- // queue: 5, 2, bs: empty.
- content.advance(2);
- assert!(content.is_empty());
+ // content > chunk size, but can send all chunks in the queue.
+ let content = new_content(15);
+ assert_eq!(15, w.write(content.clone().into()).await?);
w.close().await?;
- assert_eq!(w.inner.buf.len(), expected.len());
+ let buf = buf.lock().await;
+ assert_eq!(buf.len(), expected.len());
assert_eq!(
- format!("{:x}", Sha256::digest(&w.inner.buf)),
+ format!("{:x}", Sha256::digest(&*buf)),
format!("{:x}", Sha256::digest(&expected))
);
Ok(())
}
struct PartialWriter {
- buf: Vec<u8>,
+ buf: Arc<Mutex<Vec<u8>>>,
}
impl Write for PartialWriter {
async fn write(&mut self, mut bs: Buffer) -> Result<usize> {
+ let mut buf = self.buf.lock().await;
+
if Buffer::count(&bs) > 1 {
// Always leaves last buffer for non-contiguous buffer.
let mut written = 0;
while Buffer::count(&bs) > 1 {
let chunk = bs.chunk();
- self.buf.extend_from_slice(chunk);
+ buf.extend_from_slice(chunk);
written += chunk.len();
bs.advance(chunk.len());
}
Ok(written)
} else {
let chunk = bs.chunk();
- self.buf.extend_from_slice(chunk);
+ buf.extend_from_slice(chunk);
Ok(chunk.len())
}
}
@@ -353,7 +437,12 @@ mod tests {
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.try_init();
- let mut w = ChunkedWriter::new(PartialWriter { buf: vec![] }, 10, false);
+ let buf = Arc::new(Mutex::new(vec![]));
+ let mut w = WriteGenerator::new(
+ Box::new(PartialWriter { buf: buf.clone() }),
+ Some(10),
+ false,
+ );
let mut rng = thread_rng();
let mut expected = vec![];
@@ -374,9 +463,10 @@ mod tests {
w.close().await?;
- assert_eq!(w.inner.buf.len(), expected.len());
+ let buf = buf.lock().await;
+ assert_eq!(buf.len(), expected.len());
assert_eq!(
- format!("{:x}", Sha256::digest(&w.inner.buf)),
+ format!("{:x}", Sha256::digest(&*buf)),
format!("{:x}", Sha256::digest(&expected))
);
Ok(())
@@ -393,8 +483,13 @@ mod tests {
let mut rng = thread_rng();
let mut expected = vec![];
+ let buf = Arc::new(Mutex::new(vec![]));
let buffer_size = rng.gen_range(1..10);
- let mut writer = ChunkedWriter::new(MockWriter { buf: vec![] }, buffer_size, true);
+ let mut writer = WriteGenerator::new(
+ Box::new(MockWriter { buf: buf.clone() }),
+ Some(buffer_size),
+ true,
+ );
debug!("test_fuzz_exact_buf_writer: buffer size: {buffer_size}");
for _ in 0..1000 {
@@ -413,9 +508,10 @@ mod tests {
}
writer.close().await?;
- assert_eq!(writer.inner.buf.len(), expected.len());
+ let buf = buf.lock().await;
+ assert_eq!(buf.len(), expected.len());
assert_eq!(
- format!("{:x}", Sha256::digest(&writer.inner.buf)),
+ format!("{:x}", Sha256::digest(&*buf)),
format!("{:x}", Sha256::digest(&expected))
);
Ok(())
diff --git a/core/src/types/operator/operator.rs b/core/src/types/operator/operator.rs
index b06687292c..7cdaba59d0 100644
--- a/core/src/types/operator/operator.rs
+++ b/core/src/types/operator/operator.rs
@@ -1083,8 +1083,11 @@ impl Operator {
OperatorFuture::new(
self.inner().clone(),
path,
- OpWrite::default().merge_executor(self.default_executor.clone()),
- |inner, path, args| async move {
+ (
+ OpWrite::default().merge_executor(self.default_executor.clone()),
+ OpWriter::default(),
+ ),
+ |inner, path, (args, options)| async move {
if !validate_path(&path, EntryMode::FILE) {
return Err(
Error::new(ErrorKind::IsADirectory, "write path is a directory")
@@ -1094,7 +1097,9 @@ impl Operator {
);
}
- Writer::create(inner, &path, args).await
+ let context = WriteContext::new(inner, path, args, options);
+ let w = Writer::new(context).await?;
+ Ok(w)
},
)
}
@@ -1229,9 +1234,10 @@ impl Operator {
path,
(
OpWrite::default().merge_executor(self.default_executor.clone()),
+ OpWriter::default(),
bs,
),
- |inner, path, (args, bs)| async move {
+ |inner, path, (args, options, bs)| async move {
if !validate_path(&path, EntryMode::FILE) {
return Err(
Error::new(ErrorKind::IsADirectory, "write path is a directory")
@@ -1241,11 +1247,9 @@ impl Operator {
);
}
- let (_, w) = inner.write(&path, args).await?;
- let mut w = Writer::new(w);
- w.write(bs.clone()).await?;
- w.close().await?;
-
+ let context = WriteContext::new(inner, path, args, options);
+ let mut w = Writer::new(context).await?;
+ w.write(bs).await?;
Ok(())
},
)
diff --git a/core/src/types/operator/operator_functions.rs b/core/src/types/operator/operator_functions.rs
index 1788561cdf..d24f994fd7 100644
--- a/core/src/types/operator/operator_functions.rs
+++ b/core/src/types/operator/operator_functions.rs
@@ -87,19 +87,6 @@ impl FunctionWrite {
self
}
- /// Set the chunk size of op.
- ///
- /// If chunk size is set, the data will be chunked by the underlying writer.
- ///
- /// ## NOTE
- ///
- /// Service could have their own minimum chunk size while perform write operations like
- /// multipart uploads. So the chunk size may be larger than the given chunk size.
- pub fn chunk(mut self, v: usize) -> Self {
- self.0 = self.0.map_args(|(args, bs)| (args.with_chunk(v), bs));
- self
- }
-
/// Set the content type of option
pub fn content_type(mut self, v: &str) -> Self {
self.0 = self
@@ -153,19 +140,6 @@ impl FunctionWriter {
self
}
- /// Set the buffer size of op.
- ///
- /// If buffer size is set, the data will be buffered by the underlying writer.
- ///
- /// ## NOTE
- ///
- /// Service could have their own minimum buffer size while perform write operations like
- /// multipart uploads. So the buffer size may be larger than the given buffer size.
- pub fn buffer(mut self, v: usize) -> Self {
- self.0 = self.0.map_args(|args| args.with_chunk(v));
- self
- }
-
/// Set the content type of option
pub fn content_type(mut self, v: &str) -> Self {
self.0 = self.0.map_args(|args| args.with_content_type(v));
diff --git a/core/src/types/operator/operator_futures.rs b/core/src/types/operator/operator_futures.rs
index 2efe129041..045f211263 100644
--- a/core/src/types/operator/operator_futures.rs
+++ b/core/src/types/operator/operator_futures.rs
@@ -272,7 +272,7 @@ impl<F: Future<Output = Result<Reader>>> FutureReader<F> {
/// Future that generated by [`Operator::write_with`].
///
/// Users can add more options by public functions provided by this struct.
-pub type FutureWrite<F> = OperatorFuture<(OpWrite, Buffer), (), F>;
+pub type FutureWrite<F> = OperatorFuture<(OpWrite, OpWriter, Buffer), (), F>;
impl<F: Future<Output = Result<()>>> FutureWrite<F> {
/// Set the append mode of op.
@@ -283,7 +283,7 @@ impl<F: Future<Output = Result<()>>> FutureWrite<F> {
///
/// Service could return `Unsupported` if the underlying storage does not support append.
pub fn append(self, v: bool) -> Self {
- self.map(|(args, bs)| (args.with_append(v), bs))
+ self.map(|(args, options, bs)| (args.with_append(v), options, bs))
}
/// Set the buffer size of op.
@@ -295,39 +295,39 @@ impl<F: Future<Output = Result<()>>> FutureWrite<F> {
/// Service could have their own minimum buffer size while perform write operations like
/// multipart uploads. So the buffer size may be larger than the given buffer size.
pub fn chunk(self, v: usize) -> Self {
- self.map(|(args, bs)| (args.with_chunk(v), bs))
+ self.map(|(args, options, bs)| (args, options.with_chunk(v), bs))
}
/// Set the maximum concurrent write task amount.
pub fn concurrent(self, v: usize) -> Self {
- self.map(|(args, bs)| (args.with_concurrent(v), bs))
+ self.map(|(args, options, bs)| (args.with_concurrent(v), options, bs))
}
/// Set the content type of option
pub fn cache_control(self, v: &str) -> Self {
- self.map(|(args, bs)| (args.with_cache_control(v), bs))
+ self.map(|(args, options, bs)| (args.with_cache_control(v), options, bs))
}
/// Set the content type of option
pub fn content_type(self, v: &str) -> Self {
- self.map(|(args, bs)| (args.with_content_type(v), bs))
+ self.map(|(args, options, bs)| (args.with_content_type(v), options, bs))
}
/// Set the content disposition of option
pub fn content_disposition(self, v: &str) -> Self {
- self.map(|(args, bs)| (args.with_content_disposition(v), bs))
+ self.map(|(args, options, bs)| (args.with_content_disposition(v), options, bs))
}
/// Set the executor for this operation.
pub fn executor(self, executor: Executor) -> Self {
- self.map(|(args, bs)| (args.with_executor(executor), bs))
+ self.map(|(args, options, bs)| (args.with_executor(executor), options, bs))
}
}
/// Future that generated by [`Operator::writer_with`].
///
/// Users can add more options by public functions provided by this struct.
-pub type FutureWriter<F> = OperatorFuture<OpWrite, Writer, F>;
+pub type FutureWriter<F> = OperatorFuture<(OpWrite, OpWriter), Writer, F>;
impl<F: Future<Output = Result<Writer>>> FutureWriter<F> {
/// Set the append mode of op.
@@ -338,7 +338,7 @@ impl<F: Future<Output = Result<Writer>>> FutureWriter<F> {
///
/// Service could return `Unsupported` if the underlying storage does not support append.
pub fn append(self, v: bool) -> Self {
- self.map(|args| args.with_append(v))
+ self.map(|(args, options)| (args.with_append(v), options))
}
/// Set the chunk size of op.
@@ -357,32 +357,32 @@ impl<F: Future<Output = Result<Writer>>> FutureWriter<F> {
///
/// The services will alter the chunk size to meet their requirements.
pub fn chunk(self, v: usize) -> Self {
- self.map(|args| args.with_chunk(v))
+ self.map(|(args, options)| (args, options.with_chunk(v)))
}
/// Set the maximum concurrent write task amount.
pub fn concurrent(self, v: usize) -> Self {
- self.map(|args| args.with_concurrent(v))
+ self.map(|(args, options)| (args.with_concurrent(v), options))
}
/// Set the content type of option
pub fn cache_control(self, v: &str) -> Self {
- self.map(|args| args.with_cache_control(v))
+ self.map(|(args, options)| (args.with_cache_control(v), options))
}
/// Set the content type of option
pub fn content_type(self, v: &str) -> Self {
- self.map(|args| args.with_content_type(v))
+ self.map(|(args, options)| (args.with_content_type(v), options))
}
/// Set the content disposition of option
pub fn content_disposition(self, v: &str) -> Self {
- self.map(|args| args.with_content_disposition(v))
+ self.map(|(args, options)| (args.with_content_disposition(v), options))
}
/// Set the executor for this operation.
pub fn executor(self, executor: Executor) -> Self {
- self.map(|args| args.with_executor(executor))
+ self.map(|(args, options)| (args.with_executor(executor), options))
}
}
diff --git a/core/src/types/write/buffer_sink.rs b/core/src/types/write/buffer_sink.rs
index a230416767..46d9112530 100644
--- a/core/src/types/write/buffer_sink.rs
+++ b/core/src/types/write/buffer_sink.rs
@@ -34,9 +34,9 @@ pub struct BufferSink {
}
enum State {
- Idle(Option<oio::Writer>),
- Writing(BoxedStaticFuture<(oio::Writer, Result<usize>)>),
- Closing(BoxedStaticFuture<(oio::Writer, Result<()>)>),
+ Idle(Option<WriteGenerator<oio::Writer>>),
+ Writing(BoxedStaticFuture<(WriteGenerator<oio::Writer>, Result<usize>)>),
+ Closing(BoxedStaticFuture<(WriteGenerator<oio::Writer>, Result<()>)>),
}
/// # Safety
@@ -47,7 +47,7 @@ unsafe impl Sync for State {}
impl BufferSink {
/// Create a new sink from a [`oio::Writer`].
#[inline]
- pub(crate) fn new(w: oio::Writer) -> Self {
+ pub(crate) fn new(w: WriteGenerator<oio::Writer>) -> Self {
BufferSink {
state: State::Idle(Some(w)),
buf: Buffer::new(),
@@ -83,7 +83,7 @@ impl futures::Sink<Buffer> for BufferSink {
};
let buf = this.buf.clone();
let fut = async move {
- let res = w.write_dyn(buf).await;
+ let res = w.write(buf).await;
(w, res)
};
this.state = State::Writing(Box::pin(fut));
@@ -122,14 +122,14 @@ impl futures::Sink<Buffer> for BufferSink {
if this.buf.is_empty() {
let fut = async move {
- let res = w.close_dyn().await;
+ let res = w.close().await;
(w, res)
};
this.state = State::Closing(Box::pin(fut));
} else {
let buf = this.buf.clone();
let fut = async move {
- let res = w.write_dyn(buf).await;
+ let res = w.write(buf).await;
(w, res)
};
this.state = State::Writing(Box::pin(fut));
@@ -163,12 +163,25 @@ impl futures::Sink<Buffer> for BufferSink {
#[cfg(test)]
mod tests {
- use crate::raw::MaybeSend;
- use crate::BufferSink;
-
- #[test]
- fn test_trait() {
- let v = BufferSink::new(Box::new(()));
+ use crate::raw::*;
+ use crate::*;
+ use std::collections::HashMap;
+ use std::sync::Arc;
+
+ #[tokio::test]
+ async fn test_trait() {
+ let op = Operator::via_map(Scheme::Memory, HashMap::default()).unwrap();
+
+ let acc = op.into_inner();
+ let ctx = Arc::new(WriteContext::new(
+ acc,
+ "test".to_string(),
+ OpWrite::new(),
+ OpWriter::new().with_chunk(1),
+ ));
+ let write_gen = WriteGenerator::create(ctx).await.unwrap();
+
+ let v = BufferSink::new(write_gen);
let _: Box<dyn Unpin + MaybeSend + Sync + 'static> = Box::new(v);
}
diff --git a/core/src/types/write/futures_async_writer.rs b/core/src/types/write/futures_async_writer.rs
index 8e35db48a2..b1c932f762 100644
--- a/core/src/types/write/futures_async_writer.rs
+++ b/core/src/types/write/futures_async_writer.rs
@@ -40,7 +40,7 @@ pub struct FuturesAsyncWriter {
impl FuturesAsyncWriter {
/// NOTE: don't allow users to create directly.
#[inline]
- pub(crate) fn new(w: oio::Writer) -> Self {
+ pub(crate) fn new(w: WriteGenerator<oio::Writer>) -> Self {
FuturesAsyncWriter {
sink: BufferSink::new(w),
buf: oio::FlexBuf::new(256 * 1024),
@@ -113,10 +113,23 @@ impl AsyncWrite for FuturesAsyncWriter {
mod tests {
use super::*;
use crate::raw::MaybeSend;
-
- #[test]
- fn test_trait() {
- let v = FuturesAsyncWriter::new(Box::new(()));
+ use std::collections::HashMap;
+ use std::sync::Arc;
+
+ #[tokio::test]
+ async fn test_trait() {
+ let op = Operator::via_map(Scheme::Memory, HashMap::default()).unwrap();
+
+ let acc = op.into_inner();
+ let ctx = Arc::new(WriteContext::new(
+ acc,
+ "test".to_string(),
+ OpWrite::new(),
+ OpWriter::new().with_chunk(1),
+ ));
+ let write_gen = WriteGenerator::create(ctx).await.unwrap();
+
+ let v = FuturesAsyncWriter::new(write_gen);
let _: Box<dyn Unpin + MaybeSend + Sync + 'static> = Box::new(v);
}
diff --git a/core/src/types/write/futures_bytes_sink.rs b/core/src/types/write/futures_bytes_sink.rs
index 1643da212c..6a73013dec 100644
--- a/core/src/types/write/futures_bytes_sink.rs
+++ b/core/src/types/write/futures_bytes_sink.rs
@@ -38,7 +38,7 @@ pub struct FuturesBytesSink {
impl FuturesBytesSink {
/// Create a new sink from a [`oio::Writer`].
#[inline]
- pub(crate) fn new(w: oio::Writer) -> Self {
+ pub(crate) fn new(w: WriteGenerator<oio::Writer>) -> Self {
FuturesBytesSink {
sink: BufferSink::new(w),
}
@@ -80,10 +80,23 @@ impl futures::Sink<Bytes> for FuturesBytesSink {
mod tests {
use super::*;
use crate::raw::MaybeSend;
+ use std::collections::HashMap;
+ use std::sync::Arc;
- #[test]
- fn test_trait() {
- let v = FuturesBytesSink::new(Box::new(()));
+ #[tokio::test]
+ async fn test_trait() {
+ let op = Operator::via_map(Scheme::Memory, HashMap::default()).unwrap();
+
+ let acc = op.into_inner();
+ let ctx = Arc::new(WriteContext::new(
+ acc,
+ "test".to_string(),
+ OpWrite::new(),
+ OpWriter::new().with_chunk(1),
+ ));
+ let write_gen = WriteGenerator::create(ctx).await.unwrap();
+
+ let v = FuturesBytesSink::new(write_gen);
let _: Box<dyn Unpin + MaybeSend + Sync + 'static> = Box::new(v);
}
diff --git a/core/src/types/write/writer.rs b/core/src/types/write/writer.rs
index a378e56bf8..dc81d3d716 100644
--- a/core/src/types/write/writer.rs
+++ b/core/src/types/write/writer.rs
@@ -16,8 +16,8 @@
// under the License.
use bytes::Buf;
+use std::sync::Arc;
-use crate::raw::oio::Write;
use crate::raw::*;
use crate::*;
@@ -97,26 +97,18 @@ use crate::*;
/// - Services that doesn't support append will return [`ErrorKind::Unsupported`] error when
/// creating writer with `append` enabled.
pub struct Writer {
- inner: oio::Writer,
+ /// Keep a reference to write context in writer.
+ _ctx: Arc<WriteContext>,
+ inner: WriteGenerator<oio::Writer>,
}
impl Writer {
/// Create a new writer from an `oio::Writer`.
- pub(crate) fn new(w: oio::Writer) -> Self {
- Writer { inner: w }
- }
-
- /// Create a new writer.
- ///
- /// Create will use internal information to decide the most suitable
- /// implementation for users.
- ///
- /// We don't want to expose those details to users so keep this function
- /// in crate only.
- pub(crate) async fn create(acc: Accessor, path: &str, op: OpWrite) -> Result<Self> {
- let (_, w) = acc.write(path, op).await?;
+ pub(crate) async fn new(ctx: WriteContext) -> Result<Self> {
+ let ctx = Arc::new(ctx);
+ let inner = WriteGenerator::create(ctx.clone()).await?;
- Ok(Writer { inner: w })
+ Ok(Self { _ctx: ctx, inner })
}
/// Write [`Buffer`] into writer.
@@ -146,7 +138,7 @@ impl Writer {
pub async fn write(&mut self, bs: impl Into<Buffer>) -> Result<()> {
let mut bs = bs.into();
while !bs.is_empty() {
- let n = self.inner.write_dyn(bs.clone()).await?;
+ let n = self.inner.write(bs.clone()).await?;
bs.advance(n);
}
Ok(())
@@ -163,7 +155,7 @@ impl Writer {
let mut bs = bs;
let mut bs = Buffer::from(bs.copy_to_bytes(bs.remaining()));
while !bs.is_empty() {
- let n = self.inner.write_dyn(bs.clone()).await?;
+ let n = self.inner.write(bs.clone()).await?;
bs.advance(n);
}
Ok(())