From 96e8ac4316ac57aac90066f35422d333c532513b Mon Sep 17 00:00:00 2001 From: Viktor Seifert Date: Mon, 25 Jul 2022 15:07:40 +0200 Subject: [PATCH] RED-4653: Added code to cancel the queue subscription on application exit to queue manager so that it can exit gracefully --- pyinfra/queue/queue_manager.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/pyinfra/queue/queue_manager.py b/pyinfra/queue/queue_manager.py index f226e0b..e6e0cef 100644 --- a/pyinfra/queue/queue_manager.py +++ b/pyinfra/queue/queue_manager.py @@ -1,9 +1,12 @@ +import atexit import json import logging - -import pika, pika.exceptions +import signal from typing import Callable +import pika +import pika.exceptions + from pyinfra.config import get_config, Config CONFIG = get_config() @@ -32,6 +35,10 @@ class QueueManager(object): def __init__(self, config: Config = CONFIG): connection_params = get_connection_params(config) + atexit.register(self.cancel_queue_subscription) + signal.signal(signal.SIGTERM, self.cancel_queue_subscription) + signal.signal(signal.SIGINT, self.cancel_queue_subscription) + self._connection = pika.BlockingConnection(parameters=connection_params) self._channel = self._connection.channel() self._channel.basic_qos(prefetch_count=1) @@ -63,13 +70,14 @@ class QueueManager(object): self.cancel_queue_subscription() def cancel_queue_subscription(self): - if self._consumer_token: + if self._consumer_token and self._connection: self.logger.info(f"Cancelling subscription for consumer-tag: {self._consumer_token}") self._channel.basic_cancel(self._consumer_token) + self._connection.close() self._consumer_token = None else: - self.logger.warning("No consumer-tag found, a possible subscription cannot be cancelled") + self.logger.info("No consumer-tag found, a possible subscription cannot be cancelled") def _create_queue_callback(self, process_message_callback: Callable): def callback(_channel, frame, properties, body):