[Work] Refactor as FD work (#1049)

* Abstract out FD based work implementation

* No need of local/remote abstractions

* fix type

* Add `BaseLocalExecutor`

* Fix lint and tests

Author: abhinavsingh

docs/conf.py

@@ -322,5 +322,5 @@
     (_py_obj_role, 'proxy.core.work.threadless.T'),
     (_py_obj_role, 'proxy.core.work.work.T'),
     (_py_obj_role, 'proxy.core.base.tcp_server.T'),
-    (_py_obj_role, 'proxy.core.work.fd.T'),
+    (_py_obj_role, 'proxy.core.work.fd.fd.T'),
 ]

proxy/core/acceptor/acceptor.py

@@ -23,8 +23,9 @@
 from multiprocessing import connection
 from multiprocessing.reduction import recv_handle
 
-from ..work import LocalExecutor, start_threaded_work, delegate_work_to_pool
+from ..work import start_threaded_work, delegate_work_to_pool
 from ..event import EventQueue
+from ..work.fd import LocalFdExecutor
 from ...common.flag import flags
 from ...common.types import HostPort
 from ...common.logger import Logger
@@ -99,7 +100,7 @@ def __init__(
         # Internals
         self._total: Optional[int] = None
         self._local_work_queue: Optional['NonBlockingQueue'] = None
-        self._local: Optional[LocalExecutor] = None
+        self._local: Optional[LocalFdExecutor] = None
         self._lthread: Optional[threading.Thread] = None
 
     def accept(
@@ -195,7 +196,7 @@ def _recv_and_setup_socks(self) -> None:
     def _start_local(self) -> None:
         assert self.socks
         self._local_work_queue = NonBlockingQueue()
-        self._local = LocalExecutor(
+        self._local = LocalFdExecutor(
             iid=self.idd,
             work_queue=self._local_work_queue,
             flags=self.flags,

proxy/core/work/__init__.py

@@ -14,8 +14,7 @@
 """
 from .pool import ThreadlessPool
 from .work import Work
-from .local import LocalExecutor
-from .remote import RemoteExecutor
+from .local import BaseLocalExecutor
 from .delegate import delegate_work_to_pool
 from .threaded import start_threaded_work
 from .threadless import Threadless
@@ -24,9 +23,8 @@
 __all__ = [
     'Work',
     'Threadless',
-    'RemoteExecutor',
-    'LocalExecutor',
     'ThreadlessPool',
     'delegate_work_to_pool',
     'start_threaded_work',
+    'BaseLocalExecutor',
 ]

proxy/core/work/fd/__init__.py

@@ -0,0 +1,20 @@
+# -*- 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 .fd import ThreadlessFdExecutor
+from .local import LocalFdExecutor
+from .remote import RemoteFdExecutor
+
+
+__all__ = [
+    'ThreadlessFdExecutor',
+    'LocalFdExecutor',
+    'RemoteFdExecutor',
+]

proxy/core/work/fd.py proxy/core/work/fd/fd.pyproxy/core/work/fd.py renamed to proxy/core/work/fd/fd.py

@@ -12,9 +12,9 @@
 import logging
 from typing import Any, TypeVar, Optional
 
-from ..event import eventNames
-from .threadless import Threadless
-from ...common.types import HostPort, TcpOrTlsSocket
+from ...event import eventNames
+from ..threadless import Threadless
+from ....common.types import HostPort, TcpOrTlsSocket
 
 
 T = TypeVar('T')
@@ -23,12 +23,10 @@
 
 
 class ThreadlessFdExecutor(Threadless[T]):
+    """A threadless executor which handles file descriptors
+    and works with read/write events over a socket."""
 
-    def work(
-            self,
-            *args: Any,
-            **kwargs: Any,
-    ) -> None:
+    def work(self, **kwargs: Any) -> None:
         fileno: int = kwargs['fileno']
         addr: Optional[HostPort] = kwargs.get('addr', None)
         conn: Optional[TcpOrTlsSocket] = \

proxy/core/work/fd/local.py

@@ -0,0 +1,50 @@
+# -*- 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.
+"""
+import queue
+import asyncio
+import contextlib
+from typing import Any, Optional
+
+from .fd import ThreadlessFdExecutor
+from ....common.backports import NonBlockingQueue
+
+
+class LocalFdExecutor(ThreadlessFdExecutor[NonBlockingQueue]):
+    """A threadless executor implementation which uses a queue to receive new work."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+
+    @property
+    def loop(self) -> Optional[asyncio.AbstractEventLoop]:
+        if self._loop is None:
+            self._loop = asyncio.get_event_loop_policy().new_event_loop()
+        return self._loop
+
+    def work_queue_fileno(self) -> Optional[int]:
+        return None
+
+    def receive_from_work_queue(self) -> bool:
+        with contextlib.suppress(queue.Empty):
+            work = self.work_queue.get()
+            if isinstance(work, bool) and work is False:
+                return True
+            self.initialize(work)
+        return False
+
+    def initialize(self, work: Any) -> None:
+        assert isinstance(work, tuple)
+        conn, addr = work
+        # NOTE: Here we are assuming to receive a connection object
+        # and not a fileno because we are a LocalExecutor.
+        fileno = conn.fileno()
+        self.work(fileno=fileno, addr=addr, conn=conn)

proxy/core/work/remote.py proxy/core/work/fd/remote.pyproxy/core/work/remote.py renamed to proxy/core/work/fd/remote.py

@@ -9,18 +9,14 @@
     :license: BSD, see LICENSE for more details.
 """
 import asyncio
-import logging
 from typing import Any, Optional
 from multiprocessing import connection
 from multiprocessing.reduction import recv_handle
 
 from .fd import ThreadlessFdExecutor
 
 
-logger = logging.getLogger(__name__)
-
-
-class RemoteExecutor(ThreadlessFdExecutor[connection.Connection]):
+class RemoteFdExecutor(ThreadlessFdExecutor[connection.Connection]):
     """A threadless executor implementation which receives work over a connection.
 
     NOTE: RemoteExecutor uses ``recv_handle`` to accept file descriptors.
@@ -40,12 +36,6 @@ def loop(self) -> Optional[asyncio.AbstractEventLoop]:
             self._loop = asyncio.get_event_loop_policy().get_event_loop()
         return self._loop
 
-    def work_queue_fileno(self) -> Optional[int]:
-        return self.work_queue.fileno()
-
-    def close_work_queue(self) -> None:
-        self.work_queue.close()
-
     def receive_from_work_queue(self) -> bool:
         # Acceptor will not send address for
         # unix socket domain environments.
@@ -55,3 +45,9 @@ def receive_from_work_queue(self) -> bool:
         fileno = recv_handle(self.work_queue)
         self.work(fileno=fileno, addr=addr)
         return False
+
+    def work_queue_fileno(self) -> Optional[int]:
+        return self.work_queue.fileno()
+
+    def close_work_queue(self) -> None:
+        self.work_queue.close()

proxy/core/work/local.py

@@ -8,22 +8,14 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import queue
 import asyncio
-import logging
-import contextlib
 from typing import Any, Optional
 
-from .fd import ThreadlessFdExecutor
-from ...common.backports import (  # noqa: W0611, F401   pylint: disable=unused-import
-    NonBlockingQueue,
-)
+from .threadless import Threadless
+from ...common.backports import NonBlockingQueue
 
 
-logger = logging.getLogger(__name__)
-
-
-class LocalExecutor(ThreadlessFdExecutor['NonBlockingQueue']):
+class BaseLocalExecutor(Threadless[NonBlockingQueue]):
     """A threadless executor implementation which uses a queue to receive new work."""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -38,16 +30,3 @@ def loop(self) -> Optional[asyncio.AbstractEventLoop]:
 
     def work_queue_fileno(self) -> Optional[int]:
         return None
-
-    def receive_from_work_queue(self) -> bool:
-        with contextlib.suppress(queue.Empty):
-            work = self.work_queue.get()
-            if isinstance(work, bool) and work is False:
-                return True
-            assert isinstance(work, tuple)
-            conn, addr = work
-            # NOTE: Here we are assuming to receive a connection object
-            # and not a fileno because we are a LocalExecutor.
-            fileno = conn.fileno()
-            self.work(fileno=fileno, addr=addr, conn=conn)
-        return False

proxy/core/work/pool.py

@@ -11,16 +11,18 @@
 import logging
 import argparse
 import multiprocessing
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Type, TypeVar, Optional
 from multiprocessing import connection
 
-from .remote import RemoteExecutor
 from ...common.flag import flags
 from ...common.constants import DEFAULT_THREADLESS, DEFAULT_NUM_WORKERS
 
 
 if TYPE_CHECKING:   # pragma: no cover
     from ..event import EventQueue
+    from .threadless import Threadless
+
+T = TypeVar('T', bound='Threadless[Any]')
 
 logger = logging.getLogger(__name__)
 
@@ -70,6 +72,7 @@ class ThreadlessPool:
     def __init__(
         self,
         flags: argparse.Namespace,
+        executor_klass: Type['T'],
         event_queue: Optional['EventQueue'] = None,
     ) -> None:
         self.flags = flags
@@ -79,7 +82,9 @@ def __init__(
         self.work_pids: List[int] = []
         self.work_locks: List['multiprocessing.synchronize.Lock'] = []
         # List of threadless workers
-        self._workers: List[RemoteExecutor] = []
+        self._executor_klass = executor_klass
+        # FIXME: Instead of Any type must be the executor klass
+        self._workers: List[Any] = []
         self._processes: List[multiprocessing.Process] = []
 
     def __enter__(self) -> 'ThreadlessPool':
@@ -115,8 +120,8 @@ def _start_worker(self, index: int) -> None:
         self.work_locks.append(multiprocessing.Lock())
         pipe = multiprocessing.Pipe()
         self.work_queues.append(pipe[0])
-        w = RemoteExecutor(
-            iid=index,
+        w = self._executor_klass(
+            iid=str(index),
             work_queue=pipe[1],
             flags=self.flags,
             event_queue=self.event_queue,

proxy/core/work/threadless.py

@@ -120,7 +120,7 @@ def work_queue_fileno(self) -> Optional[int]:
         raise NotImplementedError()
 
     @abstractmethod
-    def work(self, *args: Any, **kwargs: Any) -> None:
+    def work(self, **kwargs: Any) -> None:
         raise NotImplementedError()
 
     def close_work_queue(self) -> None: