async def producer( subscription: str, message_queue: MessageQueue, subscriber_client: 'SubscriberClient', max_messages: int, metrics_client: MetricsAgent) -> None: try: while True: new_messages = [] try: pull_task = asyncio.ensure_future( subscriber_client.pull( subscription=subscription, max_messages=max_messages, # it is important to have this value reasonably # high as long lived connections may be left # hanging on a server which will cause delay in # message delivery or even false deadlettering if # it is enabled timeout=30)) new_messages = await asyncio.shield(pull_task) except (asyncio.TimeoutError, KeyError): continue metrics_client.histogram( 'pubsub.producer.batch', len(new_messages)) pulled_at = time.perf_counter() while new_messages: await message_queue.put((new_messages[-1], pulled_at)) new_messages.pop() await message_queue.join() except asyncio.CancelledError: log.info('Producer worker cancelled. Gracefully terminating...') if not pull_task.done(): # Leaving the connection hanging can result in redelivered # messages, so try to finish before shutting down try: new_messages += await asyncio.wait_for(pull_task, 5) except (asyncio.TimeoutError, KeyError): pass pulled_at = time.perf_counter() for m in new_messages: await message_queue.put((m, pulled_at)) await message_queue.join() log.info('Producer terminated gracefully.') raise
async def nacker(subscription: str, nack_queue: 'asyncio.Queue[str]', subscriber_client: 'SubscriberClient', nack_window: float, metrics_client: MetricsAgent) -> None: ack_ids: List[str] = [] while True: if not ack_ids: ack_ids.append(await nack_queue.get()) nack_queue.task_done() ack_ids += await _budgeted_queue_get(nack_queue, nack_window) # modifyAckDeadline endpoint limit is 524288 bytes # which is ~2744 ack_ids if len(ack_ids) > 2500: log.error( 'nacker is falling behind, dropping %d unacked messages', len(ack_ids) - 2500) ack_ids = ack_ids[-2500:] try: await subscriber_client.modify_ack_deadline( subscription, ack_ids=ack_ids, ack_deadline_seconds=0) except asyncio.CancelledError: # pylint: disable=try-except-raise raise except aiohttp.client_exceptions.ClientResponseError as e: if e.status == 400: log.error( 'Nack error is unrecoverable, ' 'one or more messages may be dropped', exc_info=e) async def maybe_nack(ack_id: str) -> None: try: await subscriber_client.modify_ack_deadline( subscription, ack_ids=[ack_id], ack_deadline_seconds=0) except Exception as e: log.warning('Nack failed for ack_id=%s', ack_id, exc_info=e) for ack_id in ack_ids: asyncio.ensure_future(maybe_nack(ack_id)) ack_ids = [] log.warning('Nack request failed, better luck next batch', exc_info=e) metrics_client.increment('pubsub.nacker.batch.failed') continue except Exception as e: log.warning('Nack request failed, better luck next batch', exc_info=e) metrics_client.increment('pubsub.nacker.batch.failed') continue metrics_client.histogram('pubsub.nacker.batch', len(ack_ids)) ack_ids = []
async def producer(subscription: str, message_queue: MessageQueue, subscriber_client: 'SubscriberClient', max_messages: int, metrics_client: MetricsAgent) -> None: try: while True: new_messages = [] try: new_messages = await subscriber_client.pull( subscription=subscription, max_messages=max_messages, # it is important to have this value reasonably high # as long lived connections may be left hanging # on a server which will cause delay in message # delivery or even false deadlettering if it is enabled timeout=30) except (asyncio.TimeoutError, KeyError): continue metrics_client.histogram('pubsub.producer.batch', len(new_messages)) pulled_at = time.perf_counter() while new_messages: await message_queue.put((new_messages[-1], pulled_at)) new_messages.pop() await message_queue.join() except asyncio.CancelledError: log.info('Producer worker cancelled. Gracefully terminating...') pulled_at = time.perf_counter() for m in new_messages: await message_queue.put((m, pulled_at)) await message_queue.join() log.info('Producer terminated gracefully.') raise
async def _execute_callback(message: SubscriberMessage, callback: ApplicationHandler, ack_queue: 'asyncio.Queue[str]', nack_queue: 'Optional[asyncio.Queue[str]]', metrics_client: MetricsAgent) -> None: try: start = time.perf_counter() await callback(message) await ack_queue.put(message.ack_id) metrics_client.increment('pubsub.consumer.succeeded') metrics_client.histogram('pubsub.consumer.latency.runtime', time.perf_counter() - start) except asyncio.CancelledError: if nack_queue: await nack_queue.put(message.ack_id) log.warning('Application callback was cancelled') metrics_client.increment('pubsub.consumer.cancelled') except Exception: if nack_queue: await nack_queue.put(message.ack_id) log.exception('Application callback raised an exception') metrics_client.increment('pubsub.consumer.failed')
async def subscribe(subscription: str, # pylint: disable=too-many-locals handler: ApplicationHandler, subscriber_client: SubscriberClient, *, num_producers: int = 1, max_messages_per_producer: int = 100, ack_window: float = 0.3, ack_deadline_cache_timeout: float = float('inf'), num_tasks_per_consumer: int = 1, enable_nack: bool = True, nack_window: float = 0.3, metrics_client: Optional[MetricsAgent] = None ) -> None: ack_queue: 'asyncio.Queue[str]' = asyncio.Queue( maxsize=(max_messages_per_producer * num_producers)) nack_queue: 'Optional[asyncio.Queue[str]]' = None ack_deadline_cache = AckDeadlineCache(subscriber_client, subscription, ack_deadline_cache_timeout) metrics_client = metrics_client or MetricsAgent() acker_tasks = [] consumer_tasks = [] producer_tasks = [] try: acker_tasks.append(asyncio.ensure_future( acker(subscription, ack_queue, subscriber_client, ack_window=ack_window, metrics_client=metrics_client) )) if enable_nack: nack_queue = asyncio.Queue( maxsize=(max_messages_per_producer * num_producers)) acker_tasks.append(asyncio.ensure_future( nacker(subscription, nack_queue, subscriber_client, nack_window=nack_window, metrics_client=metrics_client) )) for _ in range(num_producers): q: MessageQueue = asyncio.Queue( maxsize=max_messages_per_producer) consumer_tasks.append(asyncio.ensure_future( consumer(q, handler, ack_queue, ack_deadline_cache, num_tasks_per_consumer, nack_queue, metrics_client=metrics_client) )) producer_tasks.append(asyncio.ensure_future( producer(subscription, q, subscriber_client, max_messages=max_messages_per_producer, metrics_client=metrics_client) )) # TODO: since this is in a `not BUILD_GCLOUD_REST` section, we # shouldn't have to care about py2 support. Using splat syntax # here, though, breaks the coverage.py reporter for this file even # though it would never be loaded at runtime in py2. # all_tasks = [*producer_tasks, *consumer_tasks, *acker_tasks] all_tasks = producer_tasks + consumer_tasks + acker_tasks done, _ = await asyncio.wait(all_tasks, return_when=asyncio.FIRST_COMPLETED) for task in done: task.result() raise Exception('A subscriber worker shut down unexpectedly!') except Exception as e: log.info('Subscriber exited', exc_info=e) for task in producer_tasks: task.cancel() await asyncio.wait(producer_tasks, return_when=asyncio.ALL_COMPLETED) for task in consumer_tasks: task.cancel() await asyncio.wait(consumer_tasks, return_when=asyncio.ALL_COMPLETED) for task in acker_tasks: task.cancel() await asyncio.wait(acker_tasks, return_when=asyncio.ALL_COMPLETED) raise asyncio.CancelledError('Subscriber shut down')
async def subscribe( subscription: str, # pylint: disable=too-many-locals handler: ApplicationHandler, subscriber_client: SubscriberClient, *, num_producers: int = 1, max_messages_per_producer: int = 100, ack_window: float = 0.3, ack_deadline_cache_timeout: int = 60, num_tasks_per_consumer: int = 1, enable_nack: bool = True, nack_window: float = 0.3, metrics_client: Optional[MetricsAgent] = None) -> None: ack_queue: 'asyncio.Queue[str]' = asyncio.Queue( maxsize=(max_messages_per_producer * num_producers)) nack_queue: 'Optional[asyncio.Queue[str]]' = None ack_deadline_cache = AckDeadlineCache(subscriber_client, subscription, ack_deadline_cache_timeout) metrics_client = metrics_client or MetricsAgent() acker_tasks = [] consumer_tasks = [] producer_tasks = [] try: acker_tasks.append( asyncio.ensure_future( acker(subscription, ack_queue, subscriber_client, ack_window=ack_window, metrics_client=metrics_client))) if enable_nack: nack_queue = asyncio.Queue(maxsize=(max_messages_per_producer * num_producers)) acker_tasks.append( asyncio.ensure_future( nacker(subscription, nack_queue, subscriber_client, nack_window=nack_window, metrics_client=metrics_client))) for _ in range(num_producers): q: MessageQueue = asyncio.Queue( maxsize=max_messages_per_producer) consumer_tasks.append( asyncio.ensure_future( consumer(q, handler, ack_queue, ack_deadline_cache, num_tasks_per_consumer, nack_queue, metrics_client=metrics_client))) producer_tasks.append( asyncio.ensure_future( producer(subscription, q, subscriber_client, max_messages=max_messages_per_producer, metrics_client=metrics_client))) all_tasks = [*producer_tasks, *consumer_tasks, *acker_tasks] done, _ = await asyncio.wait(all_tasks, return_when=asyncio.FIRST_COMPLETED) for task in done: task.result() raise Exception('A subscriber worker shut down unexpectedly!') except Exception as e: log.info('Subscriber exited', exc_info=e) for task in producer_tasks: task.cancel() await asyncio.wait(producer_tasks, return_when=asyncio.ALL_COMPLETED) for task in consumer_tasks: task.cancel() await asyncio.wait(consumer_tasks, return_when=asyncio.ALL_COMPLETED) for task in acker_tasks: task.cancel() await asyncio.wait(acker_tasks, return_when=asyncio.ALL_COMPLETED) raise asyncio.CancelledError('Subscriber shut down')