def handle_in_background(self, queue: InternalQueue, handler, command, on_done: asyncio.Event): """Handle a received command by calling the provided handler This execution happens in the background. """ logger.debug(f"Handling command {command}") def when_task_finished(fut: asyncio.Future): self._running_commands.remove(fut) try: # Retrieve any error which may have occurred. # We ignore the error because we assume any exceptions which the # handler threw will have already been placed into the error queue # by the queue_exception_checker(). # Regardless, we must retrieve the result in order to keep Python happy. fut.result() except: pass # We use call_soon_threadsafe() to ensure we call the Event's set() # in a threadsafe fashion. This is because the Event object may have # been created in another thread and be attached to another event loop on_done._loop.call_soon_threadsafe(on_done.set) # fmt: off background_call_task = asyncio.ensure_future(queue_exception_checker( handler(command), self.error_queue, )) # fmt: on background_call_task.add_done_callback(when_task_finished) self._running_commands.add(background_call_task)
def start(self): """Starts the queue monitor""" # fmt: off self._queue_monitor_task = asyncio.ensure_future(queue_exception_checker( self._queue_monitor(), self.error_queue, ))
async def handle_receive_result(self, command: commands.ReceiveResultCommand): """Client wishes to receive a result from a worker""" # TODO: rpc_timeout is in three different places in the config! # Fix this. Really it makes most sense for the use if it goes on the # ApiConfig rather than having to repeat it on both the result & RPC # transports. timeout = command.options.get( "timeout", self.config.api(command.message.api_name).rpc_timeout ) result_transport = self.transport_registry.get_result_transport(command.message.api_name) logger.debug("Starting RPC result listener") task = asyncio.ensure_future( queue_exception_checker( self._result_listener( result_transport=result_transport, timeout=timeout, rpc_message=command.message, return_path=command.message.return_path, options=command.options, result_queue=command.destination_queue, ), self.error_queue, ) ) self.listener_tasks.add(task)
async def _start_listener(self, listener: "Listener"): # Setting the maxsize to 1 ensures the transport cannot load # messages faster than we can consume them queue: InternalQueue[EventMessage] = InternalQueue(maxsize=1) async def consume_events(): while True: logger.debug( "Event listener now waiting for event on the internal queue" ) event_message = await queue.get() logger.debug( "Event listener has now received an event on the internal queue, processing now" ) await self._on_message( event_message=event_message, listener=listener.callable, options=listener.options, on_error=listener.on_error, ) queue.task_done() # Start the consume_events() consumer running task = asyncio.ensure_future( queue_exception_checker(consume_events(), self.error_queue)) self._event_listener_tasks.add(task) await self.producer.send( ConsumeEventsCommand( events=listener.events, destination_queue=queue, listener_name=listener.name, options=listener.options, )).wait()
def start(self, handler: Callable): """Set the handler function and start the invoker Use `stop()` to shutdown the invoker. """ self._consumer_task = asyncio.ensure_future( queue_exception_checker(self._consumer_loop(self.queue, handler), self.error_queue) ) self._running_commands = set()
def start(self, handler: Callable): """Set the handler function and start the invoker Use `stop()` to shutdown the invoker. """ logger.debug( f"Starting consumer for handler {handler.__qualname__}(). This should report ready" " shortly...") self._consumer_task = asyncio.ensure_future( queue_exception_checker(self._consumer_loop(self.queue, handler), self.error_queue)) self._running_commands = set()
def test_queue_exception_checker_in_task(erroring_coroutine, error_queue: ErrorQueueType): coroutine = queue_exception_checker(erroring_coroutine(), error_queue) with pytest.raises(ExampleException): asyncio.run(coroutine) assert error_queue.qsize() == 1 error: Error = error_queue.get_nowait() assert error.type == ExampleException assert isinstance(error.value, ExampleException) assert "test_utilities_unit.py" in str(error) assert "ExampleException" in str(error)
async def handle_consume_rpcs(self, command: commands.ConsumeRpcsCommand): """Worker wishes for incoming RPC results to be listened for and processed""" # Not all APIs will necessarily be served by the same transport, so group them # accordingly api_names_by_transport = self.transport_registry.get_rpc_transports(command.api_names) for rpc_transport, transport_api_names in api_names_by_transport.items(): transport_apis = list(map(self.api_registry.get, transport_api_names)) # fmt: off task = asyncio.ensure_future(queue_exception_checker( self._consume_rpcs_with_transport(rpc_transport=rpc_transport, apis=transport_apis), self.error_queue, )) # fmt: on self.consumer_tasks.add(task)
async def handle_consume_events(self, command: commands.ConsumeEventsCommand): event_transports = self.transport_registry.get_event_transports( api_names=[api_name for api_name, _ in command.events]) async def listener(event_transport, events_): consumer = event_transport.consume( listen_for=events_, listener_name=command.listener_name, error_queue=self.error_queue, **command.options, ) async for event_messages in consumer: for event_message in event_messages: logger.debug( f"Putting event {event_message.id} onto the internal queue" ) await command.destination_queue.put(event_message) # Wait for the queue to be emptied before fetching more. # We will need to make this configurable if we want to support # the pre-fetching of events. This solution is a good default # though as it will ensure events are processed in a more ordered fashion # in cases where there are multiple workers, and also ensure fewer # messages need to be reclaimed in the event of a crash logger.debug( f"Waiting for messages to be consumed before adding more") await command.destination_queue.join() for event_transport_pool, api_names in event_transports.items(): # Create a listener task for each event transport, # passing each a list of events for which it should listen events = [(api_name, event_name) for api_name, event_name in command.events if api_name in api_names] # fmt: off listener_task = asyncio.ensure_future( queue_exception_checker( listener(event_transport_pool, events), self.error_queue, )) # fmt: on self.listener_tasks.add(listener_task)
async def consume( self, listen_for: List[Tuple[str, str]], listener_name: str, error_queue: ErrorQueueType, since: Union[Since, Sequence[Since]] = "$", forever=True, ) -> AsyncGenerator[List[RedisEventMessage], None]: """Consume events for the given APIs""" self._sanity_check_listen_for(listen_for) consumer_group = f"{self.service_name}-{listener_name}" if not isinstance(since, (list, tuple)): # Since has been specified as a single value. Normalise it into # the value-per-listener format. since = [since] * len(listen_for) since = map(normalise_since_value, since) stream_names = self._get_stream_names(listen_for) # Keys are stream names, values as the latest ID consumed from that stream streams = OrderedDict(zip(stream_names, since)) expected_events = {event_name for _, event_name in listen_for} logger.debug( LBullets( L( "Consuming events as consumer {} in group {} on streams", Bold(self.consumer_name), Bold(consumer_group), ), items={"{} ({})".format(*v) for v in streams.items()}, )) # Cleanup any old groups & consumers await self._cleanup(stream_names) # Here we use a queue to combine messages coming from both the # fetch messages loop and the reclaim messages loop. queue = InternalQueue(maxsize=1) initial_reclaiming_complete = asyncio.Event() async def consume_loop(): """Regular event consuming. See _fetch_new_messages()""" logger.debug( "Will begin consuming events once the initial event reclaiming is complete" ) await initial_reclaiming_complete.wait() logger.debug( "Event reclaiming is complete, beginning to consume events") async for messages in self._fetch_new_messages( streams, consumer_group, expected_events, forever): await queue.put(messages) # Wait for the queue to empty before getting trying to get another message await queue.join() retry_consume_loop = retry_on_redis_connection_failure( fn=consume_loop, retry_delay=self.consumption_restart_delay, action="consuming events") async def reclaim_loop(): """ Reclaim messages which other consumers have failed to processes in reasonable time. See _reclaim_lost_messages() """ while True: logger.debug("Checking for any events which need reclaiming") async for messages in self._reclaim_lost_messages( stream_names, consumer_group, expected_events): await queue.put(messages) # Wait for the queue to empty before getting trying to get another message await queue.join() initial_reclaiming_complete.set() await asyncio.sleep(self.reclaim_interval) consume_task = None reclaim_task = None try: # Run the two above coroutines in their own tasks consume_task = asyncio.ensure_future( queue_exception_checker(retry_consume_loop, error_queue)) reclaim_task = asyncio.ensure_future( queue_exception_checker(reclaim_loop(), error_queue)) while True: try: messages = await queue.get() logger.debug( f"Got batch of {len(messages)} message(s). Yielding messages to Lightbus" " client") yield messages logger.debug( f"Batch of {len(messages)} message(s) was processed by Lightbus client." " Marking as done.") queue.task_done() except GeneratorExit: return finally: # Make sure we cleanup the tasks we created await cancel(consume_task, reclaim_task)
async def start_worker(self): """Worker startup procedure """ # Ensure an event loop exists get_event_loop() self._worker_tasks = set() # Start monitoring for errors on the error queue error_monitor_task = asyncio.ensure_future(self.error_monitor()) self._error_monitor_task = error_monitor_task self._worker_tasks.add(self._error_monitor_task) # Features setup & logging if not self.api_registry.all() and Feature.RPCS in self.features: logger.info( "Disabling serving of RPCs as no APIs have been registered") self.features.remove(Feature.RPCS) logger.info( LBullets(f"Enabled features ({len(self.features)})", items=[f.value for f in self.features])) disabled_features = set(ALL_FEATURES) - set(self.features) logger.info( LBullets( f"Disabled features ({len(disabled_features)})", items=[f.value for f in disabled_features], )) # Api logging logger.info( LBullets( "APIs in registry ({})".format(len(self.api_registry.all())), items=self.api_registry.names(), )) # Push all registered APIs into the global schema for api in self.api_registry.all(): await self.schema.add_api(api) # We're running as a worker now (e.g. lightbus run), so # do the lazy loading immediately await self.lazy_load_now() # Setup schema monitoring monitor_task = asyncio.ensure_future( queue_exception_checker(self.schema.monitor(), self.error_queue)) logger.info("Executing before_worker_start & on_start hooks...") await self.hook_registry.execute("before_worker_start") logger.info( "Execution of before_worker_start & on_start hooks was successful") # Setup RPC consumption if Feature.RPCS in self.features: consume_rpc_task = asyncio.ensure_future( queue_exception_checker(self.consume_rpcs(), self.error_queue)) else: consume_rpc_task = None # Start off any registered event listeners if Feature.EVENTS in self.features: await self.event_client.start_registered_listeners() # Start off any background tasks if Feature.TASKS in self.features: for coroutine in self._background_coroutines: task = asyncio.ensure_future( queue_exception_checker(coroutine, self.error_queue)) self._background_tasks.append(task) self._worker_tasks.add(consume_rpc_task) self._worker_tasks.add(monitor_task)