Skip to content

Commit

Permalink
Add blocking layer
Browse files Browse the repository at this point in the history
Signed-off-by: yah01 <[email protected]>
  • Loading branch information
yah01 committed Aug 6, 2023
1 parent 338a8ae commit d18ef06
Show file tree
Hide file tree
Showing 9 changed files with 267 additions and 6 deletions.
224 changes: 224 additions & 0 deletions core/src/layers/blocking.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,224 @@
// 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 async_trait::async_trait;
use bytes;
use bytes::Bytes;
use tokio::runtime::Handle;

use crate::raw::oio::ReadExt;
use crate::raw::*;
use crate::*;

/// Add blocking API support for every operations.
///
/// # Blocking API
///
/// - This layer is auto-added to the operator if it's accessor doesn't support blocking APIs.
///
/// Tracking issue: #2678
#[derive(Debug, Clone)]
pub struct BlockingLayer {
handle: Handle,
}

impl BlockingLayer {
pub fn with_handle(handle: Handle) -> BlockingLayer {
Self { handle }
}
}

impl Default for BlockingLayer {
fn default() -> Self {
Self {
handle: Handle::current(),
}
}
}

impl<A: Accessor> Layer<A> for BlockingLayer {
type LayeredAccessor = BlockingAccessor<A>;

fn layer(&self, inner: A) -> Self::LayeredAccessor {
BlockingAccessor {
inner: inner,
handle: self.handle.clone(),
}
}
}

#[derive(Clone, Debug)]
pub struct BlockingAccessor<A: Accessor> {
inner: A,

handle: Handle,
}

#[async_trait]
impl<A: Accessor> LayeredAccessor for BlockingAccessor<A> {
type Inner = A;
type Reader = A::Reader;
type BlockingReader = BlockingWrapper<A::Reader>;
type Writer = A::Writer;
type BlockingWriter = BlockingWrapper<A::Writer>;
type Appender = A::Appender;
type Pager = A::Pager;
type BlockingPager = BlockingWrapper<A::Pager>;

fn inner(&self) -> &Self::Inner {
&self.inner
}

fn metadata(&self) -> AccessorInfo {
let info = self.inner.info();
let cap = info.capability_mut();
cap.blocking = true;
cap.read_can_seek = true;
info
}

async fn create_dir(&self, path: &str, args: OpCreateDir) -> Result<RpCreateDir> {
self.inner.create_dir(path, args).await
}

async fn read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::Reader)> {
self.inner.read(path, args).await
}

async fn write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::Writer)> {
self.inner.write(path, args).await
}

async fn append(&self, path: &str, args: OpAppend) -> Result<(RpAppend, Self::Appender)> {
self.inner.append(path, args).await
}

async fn copy(&self, from: &str, to: &str, args: OpCopy) -> Result<RpCopy> {
self.inner.copy(from, to, args).await
}

async fn rename(&self, from: &str, to: &str, args: OpRename) -> Result<RpRename> {
self.inner.rename(from, to, args).await
}

async fn stat(&self, path: &str, args: OpStat) -> Result<RpStat> {
self.inner.stat(path, args).await
}

async fn delete(&self, path: &str, args: OpDelete) -> Result<RpDelete> {
self.inner.delete(path, args).await
}

async fn list(&self, path: &str, args: OpList) -> Result<(RpList, Self::Pager)> {
self.inner.list(path, args).await
}

async fn presign(&self, path: &str, args: OpPresign) -> Result<RpPresign> {
self.inner.presign(path, args).await
}

async fn batch(&self, args: OpBatch) -> Result<RpBatch> {
self.inner.batch(args).await
}

fn blocking_create_dir(&self, path: &str, args: OpCreateDir) -> Result<RpCreateDir> {
self.handle.block_on(self.inner.create_dir(path, args))
}

fn blocking_read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::BlockingReader)> {
self.handle.block_on(async {
let (rp, reader) = self.inner.read(path, args).await?;
let blocking_reader = Self::BlockingReader::new(self.handle.clone(), reader);

Ok((rp, blocking_reader))
})
}

fn blocking_write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::BlockingWriter)> {
self.handle.block_on(async {
let (rp, writer) = self.inner.write(path, args).await?;
let blocking_writer = Self::BlockingWriter::new(self.handle.clone(), writer);
Ok((rp, blocking_writer))
})
}

fn blocking_copy(&self, from: &str, to: &str, args: OpCopy) -> Result<RpCopy> {
self.handle.block_on(self.inner.copy(from, to, args))
}

fn blocking_rename(&self, from: &str, to: &str, args: OpRename) -> Result<RpRename> {
self.handle.block_on(self.inner.rename(from, to, args))
}

fn blocking_stat(&self, path: &str, args: OpStat) -> Result<RpStat> {
self.handle.block_on(self.inner.stat(path, args))
}

fn blocking_delete(&self, path: &str, args: OpDelete) -> Result<RpDelete> {
self.handle.block_on(self.inner.delete(path, args))
}

fn blocking_list(&self, path: &str, args: OpList) -> Result<(RpList, Self::BlockingPager)> {
self.handle.block_on(async {
let (rp, pager) = self.inner.list(path, args).await?;
let blocking_pager = Self::BlockingPager::new(self.handle.clone(), pager);
Ok((rp, blocking_pager))
})
}
}

pub struct BlockingWrapper<I> {
handle: Handle,
inner: I,
}

impl<I> BlockingWrapper<I> {
fn new(handle: Handle, inner: I) -> Self {
Self { handle, inner }
}
}

