--max-sendbuf-size flag to speed up large file upload/download (#1060)

* Add `--max-sendbuf-size` flag which now defaults to 64Kb

* Use `server_recvbuf_size` flag with base tunnel implementation

* isort

* Add to readme

Author: abhinavsingh

README.md

@@ -250,7 +250,7 @@
   - See `--enable-static-server` and `--static-server-dir` flags
 
 - Optimized for large file uploads and downloads
-  - See `--client-recvbuf-size` and `--server-recvbuf-size` flag
+  - See `--client-recvbuf-size`, `--server-recvbuf-size`, `--max-sendbuf-size` flags
 
 - `IPv4` and `IPv6` support
   - See `--hostname` flag
@@ -2280,7 +2280,8 @@ usage: -m [-h] [--tunnel-hostname TUNNEL_HOSTNAME] [--tunnel-port TUNNEL_PORT]
           [--work-klass WORK_KLASS] [--pid-file PID_FILE]
           [--enable-proxy-protocol] [--enable-conn-pool] [--key-file KEY_FILE]
           [--cert-file CERT_FILE] [--client-recvbuf-size CLIENT_RECVBUF_SIZE]
-          [--server-recvbuf-size SERVER_RECVBUF_SIZE] [--timeout TIMEOUT]
+          [--server-recvbuf-size SERVER_RECVBUF_SIZE]
+          [--max-sendbuf-size MAX_SENDBUF_SIZE] [--timeout TIMEOUT]
           [--disable-http-proxy] [--disable-headers DISABLE_HEADERS]
           [--ca-key-file CA_KEY_FILE] [--ca-cert-dir CA_CERT_DIR]
           [--ca-cert-file CA_CERT_FILE] [--ca-file CA_FILE]
@@ -2391,6 +2392,9 @@ options:
   --server-recvbuf-size SERVER_RECVBUF_SIZE
                         Default: 128 KB. Maximum amount of data received from
                         the server in a single recv() operation.
+  --max-sendbuf-size MAX_SENDBUF_SIZE
+                        Default: 64 KB. Maximum amount of data to dispatch in
+                        a single send() operation.
   --timeout TIMEOUT     Default: 10.0. Number of seconds after which an
                         inactive connection must be dropped. Inactivity is
                         defined by no data sent or received by the client.

proxy/common/constants.py

@@ -74,6 +74,7 @@ def _env_threadless_compliant() -> bool:
 # Defaults
 DEFAULT_BACKLOG = 100
 DEFAULT_BASIC_AUTH = None
+DEFAULT_MAX_SEND_SIZE = 64 * 1024
 DEFAULT_BUFFER_SIZE = 128 * 1024
 DEFAULT_CA_CERT_DIR = None
 DEFAULT_CA_CERT_FILE = None
@@ -131,7 +132,6 @@ def _env_threadless_compliant() -> bool:
 DEFAULT_VERSION = False
 DEFAULT_HTTP_PORT = 80
 DEFAULT_HTTPS_PORT = 443
-DEFAULT_MAX_SEND_SIZE = 16 * 1024
 DEFAULT_WORK_KLASS = 'proxy.http.HttpProtocolHandler'
 DEFAULT_ENABLE_PROXY_PROTOCOL = False
 # 25 milliseconds to keep the loops hot

proxy/core/base/tcp_server.py

@@ -27,7 +27,8 @@
 from ...core.connection import TcpClientConnection
 from ...common.constants import (
     DEFAULT_TIMEOUT, DEFAULT_KEY_FILE, DEFAULT_CERT_FILE,
-    DEFAULT_CLIENT_RECVBUF_SIZE, DEFAULT_SERVER_RECVBUF_SIZE,
+    DEFAULT_MAX_SEND_SIZE, DEFAULT_CLIENT_RECVBUF_SIZE,
+    DEFAULT_SERVER_RECVBUF_SIZE,
 )
 
 
@@ -68,6 +69,14 @@
     'server in a single recv() operation.',
 )
 
