"""Queue worker. ================ """ import logging from queue import Empty, Full from queue import Queue from threading import Thread import time from ddtrace import tracer from ddtrace.context import Context from salessheets.connectors import sqs from salessheets.connectors.sentry import sentry_capture_exception from salessheets.errors import FailedToStartError from salessheets.logic.pdf_generation import generate_pdf RETURN_TO_QUEUE_TIMEOUT = 0 PREFETCH_QUEUE_WAIT_TIME = 1 PREFETCH_QUEUE_MAX_WAIT_TIME = 30 MAX_RETRIES = 10 class BaseWorker(Thread): """Base class for thread workers.""" @tracer.wrap('BaseWorker.__init__') def __init__(self, name, prefetch_queue, logger): """Create a new worker. Args: name (str): worker name prefetch_queue (Queue): queue for prefetching messages logger (LoggingAdapter): log adapter """ super(BaseWorker, self).__init__(name=name) self._should_work = False self.prefetch_queue = prefetch_queue self.logger = logger self.debug = False self.name = name current_span = tracer.current_span() parent_ctx = current_span.context self.trace_context = Context( trace_id=current_span.trace_id, span_id=current_span.span_id, sampling_priority=parent_ctx.sampling_priority, dd_origin=parent_ctx.dd_origin, ) @tracer.wrap('BaseWorker.start') def start(self): """Start the queue worker.""" self._log('Thread starting') self._should_work = True super(BaseWorker, self).start() self._log('Thread started') @tracer.wrap('BaseWorker.stop') def stop(self, timeout): """Stop the queue worker. Args: timeout (int): Timeout in seconds """ self._log('Stopping thread') self._should_work = False try: self.join(timeout=timeout) except RuntimeError: if sentry_capture_exception: sentry_capture_exception() self._log('Thread stopped') self._log('Thread stopped') tracer._writer.flush_queue() @tracer.wrap('BaseWorker.run') def run(self): """Run the worker operation.""" tracer.context_provider.activate(self.trace_context) with tracer.trace(name=self.name, service='BaseWorker', span_type='worker') as _: while self._should_work: try: self._do_work() except Exception: if sentry_capture_exception: sentry_capture_exception() @tracer.wrap('BaseWorker._do_work') def _do_work(self): """Actual piece of work. To be overriden in derived classes.""" pass def _log(self, message, level=logging.INFO): """Write the logs. Adds additional information (worker name) to log message, and skip message emitting for non-debug mode. Args: message (str): log message level (int): log level """ if level == logging.INFO and not self.debug: return self.logger.log(level, '{0}: {1}'.format(self.name, message)) class QueueReaderWorker(BaseWorker): """Worker for reading the ownership change request messages from SQS queue. QueueReaderWorker polls messages from SQS queue and saves them into local prefetch queue. Messages are being polled as batches (num_messages arg specifies the number of messages in a batch). If the prefetch queue becomes full, before the reader has pulled the next batch, it will wait for PREFETCH_QUEUE_WAIT_TIME seconds before trying to get the new messages. If the prefetch queue becomes full, after the batch has already been pulled and during its processing, then the visibility_timeout will be reset on messages of that batch (meaning that they will become visible in SQS queue again and could be picked up by workers again). """ @tracer.wrap('QueueReaderWorker.__init__') def __init__( self, name, prefetch_queue, sqs_queue, visibility_timeout, wait_time_seconds, num_messages, return_to_queue_timeout, max_retries, logger): """Create a new QueueReaderWorker. Args: name (str): worker name prefetch_queue (Queue): prefetch queue to put messages to sqs_queue: SQS queue to poll messaged from visibility_timeout (int): SQS visibility timeout wait_time_seconds (int): SQS long-polling interval num_messages (int): number of messages to get from SQS in a batch return_to_queue_timeout (int): visibility timeout that should be set to message if its needed to be returned to SQS queue max_retries (int): maximum number of retries of message processing logger (LoggingAdapter): log adapter """ super(QueueReaderWorker, self).__init__(name, prefetch_queue, logger) self.sqs_queue = sqs_queue self.visibility_timeout = visibility_timeout self.wait_time_seconds = wait_time_seconds self.num_messages = num_messages self.return_to_queue_timeout = return_to_queue_timeout self.max_retries = max_retries self.current_sleep_time = PREFETCH_QUEUE_WAIT_TIME @tracer.wrap('QueueReaderWorker._do_work') def _do_work(self): with tracer.trace(name=self.name, service='QueueReaderWorker', span_type='worker') as _: if self.prefetch_queue.full(): self._log('Prefetch queue full, waiting for {0} seconds'.format( self.current_sleep_time)) time.sleep(self.current_sleep_time) self.current_sleep_time += 0.5 if self.current_sleep_time > PREFETCH_QUEUE_MAX_WAIT_TIME: self.current_sleep_time = PREFETCH_QUEUE_WAIT_TIME return self.current_sleep_time = PREFETCH_QUEUE_WAIT_TIME num_messages = self.num_messages - self.prefetch_queue.qsize() self._log('receiving {0} SQS messages'.format(num_messages)) sqs_conn = sqs.get_connection() resp = sqs_conn.receive_message( QueueUrl=self.sqs_queue, MaxNumberOfMessages=num_messages, MessageAttributeNames=[ sqs.CORRELATION_ID_ATTRIBUTE, sqs.FEATURE_FLAG_USER_CONTEXT_ATTRIBUTE], VisibilityTimeout=self.visibility_timeout, WaitTimeSeconds=self.wait_time_seconds, AttributeNames=[sqs.APPROXIMATE_RECEIVE_COUNT_ATTRIBUTE] ) if 'Messages' in resp: messages = resp['Messages'] else: messages = [] self._log('{0} messages received'.format(len(messages))) for message in messages: # Don't prefetch anything, if we are shutting down. if not self._should_work: return try: self.prefetch_queue.put( message, timeout=PREFETCH_QUEUE_WAIT_TIME) self._log('Message enqueued to prefetch_queue') except Full: message.change_visibility(self.return_to_queue_timeout) self._log('Message visibility set to 0, returned to SQS') class PdfGenerationWorker(BaseWorker): """Worker that executes requests to API and persists the results. PdfGenerationWorker gets messages from prefetch queue and then processes them using the external processing module ( salessheets.logic.message). """ @tracer.wrap('PdfGenerationWorker.__init__') def __init__( self, name, prefetch_queue, sqs_queue, return_to_queue_timeout, logger): """Create a new ApiExecutorWorker. Args: name (str): worker name prefetch_queue (Queue): queue for prefetching messages sqs_queue (boto.sqs.queue.Queue): SQS queue to delete messages from return_to_queue_timeout (int): visibility timeout that should be set to message if its needed to be returned to SQS queue logger (LoggingAdapter): log adapter """ super(PdfGenerationWorker, self).__init__(name, prefetch_queue, logger) self.sqs_queue = sqs_queue self.return_to_queue_timeout = return_to_queue_timeout @tracer.wrap('PdfGenerationWorker._do_work') def _do_work(self): with tracer.trace(name=self.name, service='AppProcessingWorker', span_type='worker') as _: if not self._should_work: return try: message = self.prefetch_queue.get(timeout=PREFETCH_QUEUE_WAIT_TIME) except Empty: return try: generate_pdf(message) sqs_conn = sqs.get_connection() sqs_conn.delete_message( QueueUrl=self.sqs_queue, ReceiptHandle=message['ReceiptHandle'] ) self._log('Message deleted.') finally: self.prefetch_queue.task_done() class QueueWorkerManager(object): # todo: remove unneeded dependencies """Workers manager. Creates and starts the predefined number of workers. """ @tracer.wrap('QueueWorkerManager.__init__') def __init__( self, readers_count, executors_count, queue_name, prefetch_count, visibility_timeout, wait_time_seconds, num_messages, app_logger): """Create a new QueueWorkerManager. Args: readers_count (int): number of QueueReaderWorkers executors_count (int): number of ApiExecutorWorkers queue_name (str): SQS queue name prefetch_count (int): mx length of the prefetch queue visibility_timeout (int): SQS visibility timeout wait_time_seconds (int): SQS long-polling interval num_messages (int): SQS number of messages to poll in a batch app_logger (Logger): logging adapter """ self.readers_count = readers_count self.executors_count = executors_count self.readers_started = 0 self.executors_started = 0 self.queue_name = queue_name self.workers = [] # todo: leave only worker's launch self.prefetch_queue = Queue(prefetch_count) self.visibility_timeout = visibility_timeout self.wait_time_seconds = wait_time_seconds self.num_messages = num_messages self.logger = app_logger self.debug = False @tracer.wrap('QueueWorkerManager.start') def start(self): """Start the workers manager. Create required number of QueueReader and ApiExecutor workers and start them. Raises: FailedToStartError: raised when no readers and executors has started successfully. """ self.logger.info('Starting queue workers.') for executor_id in range(0, self.executors_count): self._start_pdf_generator(executor_id) for reader_id in range(0, self.readers_count): self._start_queue_reader(reader_id) if self.readers_started == 0 and self.executors_started == 0: self.logger.critical('No readers and executors started.') raise FailedToStartError('No readers and executors started.') self.logger.info( 'Queue workers started. {0} queue readers and {1} api executors' .format(self.readers_started, self.executors_started)) @tracer.wrap('QueueWorkerManager._start_queue_reader') def _start_queue_reader(self, reader_id): """Start the QueueReader worker. Args: reader_id (str, int): id of the worker """ worker_name = 'salessheets-queue-reader-worker{0}'.format(reader_id) sqs_queue = sqs.get_queue(self.queue_name) if not sqs_queue: self.logger.error('Worker {0} can not be started, because SQS ' 'queue is not accessible'.format(reader_id)) return worker = QueueReaderWorker( worker_name, self.prefetch_queue, sqs_queue, self.visibility_timeout, self.wait_time_seconds, self.num_messages, RETURN_TO_QUEUE_TIMEOUT, MAX_RETRIES, self.logger ) worker.debug = self.debug worker.start() self.workers.append(worker) self.readers_started += 1 @tracer.wrap('QueueWorkerManager._start_pdf_generator') def _start_pdf_generator(self, executor_id): """Start the PdfGenerationWorker. Args: executor_id (str, int): id of the worker """ worker_name = 'salessheets-queue-pdf-generator-worker{0}'\ .format(executor_id) sqs_queue = sqs.get_queue(self.queue_name) if not sqs_queue: self.logger.error('Worker {0} can not be started, because SQS ' 'queue is not accessible'.format(executor_id)) return worker = PdfGenerationWorker( worker_name, self.prefetch_queue, sqs_queue, RETURN_TO_QUEUE_TIMEOUT, self.logger ) worker.debug = self.debug worker.start() self.workers.append(worker) self.executors_started += 1 @tracer.wrap('QueueWorkerManager.stop') def stop(self, timeout): """Stop the manager. Args: timeout (int): timeout for stopping workers """ self.logger.info('Stopping queue workers...') for worker in self.workers: worker.stop(timeout) tracer._writer.flush_queue() self.logger.info('Queue workers stopped')