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 2026/09/09 06:49:04 UTC
(opendal) 01/01: Optimize Python AsyncFile writes with owned buffers
This is an automated email from the ASF dual-hosted git repository.
Xuanwo pushed a commit to branch xuanwo/python-asyncfile-owned-buffer
in repository https://gitbox.apache.org/repos/asf/opendal.git
commit 174e04ea3a9e32dc74e165af58ea91f2afd61210
Author: Xuanwo <gi...@xuanwo.io>
AuthorDate: Wed Sep 9 14:48:52 2026 +0800
Optimize Python AsyncFile writes with owned buffers
---
bindings/python/src/file.rs | 70 ++++++++++++++++++++---
bindings/python/src/operator.rs | 11 +++-
bindings/python/tests/test_file_http_requests.py | 73 +++++++++++++++++++++++-
bindings/python/tests/test_write.py | 43 ++++++++++++++
4 files changed, 185 insertions(+), 12 deletions(-)
diff --git a/bindings/python/src/file.rs b/bindings/python/src/file.rs
index ed27b15f9e..922be9b8e0 100644
--- a/bindings/python/src/file.rs
+++ b/bindings/python/src/file.rs
@@ -25,7 +25,7 @@ use std::sync::Arc;
use asyncband::mutex::Mutex;
use futures::AsyncSeekExt;
-use futures::AsyncWriteExt;
+use futures::SinkExt;
use pyo3::IntoPyObjectExt;
use pyo3::buffer::PyBuffer;
use pyo3::exceptions::PyIOError;
@@ -437,18 +437,22 @@ impl File {
pub struct AsyncFile(Arc<Mutex<AsyncFileState>>);
enum AsyncFileState {
- Reader(ocore::FuturesAsyncReader),
- Writer(ocore::FuturesAsyncWriter),
+ Reader(Box<ocore::FuturesAsyncReader>),
+ Writer(ocore::BufferSink),
Closed,
}
impl AsyncFile {
pub fn new_reader(reader: ocore::FuturesAsyncReader) -> Self {
- Self(Arc::new(Mutex::new(AsyncFileState::Reader(reader))))
+ Self(Arc::new(Mutex::new(AsyncFileState::Reader(Box::new(
+ reader,
+ )))))
}
- pub fn new_writer(writer: ocore::FuturesAsyncWriter) -> Self {
- Self(Arc::new(Mutex::new(AsyncFileState::Writer(writer))))
+ pub fn new_writer(writer: ocore::Writer) -> Self {
+ Self(Arc::new(Mutex::new(AsyncFileState::Writer(
+ writer.into_sink(),
+ ))))
}
}
#[pymethods]
@@ -519,7 +523,7 @@ impl AsyncFile {
bs: &Bound<PyBytes>,
) -> PyResult<Bound<'p, PyAny>> {
let state = self.0.clone();
- let bs = PyBackedBytes::from(bs.clone());
+ let bs = py_bytes_like_into_buffer(bs.as_any())?;
future_into_py(py, async move {
let mut guard = state.lock().await;
@@ -538,8 +542,11 @@ impl AsyncFile {
};
let len = bs.len();
+ if len == 0 {
+ return Ok(0);
+ }
writer
- .write_all(&bs)
+ .send(bs)
.await
.map(|_| len)
.map_err(|err| PyIOError::new_err(err.to_string()))
@@ -656,7 +663,7 @@ impl AsyncFile {
future_into_py(py, async move {
let mut state = state.lock().await;
if let AsyncFileState::Writer(w) = &mut *state {
- w.close().await.map_err(format_pyerr_from_io_error)?;
+ w.close().await.map_err(format_pyerr)?;
}
*state = AsyncFileState::Closed;
Ok(())
@@ -749,3 +756,48 @@ impl AsyncFile {
})
}
}
+
+#[cfg(test)]
+mod tests {
+ use asyncband::semaphore::Semaphore;
+ use futures::poll;
+ use ocore::layers::ConcurrentLimitLayer;
+
+ use super::*;
+
+ #[tokio::test]
+ async fn test_cancelled_send_keeps_owned_buffer() {
+ Python::initialize();
+ let (buffer, source_ptr) = Python::attach(|py| {
+ let source = PyBytes::new(py, b"owned Python bytes");
+ let source_ptr = source.as_bytes().as_ptr();
+ let buffer = py_bytes_like_into_buffer(source.as_any()).unwrap();
+ (buffer, source_ptr)
+ });
+ let semaphore = Arc::new(Semaphore::new(0));
+ let op = ocore::Operator::via_iter("memory", [])
+ .unwrap()
+ .layer(ConcurrentLimitLayer::with_semaphore(semaphore.clone()));
+ let file = AsyncFile::new_writer(op.writer_with("cancelled").chunk(4).await.unwrap());
+ let mut state = file.0.lock().await;
+ let AsyncFileState::Writer(sink) = &mut *state else {
+ unreachable!();
+ };
+
+ // Poll into an underlying write blocked on the semaphore, then drop
+ // the caller's future. The sink must retain both the future and input.
+ {
+ let send = sink.send(buffer);
+ futures::pin_mut!(send);
+ assert!(poll!(send).is_pending());
+ }
+ semaphore.release(1);
+ sink.close().await.unwrap();
+ drop(state);
+ drop(file);
+
+ let result = op.read("cancelled").await.unwrap();
+ assert_eq!(result.to_vec(), b"owned Python bytes");
+ assert_eq!(result.current().as_ptr(), source_ptr);
+ }
+}
diff --git a/bindings/python/src/operator.rs b/bindings/python/src/operator.rs
index 5c1fad2045..0a9c85046a 100644
--- a/bindings/python/src/operator.rs
+++ b/bindings/python/src/operator.rs
@@ -1120,12 +1120,19 @@ impl AsyncOperator {
.map_err(format_pyerr)?;
Ok(AsyncFile::new_reader(r))
} else if mode == "wb" {
+ let mut writer_opts = writer_opts;
+ // Keep small writes coalesced when neither the caller nor the
+ // service supplies a chunk size, as the AsyncWrite adapter did.
+ if writer_opts.chunk.is_none()
+ && this.info().capability().write_multi_min_size.is_none()
+ {
+ writer_opts.chunk = Some(256 * 1024);
+ }
let writer = this
.writer_options(&path, writer_opts.into())
.await
.map_err(format_pyerr)?;
- let w = writer.into_futures_async_write();
- Ok(AsyncFile::new_writer(w))
+ Ok(AsyncFile::new_writer(writer))
} else {
Err(Unsupported::new_err(format!(
"OpenDAL doesn't support mode: {mode}"
diff --git a/bindings/python/tests/test_file_http_requests.py b/bindings/python/tests/test_file_http_requests.py
index a6ce17f674..a23f238d05 100644
--- a/bindings/python/tests/test_file_http_requests.py
+++ b/bindings/python/tests/test_file_http_requests.py
@@ -39,7 +39,12 @@ class RequestState:
self.slow_release = slow_release
def _methods(self) -> list[str]:
- return list(self._requests)
+ return [
+ entry[0] if isinstance(entry, tuple) else entry for entry in self._requests
+ ]
+
+ def _uploads(self) -> list[bytes]:
+ return [entry[1] for entry in self._requests if isinstance(entry, tuple)]
def _reset(self) -> None:
self._requests[:] = []
@@ -58,6 +63,18 @@ def serve_requests(port_queue, requests, slow_started, slow_release, stop):
self.send_header("Content-Length", str(len(CONTENT)))
self.end_headers()
+ def do_PUT(self) -> None:
+ body = self.rfile.read(int(self.headers["Content-Length"]))
+ requests.append(("PUT", body))
+ if self.path == "/slow":
+ slow_started.set()
+ slow_release.wait(timeout=5)
+ uploads = [entry for entry in requests if isinstance(entry, tuple)]
+ status = 403 if self.path == "/retry" and len(uploads) == 1 else 201
+ self.send_response(status)
+ self.send_header("Content-Length", "0")
+ self.end_headers()
+
def do_GET(self) -> None:
requests.append("GET")
content = TRUNCATED_CONTENT if self.path == "/truncated" else CONTENT
@@ -376,3 +393,57 @@ async def test_async_file_cancelled_read_can_close(request_server):
assert await file.closed
assert state._methods() == ["GET"]
state.slow_release.set()
+
+
+@pytest.mark.asyncio
+async def test_async_file_coalesces_small_writes(request_server):
+ endpoint, state = request_server
+ state._reset()
+ op = opendal.AsyncOperator(
+ "webdav", endpoint=endpoint, root="/", disable_create_dir="true"
+ )
+ async with await op.open("file", "wb") as file:
+ for _ in range(1000):
+ assert await file.write(b"small") == 5
+ assert state._methods() == []
+ assert state._uploads() == [b"small" * 1000]
+
+
+@pytest.mark.asyncio
+async def test_async_file_close_error_can_retry(request_server):
+ endpoint, state = request_server
+ state._reset()
+ op = opendal.AsyncOperator(
+ "webdav", endpoint=endpoint, root="/", disable_create_dir="true"
+ )
+ file = await op.open("retry", "wb")
+ assert await file.write(CONTENT) == len(CONTENT)
+ with pytest.raises(opendal.exceptions.PermissionDenied):
+ await file.close()
+ assert not await file.closed
+ await file.close()
+ assert await file.closed
+ assert state._uploads() == [CONTENT, CONTENT]
+
+
+@pytest.mark.asyncio
+async def test_async_file_cancelled_close_resumes_request(request_server):
+ endpoint, state = request_server
+ state._reset()
+ op = opendal.AsyncOperator(
+ "webdav", endpoint=endpoint, root="/", disable_create_dir="true"
+ )
+ file = await op.open("slow", "wb")
+ assert await file.write(CONTENT) == len(CONTENT)
+ pending = asyncio.ensure_future(file.close())
+ try:
+ assert await asyncio.to_thread(state.slow_started.wait, 2)
+ pending.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await pending
+ assert not await file.closed
+ finally:
+ state.slow_release.set()
+ await asyncio.wait_for(file.close(), timeout=2)
+ assert await file.closed
+ assert state._uploads() == [CONTENT]
diff --git a/bindings/python/tests/test_write.py b/bindings/python/tests/test_write.py
index cfe42b8b0e..4875e7c8a4 100644
--- a/bindings/python/tests/test_write.py
+++ b/bindings/python/tests/test_write.py
@@ -17,12 +17,14 @@
import gc
import os
+import sys
from pathlib import Path
from random import randint
from uuid import uuid4
import pytest
+import opendal
from opendal.exceptions import NotFound
@@ -170,6 +172,47 @@ async def test_async_writer_keeps_bytes_alive(service_name, operator, async_oper
await async_operator.delete(filename)
+@pytest.mark.asyncio
+@pytest.mark.need_capability("write", "write_can_multi", "read", "delete")
+@pytest.mark.parametrize("chunk", [None, 256 * 1024, 8 * 1024 * 1024])
+async def test_async_writer_mixed_sizes(async_operator, chunk):
+ filename = f"test_file_{uuid4()}"
+ sizes = [0, 1, 17, 256 * 1024, 256 * 1024 + 1, 8 * 1024 * 1024, 31, 0]
+ expected = bytearray()
+ options = {} if chunk is None else {"chunk": chunk}
+ async with await async_operator.open(filename, "wb", **options) as file:
+ for i, size in enumerate(sizes):
+ content = bytes([i]) * size
+ expected.extend(content)
+ assert await file.write(content) == size
+ del content
+ gc.collect()
+ assert await file.closed
+ await file.close()
+ with pytest.raises(OSError, match="closed file"):
+ await file.write(b"")
+ assert await async_operator.read(filename) == expected
+ await async_operator.delete(filename)
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not hasattr(sys, "getrefcount"), reason="requires reference counts")
+@pytest.mark.parametrize("size", [17, 256 * 1024, 8 * 1024 * 1024 + 1])
+async def test_async_writer_retains_python_owner(size):
+ op = opendal.AsyncOperator("memory")
+ content = os.urandom(size)
+ references = sys.getrefcount(content)
+ file = await op.open("owner", "wb", chunk=8 * 1024 * 1024)
+ assert await file.write(content) == size
+ # A copied staging buffer would let the Python owner go after write returns.
+ assert sys.getrefcount(content) > references
+ await file.close()
+ assert await op.read("owner") == content
+ await op.delete("owner")
+ gc.collect()
+ assert sys.getrefcount(content) == references
+
+
@pytest.mark.asyncio
@pytest.mark.need_capability("write", "delete", "write_with_if_not_exists")
async def test_async_writer_options(service_name, operator, async_operator):