|
| 1 | +import asyncio |
| 2 | +import signal |
| 3 | +import time |
| 4 | +from typing import Optional, Dict |
| 5 | +from logger import CustomLogger |
| 6 | +from utils import check_ports_availability |
| 7 | + |
| 8 | + |
| 9 | +# Define the signals that will be handled by the AsyncLoop class |
| 10 | +HANDLED_SIGNALS = ( |
| 11 | + signal.SIGINT, # Unix signal 2. Sent by Ctrl+C. |
| 12 | + signal.SIGTERM, # Unix signal 15. Sent by `kill <pid>`. |
| 13 | + signal.SIGSEGV, # Unix signal 11. Caused by an invalid memory reference. |
| 14 | +) |
| 15 | + |
| 16 | + |
| 17 | +class AsyncLoop: |
| 18 | + """ |
| 19 | + Async loop to run a microservice asynchronously. |
| 20 | + This class is designed to handle the running of a microservice in an asynchronous manner. |
| 21 | + It sets up an event loop and handles certain signals to gracefully stop the service. |
| 22 | + """ |
| 23 | + |
| 24 | + def __init__(self, args: Optional[Dict] = None) -> None: |
| 25 | + """ |
| 26 | + Initialize the AsyncLoop class. |
| 27 | + This method sets up the initial state of the AsyncLoop, including setting up the event loop and signal handlers. |
| 28 | + """ |
| 29 | + self.args = args |
| 30 | + if args.get('name', None): |
| 31 | + self.name = f'{args.get("name")}/{self.__class__.__name__}' |
| 32 | + else: |
| 33 | + self.name = self.__class__.__name__ |
| 34 | + self.protocol = args.get('protocol', 'http') |
| 35 | + self.host = args.get('host', 'localhost') |
| 36 | + self.port = args.get('port', 8080) |
| 37 | + self.quiet_error = args.get('quiet_error', False) |
| 38 | + self.logger = CustomLogger(self.name) |
| 39 | + self._loop = asyncio.new_event_loop() |
| 40 | + asyncio.set_event_loop(self._loop) |
| 41 | + self.is_cancel = asyncio.Event() |
| 42 | + self.logger.info(f'Setting signal handlers') |
| 43 | + |
| 44 | + def _cancel(signum, frame): |
| 45 | + """ |
| 46 | + Signal handler for the AsyncLoop class. |
| 47 | + This method is called when a signal is received. It sets the is_cancel event to stop the event loop. |
| 48 | + """ |
| 49 | + self.logger.info(f'Received signal {signum}') |
| 50 | + self.is_cancel.set(), |
| 51 | + |
| 52 | + for sig in HANDLED_SIGNALS: |
| 53 | + signal.signal(sig, _cancel) |
| 54 | + |
| 55 | + self._start_time = time.time() |
| 56 | + self._loop.run_until_complete(self.async_setup()) |
| 57 | + |
| 58 | + def run_forever(self): |
| 59 | + """ |
| 60 | + Running method to block the main thread. |
| 61 | + This method runs the event loop until a Future is done. It is designed to be called in the main thread to keep it busy. |
| 62 | + """ |
| 63 | + self._loop.run_until_complete(self._loop_body()) |
| 64 | + |
| 65 | + def teardown(self): |
| 66 | + """ |
| 67 | + Call async_teardown() and stop and close the event loop. |
| 68 | + This method is responsible for tearing down the event loop. It first calls the async_teardown method to perform |
| 69 | + any necessary cleanup, then it stops and closes the event loop. It also logs the duration of the event loop. |
| 70 | + """ |
| 71 | + self._loop.run_until_complete(self.async_teardown()) |
| 72 | + self._loop.stop() |
| 73 | + self._loop.close() |
| 74 | + self._stop_time = time.time() |
| 75 | + self.logger.info(f"Async loop is tore down. Duration: {self._stop_time - self._start_time}") |
| 76 | + |
| 77 | + def _get_server(self): |
| 78 | + """ |
| 79 | + Get the server instance based on the protocol. |
| 80 | + This method currently only supports HTTP services. It creates an instance of the HTTPService class with the |
| 81 | + necessary arguments. |
| 82 | + In the future, it will also support gRPC services. |
| 83 | + """ |
| 84 | + if self.protocol.lower() == 'http': |
| 85 | + from http_service import HTTPService |
| 86 | + |
| 87 | + runtime_args = self.args.get('runtime_args', None) |
| 88 | + runtime_args['protocol'] = self.protocol |
| 89 | + runtime_args['host'] = self.host |
| 90 | + runtime_args['port'] = self.port |
| 91 | + return HTTPService( |
| 92 | + uvicorn_kwargs=self.args.get('uvicorn_kwargs', None), |
| 93 | + runtime_args=runtime_args, |
| 94 | + cors=self.args.get('cors', None), |
| 95 | + ) |
| 96 | + |
| 97 | + async def async_setup(self): |
| 98 | + """ |
| 99 | + The async method setup the runtime. |
| 100 | + This method is responsible for setting up the server. It first checks if the port is available, then it gets |
| 101 | + the server instance and initializes it. |
| 102 | + """ |
| 103 | + if not (check_ports_availability(self.host, self.port)): |
| 104 | + raise RuntimeError(f'port:{self.port}') |
| 105 | + |
| 106 | + self.server = self._get_server() |
| 107 | + await self.server.initialize_server() |
| 108 | + |
| 109 | + async def async_run_forever(self): |
| 110 | + """ |
| 111 | + Running method of the server. |
| 112 | + """ |
| 113 | + await self.server.execute_server() |
| 114 | + |
| 115 | + async def async_teardown(self): |
| 116 | + """ |
| 117 | + Terminate the server. |
| 118 | + """ |
| 119 | + await self.server.terminate_server() |
| 120 | + |
| 121 | + async def _wait_for_cancel(self): |
| 122 | + """ |
| 123 | + Wait for the cancellation event. |
| 124 | + This method waits for the is_cancel event to be set. If the server has a _should_exit attribute, it will also |
| 125 | + wait for that to be True. Once either of these conditions is met, it will call the async_teardown method. |
| 126 | + """ |
| 127 | + if isinstance(self.is_cancel, asyncio.Event) and not hasattr( |
| 128 | + self.server, '_should_exit' |
| 129 | + ): |
| 130 | + await self.is_cancel.wait() |
| 131 | + else: |
| 132 | + while not self.is_cancel.is_set() and not getattr( |
| 133 | + self.server, '_should_exit', False |
| 134 | + ): |
| 135 | + await asyncio.sleep(0.1) |
| 136 | + |
| 137 | + await self.async_teardown() |
| 138 | + |
| 139 | + async def _loop_body(self): |
| 140 | + """ |
| 141 | + The main body of the event loop. |
| 142 | + This method runs the async_run_forever and _wait_for_cancel methods concurrently. If a CancelledError is raised, |
| 143 | + it logs a warning message. |
| 144 | + """ |
| 145 | + try: |
| 146 | + await asyncio.gather(self.async_run_forever(), self._wait_for_cancel()) |
| 147 | + except asyncio.CancelledError: |
| 148 | + self.logger.warning('received terminate ctrl message from main process') |
| 149 | + |
| 150 | + def __enter__(self): |
| 151 | + """ |
| 152 | + Enter method for the context manager. |
| 153 | + This method simply returns the instance itself. |
| 154 | + """ |
| 155 | + return self |
| 156 | + |
| 157 | + def __exit__(self, exc_type, exc_val, exc_tb): |
| 158 | + """ |
| 159 | + Exit method for the context manager. |
| 160 | + This method handles any exceptions that occurred within the context. If a KeyboardInterrupt was raised, it logs |
| 161 | + an info message. If any other exception was raised, it logs an error message. Finally, it attempts to call the |
| 162 | + teardown method. If an OSError is raised during this, it is ignored. Any other exceptions are logged as errors. |
| 163 | + """ |
| 164 | + if exc_type == KeyboardInterrupt: |
| 165 | + self.logger.info(f'{self!r} is interrupted by user') |
| 166 | + elif exc_type and issubclass(exc_type, Exception): |
| 167 | + self.logger.error( |
| 168 | + ( |
| 169 | + f'{exc_val!r} during {self.run_forever!r}' |
| 170 | + + f'\n add "--quiet-error" to suppress the exception details' |
| 171 | + if not self.quiet_error |
| 172 | + else '' |
| 173 | + ), |
| 174 | + exc_info=not self.quiet_error, |
| 175 | + ) |
| 176 | + else: |
| 177 | + self.logger.info(f'{self!r} is ended') |
| 178 | + |
| 179 | + return True |
0 commit comments