Skip to content

Commit

Permalink
feat: allow to set custom buf chunk size for ServeDir
Browse files Browse the repository at this point in the history
  • Loading branch information
ttys3 committed Sep 4, 2021
1 parent 39625b0 commit 63d6179
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 9 deletions.
4 changes: 0 additions & 4 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,3 @@ members = [
"examples/warp-key-value-store",
"examples/tonic-key-value-store",
]

[patch.crates-io]
tokio = { git = "https://github.com/tokio-rs/tokio.git", branch = "master"}
tokio-util = { git = "https://github.com/tokio-rs/tokio.git", branch = "master" }
20 changes: 15 additions & 5 deletions tower-http/src/services/fs/serve_dir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use std::{
};
use tokio::fs::File;
use tower_service::Service;
use crate::services::fs::DEFAULT_CAPACITY;

/// Service that serves files from a given directory and all its sub directories.
///
Expand All @@ -28,6 +29,7 @@ use tower_service::Service;
pub struct ServeDir {
base: PathBuf,
append_index_html_on_directories: bool,
buf_chunk_size: usize,
}

impl ServeDir {
Expand All @@ -39,6 +41,7 @@ impl ServeDir {
Self {
base,
append_index_html_on_directories: true,
buf_chunk_size: DEFAULT_CAPACITY,
}
}

Expand All @@ -51,6 +54,12 @@ impl ServeDir {
self.append_index_html_on_directories = append;
self
}

/// set custom buffer chunk size.
pub fn with_buf_chunk_size(mut self, chunk_size: usize) -> Self {
self.buf_chunk_size = chunk_size;
self
}
}

impl<ReqBody> Service<Request<ReqBody>> for ServeDir {
Expand Down Expand Up @@ -87,6 +96,7 @@ impl<ReqBody> Service<Request<ReqBody>> for ServeDir {
}

let append_index_html_on_directories = self.append_index_html_on_directories;
let buf_chunk_size = self.buf_chunk_size;
let uri = req.uri().clone();

let open_file_future = Box::pin(async move {
Expand All @@ -113,7 +123,7 @@ impl<ReqBody> Service<Request<ReqBody>> for ServeDir {
});

let file = File::open(full_path).await?;
Ok(Output::File(file, mime))
Ok(Output::File(file, mime, buf_chunk_size))
});

ResponseFuture {
Expand Down Expand Up @@ -158,7 +168,7 @@ fn append_slash_on_path(uri: Uri) -> Uri {
}

enum Output {
File(File, HeaderValue),
File(File, HeaderValue, usize),
Redirect(HeaderValue),
NotFound,
}
Expand All @@ -181,8 +191,8 @@ impl Future for ResponseFuture {
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match &mut self.inner {
Inner::Valid(open_file_future) => {
let (file, mime) = match ready!(Pin::new(open_file_future).poll(cx)) {
Ok(Output::File(file, mime)) => (file, mime),
let (file, mime, chunk_size) = match ready!(Pin::new(open_file_future).poll(cx)) {
Ok(Output::File(file, mime, chunk_size)) => (file, mime, chunk_size),

Ok(Output::Redirect(location)) => {
let res = Response::builder()
Expand All @@ -208,7 +218,7 @@ impl Future for ResponseFuture {
)
}
};
let body = AsyncReadBody::new(file).boxed();
let body = AsyncReadBody::with_capacity(file, chunk_size).boxed();
let body = ResponseBody(body);

let mut res = Response::new(body);
Expand Down

0 comments on commit 63d6179

Please sign in to comment.