From dc5af2e7e64cc40662c8d87f013b11367c600592 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 24 Jun 2024 10:36:51 -0400
Subject: [PATCH] more buffering

---
 src/unix/async_io.ml  | 21 ++++++++++++---------
 src/unix/async_io.mli |  3 +++
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/unix/async_io.ml b/src/unix/async_io.ml
index 5dd4766e..0d8c944e 100644
--- a/src/unix/async_io.ml
+++ b/src/unix/async_io.ml
@@ -161,8 +161,8 @@ module TCP_server = struct
       Ev_loop.wait_readable sock Cancel_handle.dummy ignore;
       accept_ sock
 
-  class base_server ?(listen = 32) ?(buf_pool = Buf_pool.dummy) ~runner
-    ~(handle : conn_handler) (addr : Sockaddr.t) :
+  class base_server ?(listen = 32) ?(buf_pool = Buf_pool.dummy)
+    ?(buf_size = 4096) ~runner ~(handle : conn_handler) (addr : Sockaddr.t) :
     t =
     let n_active_ = A.make 0 in
     let st = A.make Created in
@@ -214,8 +214,8 @@ module TCP_server = struct
                         Fd.close_noerr client_fd)
                   in
 
-                  let@ buf_in = buf_pool.with_buf 4096 in
-                  let@ buf_out = buf_pool.with_buf 4096 in
+                  let@ buf_in = buf_pool.with_buf buf_size in
+                  let@ buf_out = buf_pool.with_buf buf_size in
 
                   let ic = Buf_reader.of_fd client_fd ~buf:buf_in in
                   let oc = Buf_writer.of_fd client_fd ~buf:buf_out in
@@ -226,9 +226,11 @@ module TCP_server = struct
         )
     end
 
-  let create ?(after_init = ignore) ?listen ?buf_pool ~runner
+  let create ?(after_init = ignore) ?listen ?buf_pool ?buf_size ~runner
       ~(handle : conn_handler) (addr : Sockaddr.t) : t =
-    let self = new base_server ?listen ?buf_pool ~runner ~handle addr in
+    let self =
+      new base_server ?listen ?buf_pool ?buf_size ~runner ~handle addr
+    in
     after_init self;
     self
 end
@@ -256,10 +258,11 @@ module TCP_client = struct
     let@ () = Fun.protect ~finally in
     f sock
 
-  let with_connect ?(buf_pool = Buf_pool.dummy) addr (f : _ -> _ -> 'a) : 'a =
+  let with_connect ?(buf_pool = Buf_pool.dummy) ?(buf_size = 4096) addr
+      (f : _ -> _ -> 'a) : 'a =
     with_connect' addr (fun sock ->
-        let@ buf_in = buf_pool.with_buf 4096 in
-        let@ buf_out = buf_pool.with_buf 4096 in
+        let@ buf_in = buf_pool.with_buf buf_size in
+        let@ buf_out = buf_pool.with_buf buf_size in
 
         let ic = Buf_reader.of_fd ~buf:buf_in sock in
         let oc = Buf_writer.of_fd ~buf:buf_out sock in
diff --git a/src/unix/async_io.mli b/src/unix/async_io.mli
index 5860eee4..f044ee3f 100644
--- a/src/unix/async_io.mli
+++ b/src/unix/async_io.mli
@@ -53,6 +53,7 @@ module TCP_client : sig
 
   val with_connect :
     ?buf_pool:Buf_pool.t ->
+    ?buf_size:int ->
     Sockaddr.t ->
     (Buf_reader.t -> Buf_writer.t -> 'a) ->
     'a
@@ -88,6 +89,7 @@ module TCP_server : sig
   class base_server :
     ?listen:int ->
     ?buf_pool:Buf_pool.t ->
+    ?buf_size:int ->
     runner:Moonpool.Runner.t ->
     handle:conn_handler ->
     Sockaddr.t ->
@@ -97,6 +99,7 @@ module TCP_server : sig
     ?after_init:(t -> unit) ->
     ?listen:int ->
     ?buf_pool:Buf_pool.t ->
+    ?buf_size:int ->
     runner:Moonpool.Runner.t ->
     handle:conn_handler ->
     Sockaddr.t ->