Refactor into separate Work module (#977)

* work module

* Fix imports

* String based typing for multiprocessing.synchronize

* Fix `test_accepts_client_from_server_socket`

* Move staticmethod outside of threadless pool class

* Fix doc build

* Fix test mock

* mp grouped together

* pylint happy

* import only for type checking

* doc build

* wrong import order

Author: abhinavsingh

docs/conf.py

@@ -296,8 +296,9 @@
     (_py_class_role, 'proxy.plugin.cache.store.base.CacheStore'),
     (_py_class_role, 'proxy.core.pool.AcceptorPool'),
     (_py_class_role, 'proxy.core.executors.ThreadlessPool'),
-    (_py_class_role, 'proxy.core.acceptor.threadless.T'),
-    (_py_class_role, 'proxy.core.acceptor.work.T'),
+    (_py_class_role, 'proxy.core.work.threadless.T'),
+    (_py_class_role, 'proxy.core.work.work.T'),
+    (_py_class_role, 'proxy.core.acceptor.threadless.Threadless'),
     (_py_class_role, 'queue.Queue[Any]'),
     (_py_class_role, 'SelectableEvents'),
     (_py_class_role, 'TcpClientConnection'),
@@ -309,6 +310,8 @@
     (_py_class_role, 'Url'),
     (_py_class_role, 'WebsocketFrame'),
     (_py_class_role, 'Work'),
-    (_py_obj_role, 'proxy.core.acceptor.threadless.T'),
-    (_py_obj_role, 'proxy.core.acceptor.work.T'),
+    (_py_class_role, 'proxy.core.acceptor.work.Work'),
+    (_py_class_role, 'connection.Connection'),
+    (_py_obj_role, 'proxy.core.work.threadless.T'),
+    (_py_obj_role, 'proxy.core.work.work.T'),
 ]

examples/web_scraper.py

@@ -12,7 +12,7 @@
 
 from proxy import Proxy
 from proxy.common.types import Readables, Writables, SelectableEvents
-from proxy.core.acceptor import Work
+from proxy.core.work import Work
 from proxy.core.connection import TcpClientConnection
 
 

proxy/core/acceptor/__init__.py

