|
8 | 8 | :copyright: (c) 2013-present by Abhinav Singh and contributors. |
9 | 9 | :license: BSD, see LICENSE for more details. |
10 | 10 | """ |
11 | | -import time |
| 11 | +import sys |
12 | 12 | import argparse |
13 | | -import threading |
14 | | -import multiprocessing |
15 | | -from typing import Any |
16 | 13 |
|
17 | | -from proxy.core.work import ( |
18 | | - Work, ThreadlessPool, BaseLocalExecutor, BaseRemoteExecutor, |
19 | | -) |
| 14 | +from proxy.core.work import ThreadlessPool |
20 | 15 | from proxy.common.flag import FlagParser |
21 | | -from proxy.common.backports import NonBlockingQueue |
22 | | - |
23 | | - |
24 | | -class Task: |
25 | | - """This will be our work object.""" |
26 | | - |
27 | | - def __init__(self, payload: bytes) -> None: |
28 | | - self.payload = payload |
29 | | - print(payload) |
30 | | - |
31 | | - |
32 | | -class TaskWork(Work[Task]): |
33 | | - """This will be our handler class, created for each received work.""" |
34 | | - |
35 | | - @staticmethod |
36 | | - def create(*args: Any) -> Task: |
37 | | - """Work core doesn't know how to create work objects for us, so |
38 | | - we must provide an implementation of create method here.""" |
39 | | - return Task(*args) |
40 | | - |
41 | | - |
42 | | -class LocalTaskExecutor(BaseLocalExecutor): |
43 | | - """We'll define a local executor which is capable of receiving |
44 | | - log lines over a non blocking queue.""" |
45 | | - |
46 | | - def work(self, *args: Any) -> None: |
47 | | - task_id = int(time.time()) |
48 | | - uid = '%s-%s' % (self.iid, task_id) |
49 | | - self.works[task_id] = self.create(uid, *args) |
50 | | - |
51 | | - |
52 | | -class RemoteTaskExecutor(BaseRemoteExecutor): |
53 | | - |
54 | | - def work(self, *args: Any) -> None: |
55 | | - task_id = int(time.time()) |
56 | | - uid = '%s-%s' % (self.iid, task_id) |
57 | | - self.works[task_id] = self.create(uid, *args) |
58 | | - |
59 | | - |
60 | | -def start_local(flags: argparse.Namespace) -> None: |
61 | | - work_queue = NonBlockingQueue() |
62 | | - executor = LocalTaskExecutor(iid=1, work_queue=work_queue, flags=flags) |
| 16 | +from proxy.core.work.task import ( |
| 17 | + RemoteTaskExecutor, ThreadedTaskExecutor, SingleProcessTaskExecutor, |
| 18 | +) |
63 | 19 |
|
64 | | - t = threading.Thread(target=executor.run) |
65 | | - t.daemon = True |
66 | | - t.start() |
67 | 20 |
|
68 | | - try: |
| 21 | +def start_local_thread(flags: argparse.Namespace) -> None: |
| 22 | + with ThreadedTaskExecutor(flags=flags) as thread: |
69 | 23 | i = 0 |
70 | 24 | while True: |
71 | | - work_queue.put(('%d' % i).encode('utf-8')) |
| 25 | + thread.executor.work_queue.put(('%d' % i).encode('utf-8')) |
72 | 26 | i += 1 |
73 | | - except KeyboardInterrupt: |
74 | | - pass |
75 | | - finally: |
76 | | - executor.running.set() |
77 | | - t.join() |
78 | 27 |
|
79 | 28 |
|
80 | | -def start_remote(flags: argparse.Namespace) -> None: |
81 | | - pipe = multiprocessing.Pipe() |
82 | | - work_queue = pipe[0] |
83 | | - executor = RemoteTaskExecutor(iid=1, work_queue=pipe[1], flags=flags) |
| 29 | +def start_remote_process(flags: argparse.Namespace) -> None: |
| 30 | + with SingleProcessTaskExecutor(flags=flags) as process: |
| 31 | + i = 0 |
| 32 | + while True: |
| 33 | + process.work_queue.send(('%d' % i).encode('utf-8')) |
| 34 | + i += 1 |
84 | 35 |
|
85 | | - p = multiprocessing.Process(target=executor.run) |
86 | | - p.daemon = True |
87 | | - p.start() |
88 | 36 |
|
89 | | - try: |
| 37 | +def start_remote_pool(flags: argparse.Namespace) -> None: |
| 38 | + with ThreadlessPool(flags=flags, executor_klass=RemoteTaskExecutor) as pool: |
90 | 39 | i = 0 |
91 | 40 | while True: |
| 41 | + work_queue = pool.work_queues[i % flags.num_workers] |
92 | 42 | work_queue.send(('%d' % i).encode('utf-8')) |
93 | 43 | i += 1 |
94 | | - except KeyboardInterrupt: |
95 | | - pass |
96 | | - finally: |
97 | | - executor.running.set() |
98 | | - p.join() |
99 | 44 |
|
100 | 45 |
|
101 | | -def start_remote_pool(flags: argparse.Namespace) -> None: |
102 | | - with ThreadlessPool(flags=flags, executor_klass=RemoteTaskExecutor) as pool: |
103 | | - try: |
104 | | - i = 0 |
105 | | - while True: |
106 | | - work_queue = pool.work_queues[i % flags.num_workers] |
107 | | - work_queue.send(('%d' % i).encode('utf-8')) |
108 | | - i += 1 |
109 | | - except KeyboardInterrupt: |
110 | | - pass |
| 46 | +def main() -> None: |
| 47 | + try: |
| 48 | + flags = FlagParser.initialize( |
| 49 | + sys.argv[2:] + ['--disable-http-proxy'], |
| 50 | + work_klass='proxy.core.work.task.TaskHandler', |
| 51 | + ) |
| 52 | + globals()['start_%s' % sys.argv[1]](flags) |
| 53 | + except KeyboardInterrupt: |
| 54 | + pass |
111 | 55 |
|
112 | 56 |
|
113 | 57 | # TODO: TaskWork, LocalTaskExecutor, RemoteTaskExecutor |
114 | 58 | # should not be needed, abstract those pieces out in the core |
115 | 59 | # for stateless tasks. |
116 | 60 | if __name__ == '__main__': |
117 | | - flags = FlagParser.initialize( |
118 | | - ['--disable-http-proxy'], |
119 | | - work_klass=TaskWork, |
120 | | - ) |
121 | | - start_remote_pool(flags) |
122 | | - # start_remote(flags) |
123 | | - # start_local(flags) |
| 61 | + if len(sys.argv) < 2: |
| 62 | + print( |
| 63 | + '\n'.join([ |
| 64 | + 'Usage:', |
| 65 | + ' %s <execution-mode>' % sys.argv[0], |
| 66 | + ' execution-mode can be one of the following:', |
| 67 | + ' "remote_pool", "remote_process", "local_thread"', |
| 68 | + ]), |
| 69 | + ) |
| 70 | + sys.exit(1) |
| 71 | + main() |
![pre-commit-ci[bot]](https://avatars.githubusercontent.com/in/68672?v=4&size=40){kind=avatar}
0 commit comments