+flags.add_argument(
+    '--max-sendbuf-size',
+    type=int,
+    default=DEFAULT_MAX_SEND_SIZE,
+    help='Default: ' + str(int(DEFAULT_MAX_SEND_SIZE / 1024)) +
+    ' KB. Maximum amount of data to dispatch in a single send() operation.',
+)
+
 flags.add_argument(
     '--timeout',
     type=int,
@@ -164,7 +173,7 @@ async def handle_writables(self, writables: Writables) -> bool:
             logger.debug(
                 'Flushing buffer to client {0}'.format(self.work.address),
             )
-            self.work.flush()
+            self.work.flush(self.flags.max_sendbuf_size)
             if self.must_flush_before_shutdown is True and \
                     not self.work.has_buffer():
                 teardown = True

proxy/core/base/tcp_tunnel.py

@@ -89,15 +89,15 @@ async def handle_events(
             return do_shutdown
         # Handle server events
         if self.upstream and self.upstream.connection.fileno() in readables:
-            data = self.upstream.recv()
+            data = self.upstream.recv(self.flags.server_recvbuf_size)
             if data is None:
                 # Server closed connection
                 logger.debug('Connection closed by server')
                 return True
             # tunnel data to client
             self.work.queue(data)
         if self.upstream and self.upstream.connection.fileno() in writables:
-            self.upstream.flush()
+            self.upstream.flush(self.flags.max_sendbuf_size)
         return False
 
     def connect_upstream(self) -> None:

proxy/core/base/tcp_upstream.py

@@ -96,6 +96,7 @@ async def write_to_descriptors(self, w: Writables) -> bool:
                 self.upstream.connection.fileno() in w and \
                 self.upstream.has_buffer():
             try:
+                # TODO: max sendbuf size flag currently not used here
                 self.upstream.flush()
             except ssl.SSLWantWriteError:   # pragma: no cover
                 logger.info('Upstream SSLWantWriteError, will retry')

proxy/core/connection/connection.py

@@ -79,12 +79,15 @@ def queue(self, mv: memoryview) -> None:
         self.buffer.append(mv)
         self._num_buffer += 1
 
-    def flush(self) -> int:
+    def flush(self, max_send_size: Optional[int] = None) -> int:
         """Users must handle BrokenPipeError exceptions"""
         if not self.has_buffer():
             return 0
         mv = self.buffer[0].tobytes()
-        sent: int = self.send(mv[:DEFAULT_MAX_SEND_SIZE])
+        max_send_size = max_send_size or DEFAULT_MAX_SEND_SIZE
+        # TODO: Assemble multiple packets if total
+        # size remains below max send size.
+        sent: int = self.send(mv[:max_send_size])
         if sent == len(mv):
             self.buffer.pop(0)
             self._num_buffer -= 1

proxy/http/handler.py

@@ -396,7 +396,7 @@ def _flush(self) -> None:
                 ] = self.selector.select(timeout=DEFAULT_SELECTOR_SELECT_TIMEOUT)
                 if len(ev) == 0:
                     continue
-                self.work.flush()
+                self.work.flush(self.flags.max_sendbuf_size)
         except BrokenPipeError:
             pass
         finally:

proxy/http/proxy/server.py

@@ -188,7 +188,7 @@ async def write_to_descriptors(self, w: Writables) -> bool:
                 self.upstream.connection.fileno() in w:
             logger.debug('Server is write ready, flushing...')
             try:
-                self.upstream.flush()
+                self.upstream.flush(self.flags.max_sendbuf_size)
             except ssl.SSLWantWriteError:
                 logger.warning(
                     'SSLWantWriteError while trying to flush to server, will retry',

proxy/http/websocket/client.py

@@ -94,6 +94,7 @@ def run_once(self) -> bool:
         self.selector.unregister(self.sock)
         for _, mask in events:
             if mask & selectors.EVENT_READ and self.on_message:
+                # TODO: client recvbuf size flag currently not used here
                 raw = self.recv()
                 if raw is None or raw.tobytes() == b'':
                     self.closed = True
@@ -104,6 +105,7 @@ def run_once(self) -> bool:
                 frame.parse(raw.tobytes())
                 self.on_message(frame)
             elif mask & selectors.EVENT_WRITE:
+                # TODO: max sendbuf size flag currently not used here
                 self.flush()
         return False