@@ -12,22 +12,12 @@
 
        pre
 """
+from .listener import Listener
 from .acceptor import Acceptor
 from .pool import AcceptorPool
-from .work import Work
-from .threadless import Threadless
-from .remote import RemoteExecutor
-from .local import LocalExecutor
-from .executors import ThreadlessPool
-from .listener import Listener
 
 __all__ = [
+    'Listener',
     'Acceptor',
     'AcceptorPool',
-    'Work',
-    'Threadless',
-    'RemoteExecutor',
-    'LocalExecutor',
-    'ThreadlessPool',
-    'Listener',
 ]

proxy/core/acceptor/acceptor.py

@@ -32,8 +32,7 @@
 
 from ..event import EventQueue
 
-from .local import LocalExecutor
-from .executors import ThreadlessPool
+from ..work import LocalExecutor, delegate_work_to_pool, start_threaded_work
 
 logger = logging.getLogger(__name__)
 
@@ -72,11 +71,11 @@ def __init__(
             idd: int,
             fd_queue: connection.Connection,
             flags: argparse.Namespace,
-            lock: multiprocessing.synchronize.Lock,
+            lock: 'multiprocessing.synchronize.Lock',
             # semaphore: multiprocessing.synchronize.Semaphore,
             executor_queues: List[connection.Connection],
             executor_pids: List[int],
-            executor_locks: List[multiprocessing.synchronize.Lock],
+            executor_locks: List['multiprocessing.synchronize.Lock'],
             event_queue: Optional[EventQueue] = None,
     ) -> None:
         super().__init__()
@@ -214,7 +213,7 @@ def _work(self, conn: socket.socket, addr: Optional[Tuple[str, int]]) -> None:
             # 1st workers.  To randomize, we offset index by idd.
             index = (self._total + self.idd) % self.flags.num_workers
             thread = threading.Thread(
-                target=ThreadlessPool.delegate,
+                target=delegate_work_to_pool,
                 args=(
                     self.executor_pids[index],
                     self.executor_queues[index],
@@ -231,7 +230,7 @@ def _work(self, conn: socket.socket, addr: Optional[Tuple[str, int]]) -> None:
                 ),
             )
         else:
-            _, thread = ThreadlessPool.start_threaded_work(
+            _, thread = start_threaded_work(
                 self.flags,
                 conn,
                 addr,

proxy/core/acceptor/pool.py

@@ -59,7 +59,7 @@ class AcceptorPool:
             while True:
                 time.sleep(1)
 
-    `flags.work_klass` must implement `work.Work` class.
+    `flags.work_klass` must implement :py:class:`~proxy.core.work.Work` class.
     """
 
     def __init__(
@@ -68,7 +68,7 @@ def __init__(
             listener: Listener,
             executor_queues: List[connection.Connection],
             executor_pids: List[int],
-            executor_locks: List[multiprocessing.synchronize.Lock],
+            executor_locks: List['multiprocessing.synchronize.Lock'],
             event_queue: Optional[EventQueue] = None,
     ) -> None:
         self.flags = flags
@@ -77,7 +77,7 @@ def __init__(
         # Available executors
         self.executor_queues: List[connection.Connection] = executor_queues
         self.executor_pids: List[int] = executor_pids
-        self.executor_locks: List[multiprocessing.synchronize.Lock] = executor_locks
+        self.executor_locks: List['multiprocessing.synchronize.Lock'] = executor_locks
         # Eventing core queue
         self.event_queue: Optional[EventQueue] = event_queue
         # Acceptor process instances

proxy/core/base/tcp_server.py

@@ -18,7 +18,7 @@
 from abc import abstractmethod
 from typing import Any, Optional
 
-from ...core.acceptor import Work
+from ...core.work import Work
 from ...core.connection import TcpClientConnection
 from ...common.types import Readables, SelectableEvents, Writables
 

proxy/core/connection/pool.py

@@ -21,7 +21,7 @@
 from ...common.flag import flags
 from ...common.types import Readables, SelectableEvents, Writables
 
-from ..acceptor.work import Work
+from ..work import Work
 
 from .server import TcpServerConnection
 

proxy/core/work/__init__.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+
+    .. spelling::
+
+       pre
+"""
+from .work import Work
+from .threadless import Threadless
+from .remote import RemoteExecutor
+from .local import LocalExecutor
+from .pool import ThreadlessPool
+from .delegate import delegate_work_to_pool
+from .threaded import start_threaded_work
+
+__all__ = [
+    'Work',
+    'Threadless',
+    'RemoteExecutor',
+    'LocalExecutor',
+    'ThreadlessPool',
+    'delegate_work_to_pool',
+    'start_threaded_work',
+]

proxy/core/work/delegate.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+from typing import TYPE_CHECKING, Optional, Tuple
+from multiprocessing.reduction import send_handle
+
+if TYPE_CHECKING:
+    import socket
+    import multiprocessing
+    from multiprocessing import connection
+
+
+def delegate_work_to_pool(
+        worker_pid: int,
+        work_queue: 'connection.Connection',
+        work_lock: 'multiprocessing.synchronize.Lock',
+        conn: 'socket.socket',
+        addr: Optional[Tuple[str, int]],
+        unix_socket_path: Optional[str] = None,
+) -> None:
+    """Utility method to delegate a work to threadless executor pool."""
+    with work_lock:
+        # Accepted client address is empty string for
+        # unix socket domain, avoid sending empty string
+        # for optimization.
+        if not unix_socket_path:
+            work_queue.send(addr)
+        send_handle(
+            work_queue,
+            conn.fileno(),
+            worker_pid,
+        )
+        conn.close()