Differential D2864 Diff 10206 swh/scheduler/celery_backend/pika_listener.py

Changeset View

Standalone View

swh/scheduler/celery_backend/pika_listener.py

This file was added.

				# Copyright (C) 2020 The Software Heritage developers
				ardumontUnsubmitted Done Inline Actions -2020? ardumont: -2020?
				# See the AUTHORS file at the top-level directory of this distribution
				# License: GNU General Public License version 3, or any later version
				# See top-level LICENSE file for more information

				import datetime
				import json
				import logging
				import sys

				import pika

				from swh.core.statsd import statsd
				from swh.scheduler import get_scheduler

				logger = logging.getLogger(__name__)


				def utcnow():
				return datetime.datetime.now(tz=datetime.timezone.utc)


				def get_listener(broker_url, queue_name, scheduler_backend):
				connection = pika.BlockingConnection(
				pika.URLParameters(broker_url)
				)
				channel = connection.channel()

				channel.queue_declare(queue=queue_name, durable=True)

				exchange = 'celeryev'
				routing_key = '#'
				channel.queue_bind(queue=queue_name, exchange=exchange,
				routing_key=routing_key)

				channel.basic_qos(prefetch_count=1000)

				channel.basic_consume(
				queue=queue_name,
				on_message_callback=get_on_message(scheduler_backend),
				)

				return channel


				def get_on_message(scheduler_backend):
				def on_message(channel, method_frame, properties, body):
				try:
				events = json.loads(body)
				except Exception:
				logger.warning('Could not parse body %r', body)
				events = []

				if not isinstance(events, list):
				events = [events]

				for event in events:
				logger.debug('Received event %r', event)
				process_event(event, scheduler_backend)

				channel.basic_ack(delivery_tag=method_frame.delivery_tag)

				return on_message


				def process_event(event, scheduler_backend):
				uuid = event.get('uuid')
				if not uuid:
				return

				event_type = event['type']
				statsd.increment('swh_scheduler_listener_handled_event_total',
				tags={'event_type': event_type})

				if event_type == 'task-started':
				scheduler_backend.start_task_run(
				uuid, timestamp=utcnow(),
				metadata={'worker': event.get('hostname')},
				)
				elif event_type == 'task-result':
				result = event['result']

				status = None

				if isinstance(result, dict) and 'status' in result:
				status = result['status']
				if status == 'success':
				status = 'eventful' if result.get('eventful') else 'uneventful'

				if status is None:
				status = 'eventful' if result else 'uneventful'

				scheduler_backend.end_task_run(uuid, timestamp=utcnow(),
				status=status, result=result)
				elif event_type == 'task-failed':
				scheduler_backend.end_task_run(uuid, timestamp=utcnow(),
				status='failed')


				if __name__ == '__main__':
				url = sys.argv[1]
				logging.basicConfig(level=logging.DEBUG)
				scheduler_backend = get_scheduler('local', args={
				'db': 'service=swh-scheduler'
				})
				channel = get_listener(url, 'celeryev.test', scheduler_backend)
				logger.info('Start consuming')
				channel.start_consuming()