Add --num-acceptors flag + Allow work_klass via Proxy context manager kwargs

examples/https_connect_tunnel.py

@@ -12,12 +12,11 @@
 
 from typing import Any, Optional
 
-from proxy.common.flag import FlagParser
+from proxy import Proxy
 from proxy.common.utils import build_http_response
 from proxy.http.codes import httpStatusCodes
 from proxy.http.parser import httpParserStates
 from proxy.http.methods import httpMethods
-from proxy.core.acceptor import AcceptorPool
 from proxy.core.base import BaseTcpTunnelHandler
 
 
@@ -75,11 +74,11 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
 
 def main() -> None:
     # This example requires `threadless=True`
-    with AcceptorPool(
-            flags=FlagParser.initialize(
-                port=12345, num_workers=1, threadless=True,
-            ),
-            work_klass=HttpsConnectTunnelHandler,
+    with Proxy(
+        work_klass=HttpsConnectTunnelHandler,
+        threadless=True,
+        num_workers=1,
+        port=12345,
     ):
         try:
             while True:

examples/ssl_echo_server.py

@@ -11,9 +11,8 @@
 import time
 from typing import Optional
 
-from proxy.common.flag import FlagParser
+from proxy import Proxy
 from proxy.common.utils import wrap_socket
-from proxy.core.acceptor import AcceptorPool
 from proxy.core.connection import TcpClientConnection
 
 from proxy.core.base import BaseTcpServerHandler
@@ -45,15 +44,13 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
 
 def main() -> None:
     # This example requires `threadless=True`
-    with AcceptorPool(
-        flags=FlagParser.initialize(
-            port=12345,
-            num_workers=1,
-            threadless=True,
-            keyfile='https-key.pem',
-            certfile='https-signed-cert.pem',
-        ),
+    with Proxy(
         work_klass=EchoSSLServerHandler,
+        threadless=True,
+        num_workers=1,
+        port=12345,
+        keyfile='https-key.pem',
+        certfile='https-signed-cert.pem',
     ):
         try:
             while True:

examples/tcp_echo_server.py

@@ -11,8 +11,7 @@
 import time
 from typing import Optional
 
-from proxy.common.flag import FlagParser
-from proxy.core.acceptor import AcceptorPool
+from proxy import Proxy
 from proxy.core.base import BaseTcpServerHandler
 
 
@@ -30,11 +29,11 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
 
 def main() -> None:
     # This example requires `threadless=True`
-    with AcceptorPool(
-        flags=FlagParser.initialize(
-            port=12345, num_workers=1, threadless=True,
-        ),
+    with Proxy(
         work_klass=EchoServerHandler,
+        threadless=True,
+        num_workers=1,
+        port=12345,
     ):
         try:
             while True:

examples/web_scraper.py

@@ -13,8 +13,8 @@
 
 from typing import Dict
 
-from proxy.common.flag import FlagParser
-from proxy.core.acceptor import Work, AcceptorPool
+from proxy import Proxy
+from proxy.core.acceptor import Work
 from proxy.common.types import Readables, Writables
 
 
@@ -56,15 +56,11 @@ def handle_events(
 
 
 if __name__ == '__main__':
-    with AcceptorPool(
-        flags=FlagParser.initialize(
-            port=12345,
-            num_workers=1,
-            threadless=True,
-            keyfile='https-key.pem',
-            certfile='https-signed-cert.pem',
-        ),
+    with Proxy(
         work_klass=WebScraper,
+        threadless=True,
+        num_workers=1,
+        port=12345,
     ) as pool:
         while True:
             time.sleep(1)

proxy/common/constants.py

@@ -87,6 +87,7 @@ def _env_threadless_compliant() -> bool:
 DEFAULT_HTTPS_ACCESS_LOG_FORMAT = '{client_ip}:{client_port} - ' + \
     '{request_method} {server_host}:{server_port} - ' + \
     '{response_bytes} bytes - {connection_time_ms} ms'
+DEFAULT_NUM_ACCEPTORS = 0
 DEFAULT_NUM_WORKERS = 0
 DEFAULT_OPEN_FILE_LIMIT = 1024
 DEFAULT_PAC_FILE = None
@@ -102,6 +103,7 @@ def _env_threadless_compliant() -> bool:
 DEFAULT_HTTP_PORT = 80
 DEFAULT_HTTPS_PORT = 443
 DEFAULT_MAX_SEND_SIZE = 16 * 1024
+DEFAULT_WORK_KLASS = 'proxy.http.HttpProtocolHandler'
 
 DEFAULT_DEVTOOLS_DOC_URL = 'http://proxy'
 DEFAULT_DEVTOOLS_FRAME_ID = secrets.token_hex(8)

proxy/common/flag.py

@@ -22,7 +22,7 @@
 from .plugins import Plugins
 from .types import IpAddress
 from .utils import text_, bytes_, is_py2, set_open_file_limit
-from .constants import COMMA, DEFAULT_DATA_DIRECTORY_PATH, DEFAULT_NUM_WORKERS
+from .constants import COMMA, DEFAULT_DATA_DIRECTORY_PATH, DEFAULT_NUM_ACCEPTORS, DEFAULT_NUM_WORKERS
 from .constants import DEFAULT_DEVTOOLS_WS_PATH, DEFAULT_DISABLE_HEADERS, PY2_DEPRECATION_MESSAGE
 from .constants import PLUGIN_DASHBOARD, PLUGIN_DEVTOOLS_PROTOCOL
 from .constants import PLUGIN_HTTP_PROXY, PLUGIN_INSPECT_TRAFFIC, PLUGIN_PAC_FILE
@@ -84,8 +84,8 @@ def parse_args(
 
     @staticmethod
     def initialize(
-        input_args: Optional[List[str]]
-        = None, **opts: Any,
+        input_args: Optional[List[str]] = None,
+        **opts: Any,
     ) -> argparse.Namespace:
         if input_args is None:
             input_args = []
@@ -107,13 +107,35 @@ def initialize(
             print(__version__)
             sys.exit(0)
 
+        # proxy.py currently cannot serve over HTTPS and also perform TLS interception
+        # at the same time.  Check if user is trying to enable both feature
+        # at the same time.
+        if (args.cert_file and args.key_file) and \
+                (args.ca_key_file and args.ca_cert_file and args.ca_signing_key_file):
+            print(
+                'You can either enable end-to-end encryption OR TLS interception,'
+                'not both together.',
+            )
+            sys.exit(1)
+
         # Setup logging module
         Logger.setup_logger(args.log_file, args.log_level, args.log_format)
 
         # Setup limits
         set_open_file_limit(args.open_file_limit)
 
-        # Load plugins
+        # Load work_klass
+        work_klass = opts.get('work_klass', args.work_klass)
+        work_klass = Plugins.importer(bytes_(work_klass))[0] \
+            if isinstance(work_klass, str) \
+            else work_klass
+
+        # Generate auth_code required for basic authentication if enabled
+        auth_code = None
+        if args.basic_auth:
+            auth_code = base64.b64encode(bytes_(args.basic_auth))
+
+        # Load default plugins along with user provided --plugins
         default_plugins = [
             bytes_(p)
             for p in FlagParser.get_default_plugins(args)
@@ -123,31 +145,14 @@ def initialize(
             for p in opts.get('plugins', args.plugins.split(text_(COMMA)))
             if not (isinstance(p, str) and len(p) == 0)
         ]
-
-        # Load default plugins along with user provided --plugins
         plugins = Plugins.load(default_plugins + extra_plugins)
 
-        # proxy.py currently cannot serve over HTTPS and also perform TLS interception
-        # at the same time.  Check if user is trying to enable both feature
-        # at the same time.
-        if (args.cert_file and args.key_file) and \
-                (args.ca_key_file and args.ca_cert_file and args.ca_signing_key_file):
-            print(
-                'You can either enable end-to-end encryption OR TLS interception,'
-                'not both together.',
-            )
-            sys.exit(1)
-
-        # Generate auth_code required for basic authentication if enabled
-        auth_code = None
-        if args.basic_auth:
-            auth_code = base64.b64encode(bytes_(args.basic_auth))
-
         # https://github.com/python/mypy/issues/5865
         #
         # def option(t: object, key: str, default: Any) -> Any:
         #     return cast(t, opts.get(key, default))
 
+        args.work_klass = work_klass
         args.plugins = plugins
         args.auth_code = cast(
             Optional[bytes],
@@ -235,18 +240,33 @@ def initialize(
                 socket.AF_INET6 if args.hostname.version == 6 else socket.AF_INET
             )
         else:
-            # FIXME: Not true for tests, as this value will be mock
+            # FIXME: Not true for tests, as this value will be a mock.
+            #
             # It's a problem only on Windows.  Instead of a proper
-            # test level fix, simply commenting this for now.
+            # fix in the tests, simply commenting this line of assertion
+            # for now.
+            #
             # assert args.unix_socket_path is None
             args.family = socket.AF_INET6 if args.hostname.version == 6 else socket.AF_INET
         args.port = cast(int, opts.get('port', args.port))
         args.backlog = cast(int, opts.get('backlog', args.backlog))
         num_workers = opts.get('num_workers', args.num_workers)
-        num_workers = num_workers if num_workers is not None else DEFAULT_NUM_WORKERS
         args.num_workers = cast(
             int, num_workers if num_workers > 0 else multiprocessing.cpu_count(),
         )
+        num_acceptors = opts.get('num_acceptors', args.num_acceptors)
+        # See https://github.com/abhinavsingh/proxy.py/pull/714 description
+        # to understand rationale behind the following logic.
+        #
+        # --num-workers flag or option was found. We will use
+        # the same value for num_acceptors when --num-acceptors flag
+        # is absent.
+        if num_workers != DEFAULT_NUM_WORKERS and num_acceptors == DEFAULT_NUM_ACCEPTORS:
+            args.num_acceptors = args.num_workers
+        else:
+            args.num_acceptors = cast(
+                int, num_acceptors if num_acceptors > 0 else multiprocessing.cpu_count(),
+            )
         args.static_server_dir = cast(
             str,
             opts.get(

proxy/core/acceptor/__init__.py

@@ -12,10 +12,12 @@
 from .pool import AcceptorPool
 from .work import Work
 from .threadless import Threadless
+from .executors import ThreadlessPool
 
 __all__ = [
     'Acceptor',
     'AcceptorPool',
     'Work',
     'Threadless',
+    'ThreadlessPool',
 ]