impl<I: oio::Read + 'static> oio::BlockingRead for BlockingWrapper<I> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
self.handle
.block_on(oio::ReadExt::read(&mut self.inner, buf))
}

fn seek(&mut self, pos: std::io::SeekFrom) -> Result<u64> {
self.handle.block_on(self.inner.seek(pos))
}

fn next(&mut self) -> Option<Result<Bytes>> {
self.handle.block_on(self.inner.next())
}
}

impl<I: oio::Write + 'static> oio::BlockingWrite for BlockingWrapper<I> {
fn write(&mut self, bs: Bytes) -> Result<()> {
self.handle.block_on(oio::Write::write(&mut self.inner, bs))
}

fn close(&mut self) -> Result<()> {
self.handle.block_on(oio::Write::close(&mut self.inner))
}
}

impl<I: oio::Page> oio::BlockingPage for BlockingWrapper<I> {
fn next(&mut self) -> Result<Option<Vec<oio::Entry>>> {
self.handle.block_on(self.inner.next())
}
}
3 changes: 3 additions & 0 deletions core/src/layers/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ pub use logging::LoggingLayer;
mod timeout;
pub use timeout::TimeoutLayer;

mod blocking;
pub use blocking::BlockingLayer;

#[cfg(feature = "layers-chaos")]
mod chaos;
#[cfg(feature = "layers-chaos")]
Expand Down
1 change: 0 additions & 1 deletion core/src/types/operator/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -370,7 +370,6 @@ impl<A: Accessor> OperatorBuilder<A> {
/// Finish the building to construct an Operator.
pub fn finish(self) -> Operator {
let ob = self.layer(TypeEraseLayer);

Operator::from_inner(Arc::new(ob.accessor) as FusedAccessor)
}
}
9 changes: 8 additions & 1 deletion core/tests/behavior/blocking_copy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,17 @@ use crate::*;
pub fn behavior_blocking_copy_tests(op: &Operator) -> Vec<Trial> {
let cap = op.info().capability();

if !(cap.read && cap.write && cap.copy && cap.blocking) {
if !(cap.read && cap.write && cap.copy) {
return vec![];
}

let mut op = op.clone();
if !cap.blocking {
op = op.layer(layers::BlockingLayer::with_handle(RUNTIME.handle().clone()));
}

let op = &op;

blocking_trials!(
op,
test_blocking_copy_file,
Expand Down
9 changes: 8 additions & 1 deletion core/tests/behavior/blocking_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,17 @@ use crate::*;
pub fn behavior_blocking_list_tests(op: &Operator) -> Vec<Trial> {
let cap = op.info().capability();

if !(cap.read && cap.write && cap.copy && cap.blocking && cap.list) {
if !(cap.read && cap.write && cap.copy && cap.list) {
return vec![];
}

let mut op = op.clone();
if !cap.blocking {
op = op.layer(layers::BlockingLayer::with_handle(RUNTIME.handle().clone()));
}

let op = &op;

blocking_trials!(
op,
test_blocking_list_dir,
Expand Down
8 changes: 7 additions & 1 deletion core/tests/behavior/blocking_read_only.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,16 @@ use crate::*;
pub fn behavior_blocking_read_only_tests(op: &Operator) -> Vec<Trial> {
let cap = op.info().capability();

if !(cap.read && !cap.write && cap.blocking) {
if !(cap.read && !cap.write) {
return vec![];
}

let mut op = op.clone();
if !cap.blocking {
op = op.layer(layers::BlockingLayer::with_handle(RUNTIME.handle().clone()));
}

let op = &op;
blocking_trials!(
op,
test_blocking_read_only_stat_file_and_dir,
Expand Down
9 changes: 8 additions & 1 deletion core/tests/behavior/blocking_rename.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,17 @@ use crate::*;
pub fn behavior_blocking_rename_tests(op: &Operator) -> Vec<Trial> {
let cap = op.info().capability();

if !(cap.read && cap.write && cap.copy && cap.blocking && cap.rename) {
if !(cap.read && cap.write && cap.copy && cap.rename) {
return vec![];
}

let mut op = op.clone();
if !cap.blocking {
op = op.layer(layers::BlockingLayer::with_handle(RUNTIME.handle().clone()));
}

let op = &op;

blocking_trials!(
op,
test_blocking_rename_file,
Expand Down
9 changes: 8 additions & 1 deletion core/tests/behavior/blocking_write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,17 @@ use crate::*;
pub fn behavior_blocking_write_tests(op: &Operator) -> Vec<Trial> {
let cap = op.info().capability();

if !(cap.read && cap.write && cap.blocking) {
if !(cap.read && cap.write) {
return vec![];
}

let mut op = op.clone();
if !cap.blocking {
op = op.layer(layers::BlockingLayer::with_handle(RUNTIME.handle().clone()));
}

let op = &op;

blocking_trials!(
op,
test_blocking_create_dir,
Expand Down
1 change: 1 addition & 0 deletions core/tests/behavior/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ fn behavior_test<B: Builder>() -> Vec<Trial> {
trials.extend(behavior_blocking_read_only_tests(&operator));
trials.extend(behavior_blocking_rename_tests(&operator));
trials.extend(behavior_blocking_write_tests(&operator));

// Async tests
trials.extend(behavior_append_tests(&operator));
trials.extend(behavior_copy_tests(&operator));
Expand Down

0 comments on commit d18ef06

Please sign in to comment.