use std::cmp::Ordering;
use std::fmt::Formatter;
use std::sync::Arc;
use crate::raw::oio;
use crate::raw::*;
use crate::*;
pub struct CompleteLayer;
impl std::fmt::Debug for CompleteLayer {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CompleteLayer").finish()
}
}
impl Layer for CompleteLayer {
fn apply_service(&self, inner: Servicer) -> Servicer {
Arc::new(self.layer(inner))
}
}
impl CompleteLayer {
fn layer(&self, inner: Servicer) -> CompleteService {
CompleteService { inner }
}
}
pub struct CompleteService {
inner: Servicer,
}
impl std::fmt::Debug for CompleteService {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
self.inner.fmt(f)
}
}
impl Service for CompleteService {
type Reader = CompleteReader<oio::Reader>;
type Writer = CompleteWriter<oio::Writer>;
type Lister = CompleteLister;
type Deleter = oio::Deleter;
type Copier = oio::Copier;
fn info(&self) -> ServiceInfo {
self.inner.info()
}
fn capability(&self) -> Capability {
self.inner.capability()
}
async fn create_dir(
&self,
ctx: &OperationContext,
path: &str,
args: OpCreateDir,
) -> Result<RpCreateDir> {
self.inner.create_dir(ctx, path, args).await
}
fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result<Self::Reader> {
let reader = self.inner.read(ctx, path, args)?;
Ok(CompleteReader::new(reader))
}
fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result<Self::Writer> {
let append = args.append();
let w = self.inner.write(ctx, path, args)?;
Ok(CompleteWriter::new(w, append))
}
fn copy(
&self,
ctx: &OperationContext,
from: &str,
to: &str,
args: OpCopy,
opts: OpCopier,
) -> Result<Self::Copier> {
self.inner.copy(ctx, from, to, args, opts)
}
async fn rename(
&self,
ctx: &OperationContext,
from: &str,
to: &str,
args: OpRename,
) -> Result<RpRename> {
self.inner.rename(ctx, from, to, args).await
}
async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result<RpStat> {
self.inner.stat(ctx, path, args).await
}
fn delete(&self, ctx: &OperationContext) -> Result<Self::Deleter> {
self.inner.delete(ctx)
}
fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result<Self::Lister> {
let lister = self.inner.list(ctx, path, args)?;
Ok(CompleteLister::new(ctx.clone(), self.inner.clone(), lister))
}
async fn presign(
&self,
ctx: &OperationContext,
path: &str,
args: OpPresign,
) -> Result<RpPresign> {
self.inner.presign(ctx, path, args).await
}
}
pub struct CompleteLister {
inner: oio::Lister,
ctx: OperationContext,
srv: Servicer,
}
impl CompleteLister {
fn new(ctx: OperationContext, srv: Servicer, inner: oio::Lister) -> Self {
Self { inner, ctx, srv }
}
async fn ensure_file_content_length(&self, entry: &mut oio::Entry) -> Result<()> {
let path = entry.path().to_string();
let version = entry.metadata().version().map(str::to_owned);
let mut op = OpStat::new();
if let Some(version) = version.as_deref() {
op = op.with_version(version);
}
let stat_metadata = self.srv.stat(&self.ctx, &path, op).await?.into_metadata();
if !stat_metadata.has_content_length() {
return Err(Error::new(
ErrorKind::Unexpected,
"content length is required for list file entries",
)
.with_operation("CompleteLister::ensure_file_content_length")
.with_context("service", self.srv.info().scheme().to_string())
.with_context("path", path));
}
entry
.metadata_mut()
.set_content_length(stat_metadata.content_length());
Ok(())
}
}
impl oio::List for CompleteLister {
async fn next(&mut self) -> Result<Option<oio::Entry>> {
loop {
let Some(mut entry) = self.inner.next().await? else {
return Ok(None);
};
if !entry.mode().is_file()
|| entry.metadata().is_deleted()
|| entry.metadata().has_content_length()
{
return Ok(Some(entry));
}
match self.ensure_file_content_length(&mut entry).await {
Ok(()) => return Ok(Some(entry)),
Err(err) if err.kind() == ErrorKind::NotFound => continue,
Err(err) => return Err(err),
}
}
}
}
pub struct CompleteReader<R> {
inner: R,
}
impl<R> CompleteReader<R> {
pub fn new(inner: R) -> Self {
Self { inner }
}
}
impl<R: oio::Read> oio::Read for CompleteReader<R> {
async fn open(&self, range: BytesRange) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
let size = if range.is_suffix() {
None
} else {
range.size()
};
self.inner.open(range).await.map(|(rp, stream)| {
(
rp,
Box::new(CompleteReadStream::new(stream, size)) as Box<dyn oio::ReadStreamDyn>,
)
})
}
async fn read(&self, range: BytesRange) -> Result<(RpRead, Buffer)> {
let size = if range.is_suffix() {
None
} else {
range.size()
};
let (rp, buffer) = self.inner.read(range).await?;
check_complete(size, buffer.len() as u64)?;
Ok((rp, buffer))
}
}
pub struct CompleteReadStream<R> {
inner: R,
size: Option<u64>,
read: u64,
}
impl<R> CompleteReadStream<R> {
pub fn new(inner: R, size: Option<u64>) -> Self {
Self {
inner,
size,
read: 0,
}
}
pub fn check(&self) -> Result<()> {
check_complete(self.size, self.read)
}
}
impl<R: oio::ReadStream> oio::ReadStream for CompleteReadStream<R> {
async fn read(&mut self) -> Result<Buffer> {
let buf = self.inner.read().await?;
if buf.is_empty() {
self.check()?;
} else {
self.read += buf.len() as u64;
}
Ok(buf)
}
}
fn check_complete(size: Option<u64>, actual: u64) -> Result<()> {
let Some(size) = size else {
return Ok(());
};
match actual.cmp(&size) {
Ordering::Equal => Ok(()),
Ordering::Less => Err(
Error::new(ErrorKind::Unexpected, "reader got too little data")
.with_context("expect", size)
.with_context("actual", actual),
),
Ordering::Greater => Err(
Error::new(ErrorKind::Unexpected, "reader got too much data")
.with_context("expect", size)
.with_context("actual", actual),
),
}
}
#[derive(Debug, PartialEq, Eq)]
enum CompleteState {
Open,
Written,
Closed,
Error,
}
impl CompleteState {
fn transition(&mut self, destination: CompleteState) {
if *self != CompleteState::Error {
*self = destination
}
}
}
pub struct CompleteWriter<W> {
inner: Option<W>,
append: bool,
size: u64,
state: CompleteState,
}
impl<W> CompleteWriter<W> {
pub fn new(inner: W, append: bool) -> CompleteWriter<W> {
CompleteWriter {
inner: Some(inner),
append,
size: 0,
state: CompleteState::Open,
}
}
fn check(&self, content_length: u64) -> Result<()> {
if self.append || content_length == 0 {
return Ok(());
}
match self.size.cmp(&content_length) {
Ordering::Equal => Ok(()),
Ordering::Less => Err(
Error::new(ErrorKind::Unexpected, "writer got too little data")
.with_context("expect", content_length)
.with_context("actual", self.size),
),
Ordering::Greater => Err(
Error::new(ErrorKind::Unexpected, "writer got too much data")
.with_context("expect", content_length)
.with_context("actual", self.size),
),
}
}
}
#[cfg(debug_assertions)]
impl<W> Drop for CompleteWriter<W> {
fn drop(&mut self) {
if self.state == CompleteState::Written {
log::warn!(
"writer has not been closed or aborted after successful write operation, must be a bug"
)
}
}
}
impl<W> oio::Write for CompleteWriter<W>
where
W: oio::Write,
{
async fn write(&mut self, bs: Buffer) -> Result<()> {
let w = self.inner.as_mut().ok_or_else(|| {
debug_assert_ne!(
self.state,
CompleteState::Open,
"bug: inner is empty, but state is Open"
);
Error::new(ErrorKind::Unexpected, "writer has been closed or aborted")
})?;
let len = bs.len();
w.write(bs)
.await
.inspect_err(|_| self.state.transition(CompleteState::Error))?;
self.size += len as u64;
self.state.transition(CompleteState::Written);
Ok(())
}
async fn close(&mut self) -> Result<Metadata> {
let w = self.inner.as_mut().ok_or_else(|| {
debug_assert_ne!(
self.state,
CompleteState::Open,
"bug: inner is empty, but state is Open"
);
Error::new(ErrorKind::Unexpected, "writer has been closed or aborted")
})?;
let mut ret = w
.close()
.await
.inspect_err(|_| self.state.transition(CompleteState::Error))?;
self.check(ret.content_length())
.inspect_err(|_| self.state.transition(CompleteState::Error))?;
if ret.content_length() == 0 {
ret = ret.with_content_length(self.size);
}
self.inner = None;
self.state.transition(CompleteState::Closed);
Ok(ret)
}
async fn abort(&mut self) -> Result<()> {
let w = self.inner.as_mut().ok_or_else(|| {
Error::new(ErrorKind::Unexpected, "writer has been closed or aborted")
})?;
w.abort()
.await
.inspect_err(|_| self.state.transition(CompleteState::Error))?;
self.inner = None;
self.state.transition(CompleteState::Closed);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
struct MockReadReader {
buffer: Buffer,
}
impl oio::Read for MockReadReader {
async fn open(&self, _: BytesRange) -> Result<(RpRead, Box<dyn oio::ReadStreamDyn>)> {
Err(Error::new(ErrorKind::Unsupported, "open is not supported"))
}
async fn read(&self, _: BytesRange) -> Result<(RpRead, Buffer)> {
Ok((
RpRead::new(
Metadata::new(EntryMode::FILE).with_content_length(self.buffer.len() as u64),
),
self.buffer.clone(),
))
}
}
fn new_test_reader(buffer: impl Into<Buffer>) -> CompleteReader<MockReadReader> {
CompleteReader::new(MockReadReader {
buffer: buffer.into(),
})
}
#[tokio::test]
async fn test_read_rejects_short_buffer() {
let reader = new_test_reader("a");
let err = oio::Read::read(&reader, BytesRange::from(0_u64..2))
.await
.expect_err("read should reject short buffer");
assert_eq!(err.kind(), ErrorKind::Unexpected);
}
#[tokio::test]
async fn test_read_rejects_extra_buffer() {
let reader = new_test_reader("ab");
let err = oio::Read::read(&reader, BytesRange::from(0_u64..1))
.await
.expect_err("read should reject extra buffer");
assert_eq!(err.kind(), ErrorKind::Unexpected);
}
}