From 63d6179585a6677f3de7934e3d99292e206a0c2b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=8D=92=E9=87=8E=E7=84=A1=E7=87=88?= <ttys3@outlook.com>
Date: Sun, 5 Sep 2021 00:01:39 +0800
Subject: [PATCH] feat: allow to set custom buf chunk size for ServeDir

---
 Cargo.toml                              |  4 ----
 tower-http/src/services/fs/serve_dir.rs | 20 +++++++++++++++-----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 39e2c1ff..00a09395 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -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" }
\ No newline at end of file
diff --git a/tower-http/src/services/fs/serve_dir.rs b/tower-http/src/services/fs/serve_dir.rs
index a03166bf..d65ee36c 100644
--- a/tower-http/src/services/fs/serve_dir.rs
+++ b/tower-http/src/services/fs/serve_dir.rs
@@ -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.
 ///
@@ -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 {
@@ -39,6 +41,7 @@ impl ServeDir {
         Self {
             base,
             append_index_html_on_directories: true,
+            buf_chunk_size: DEFAULT_CAPACITY,
         }
     }
 
@@ -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 {
@@ -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 {
@@ -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 {
@@ -158,7 +168,7 @@ fn append_slash_on_path(uri: Uri) -> Uri {
 }
 
 enum Output {
-    File(File, HeaderValue),
+    File(File, HeaderValue, usize),
     Redirect(HeaderValue),
     NotFound,
 }
@@ -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()
@@ -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);