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')
