"""Queue worker.""" import gc import logging from queue import Empty from queue import Full from queue import Queue from threading import Thread import time from sentry_sdk import capture_exception from sentry_sdk import capture_message from transcoding.connectors import sqs from transcoding.constants.exceptions import FailedToStartError from transcoding.logic import processing 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.""" 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 def start(self): """Start the queue worker.""" self._log('Thread starting') self._should_work = True super(BaseWorker, self).start() self._log('Thread started') 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 as e: capture_exception(e) self._log('Thread stopped') self._log('Thread stopped') def run(self): """Run the worker operation.""" while self._should_work: try: self._do_work() except Exception as e: capture_exception(e) def _do_work(self): """Actual piece of work. To be overriden in derived classes.""" pass def _log(self, message, level=logging.INFO): """Log adapter wrapper. 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). """ 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 def _do_work(self): 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.prefetch_queue.maxsize - self.prefetch_queue.qsize() num_messages = num_messages if num_messages <= 10 else 10 if num_messages <= 0: self._log('Prefetch queue is full') return self._log('receiving {0} SQS messages'.format(num_messages)) messages = self.sqs_queue.receive_messages( MaxNumberOfMessages=num_messages, VisibilityTimeout=self.visibility_timeout, WaitTimeSeconds=self.wait_time_seconds, ) self._log('{} messages requested, {} messages received'.format( num_messages, len(messages))) processed_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') processed_messages.append(message) except Full: message.change_visibility( VisibilityTimeout=self.return_to_queue_timeout) self._log('Message visibility set to 0, returned to SQS') self._delete_processed_messages(processed_messages) gc.collect() def _delete_processed_messages(self, processed_messages): """Remove processed messages and notify sentry in case of failure. Args: processed_messages (list): List of SQS messages. """ if processed_messages: entries = [{ 'Id': msg.message_id, 'ReceiptHandle': msg.receipt_handle} for msg in processed_messages] delete_messages_response = self.sqs_queue.delete_messages( Entries=entries) failed_to_delete = delete_messages_response.get('Failed', []) if failed_to_delete: error_message = ( 'Failed to delete sqs messages: {messages}'.format( messages=str(failed_to_delete))) self._log(error_message) capture_message(error_message) else: self._log('{} SQS messages were deleted'.format( len(processed_messages))) class AppProcessingWorker(BaseWorker): """Worker that executes any work in application. AppProcessingWorker gets messages from prefetch queue and then processes them using the external processing module. """ def __init__( self, name, prefetch_queue, sqs_queue, return_to_queue_timeout, logger, worker_id): """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 worker_id (int): Unique worker id. """ super(AppProcessingWorker, self).__init__(name, prefetch_queue, logger) self.sqs_queue = sqs_queue self.return_to_queue_timeout = return_to_queue_timeout self.worker_id = worker_id def _do_work(self): if not self._should_work: return try: message = self.prefetch_queue.get(timeout=PREFETCH_QUEUE_WAIT_TIME) except Empty: return try: processing_result, details = processing.process_transcoding_job( message, self.worker_id) if not processing_result: self._log( 'Failed to complete transcoding job: {details}'.format( details=details)) message.delete() self._log('Message deleted.') finally: self.prefetch_queue.task_done() class QueueWorkerManager(object): """Workers manager. Creates and starts the predefined number of workers. """ def __init__( self, readers_count, executors_count, queue_url, 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_url (str): SQS queue identifier. 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_url = queue_url self.workers = [] 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 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_app_processing(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)) def _start_queue_reader(self, reader_id): """Start the QueueReader worker. Args: reader_id (str, int): id of the worker """ worker_name = 'lce-queue-reader-worker{0}'.format(reader_id) sqs_queue = sqs.get_queue(self.queue_url) 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 def _start_app_processing(self, executor_id): """Start the AppProcessingWorker. Args: executor_id (str, int): id of the worker """ worker_name = 'transcoding-queue-worker{0}'\ .format(executor_id) sqs_queue = sqs.get_queue(self.queue_url) if not sqs_queue: self.logger.error('Worker {0} can not be started, because SQS ' 'queue is not accessible'.format(executor_id)) return worker = AppProcessingWorker( worker_name, self.prefetch_queue, sqs_queue, RETURN_TO_QUEUE_TIMEOUT, self.logger, executor_id ) worker.debug = self.debug worker.start() self.workers.append(worker) self.executors_started += 1 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) self.logger.info('Queue workers stopped')