""" Queue worker ============ """ import logging from queue import Empty, Full from queue import Queue from threading import Thread import time from ytownership.connectors import sqs from ytownership.connectors import youtube from ytownership.connectors.sentry import sentry_client from ytownership.connectors.sqs import JSONMessageExt from ytownership.errors import FailedToStartError from ytownership.logic import message as message_logic RETURN_TO_QUEUE_TIMEOUT = 0 PREFETCH_QUEUE_WAIT_TIME = 1 PREFETCH_QUEUE_MAX_WAIT_TIME = 30 MAX_RETRIES = 3 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: if sentry_client: sentry_client.captureException() 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: if sentry_client: sentry_client.captureException() def _do_work(self): """Actual piece of work. Must be overriden in derived classes. """ pass def _log(self, message, level=logging.INFO): """Wrapper for log adapter 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.num_messages - self.prefetch_queue.qsize() self._log('receiving {0} SQS messages'.format(num_messages)) messages = self.sqs_queue.receive_messages( MaxNumberOfMessages=num_messages, WaitTimeSeconds=self.wait_time_seconds, VisibilityTimeout=self.visibility_timeout, MessageAttributeNames=[sqs.CORRELATION_ID_ATTRIBUTE], AttributeNames=[sqs.APPROXIMATE_RECEIVE_COUNT_ATTRIBUTE]) for message in messages: # Don't prefetch anything, if we are shutting down if not self._should_work: return try: message = JSONMessageExt(message) 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 ApiExecutorWorker(BaseWorker): """Worker that executes requests to YouTube API and persists the results ApiExecutorWorker gets messages from prefetch queue and them process them using the external processing module (ytownership.logic.message). """ def __init__( self, name, prefetch_queue, sqs_queue, gapi_client, 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 gapi_client: Google API Client instance 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(ApiExecutorWorker, self).__init__(name, prefetch_queue, logger) self.sqs_queue = sqs_queue self.gapi_client = gapi_client self.return_to_queue_timeout = return_to_queue_timeout 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: message.context.youtube_client = self.gapi_client result = message_logic.process_message(message) if result == message_logic.MessageProcessingResult.SUCCESS: message.delete() self._log('Message deleted.') elif result == message_logic.MessageProcessingResult.ERROR: message.change_visibility(self.return_to_queue_timeout) self._log('Message returned to the queue.') else: message.delete() self._log('Message processing failed, 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_name, prefetch_count, visibility_timeout, wait_time_seconds, num_messages, key_file, 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 key_file (str): Google API Account key file 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 = [] 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 self.key_file = key_file 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_api_executor(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 = 'yt-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 def _start_api_executor(self, executor_id): """Start the ApiExecutor worker Args: executor_id (str, int): id of the worker """ worker_name = 'yt-queue-api-executor-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 youtube_client = youtube.ServiceMapping(self.key_file) worker = ApiExecutorWorker( worker_name, self.prefetch_queue, sqs_queue, youtube_client.youtube_partner, RETURN_TO_QUEUE_TIMEOUT, self.logger ) 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')