add support for awx-manage run_callback_receiver --status

awx/main/dispatch/control.py

@@ -15,18 +15,19 @@ class Control(object):
     services = ('dispatcher', 'callback_receiver')
     result = None
 
-    def __init__(self, service, host=None):
+    def __init__(self, service, queuename=None, routing_key=None):
         if service not in self.services:
             raise RuntimeError('{} must be in {}'.format(service, self.services))
         self.service = service
-        self.queuename = host or get_local_queuename()
-        self.queue = Queue(self.queuename, Exchange(self.queuename), routing_key=self.queuename)
+        self.queuename = queuename or get_local_queuename()
+        self.routing_key = routing_key or self.queuename
+        self.queue = Queue(self.queuename, Exchange(self.queuename), routing_key=self.routing_key)
 
     def publish(self, msg, conn, **kwargs):
         producer = Producer(
             exchange=self.queue.exchange,
             channel=conn,
-            routing_key=self.queuename
+            routing_key=self.routing_key
         )
         producer.publish(msg, expiration=5, **kwargs)
 

awx/main/dispatch/pool.py

@@ -280,6 +280,11 @@ class WorkerPool(object):
             logger.exception('could not kill {}'.format(worker.pid))
 
 
+    def cleanup(self):
+        for worker in self.workers:
+            worker.calculate_managed_tasks()
+
+
 class AutoscalePool(WorkerPool):
     '''
     An extended pool implementation that automatically scales workers up and

awx/main/dispatch/worker/base.py

@@ -56,8 +56,18 @@ class AWXConsumer(ConsumerMixin):
 
     @property
     def listening_on(self):
+        def qname(q):
+            if q.routing_key != q.name:
+                return ':'.join([q.name, q.routing_key])
+            return q.name
+
+        def qtype(q):
+            if q.exchange.type != 'direct':
+                return ' [{}]'.format(q.exchange.type)
+            return ''
+
         return 'listening on {}'.format([
-            '{} [{}]'.format(q.name, q.exchange.type) for q in self.queues
+            '{}{}'.format(qname(q), qtype(q)) for q in self.queues
         ])
 
     def control(self, body, message):

awx/main/management/commands/run_callback_receiver.py

@@ -5,6 +5,8 @@ from django.conf import settings
 from django.core.management.base import BaseCommand
 from kombu import Exchange, Queue
 
+from awx.main.dispatch import get_local_queuename
+from awx.main.dispatch.control import Control
 from awx.main.dispatch.kombu import Connection
 from awx.main.dispatch.worker import AWXConsumer, CallbackBrokerWorker
 
@@ -17,7 +19,20 @@ class Command(BaseCommand):
     '''
     help = 'Launch the job callback receiver'
 
+    def add_arguments(self, parser):
+        parser.add_argument('--status', dest='status', action='store_true',
+                            help='print the internal state of any running callback receiver')
+
     def handle(self, *arg, **options):
+        control_routing_key = 'callback_receiver-{}-control'.format(get_local_queuename())
+        if options.get('status'):
+            print(Control(
+                'callback_receiver',
+                queuename=settings.CALLBACK_QUEUE,
+                routing_key=control_routing_key
+            ).status())
+            return
+
         with Connection(settings.BROKER_URL) as conn:
             consumer = None
             try:
@@ -29,8 +44,9 @@ class Command(BaseCommand):
                         Queue(
                             settings.CALLBACK_QUEUE,
                             Exchange(settings.CALLBACK_QUEUE, type='direct'),
-                            routing_key=settings.CALLBACK_QUEUE
+                            routing_key=key
                         )
+                        for key in [settings.CALLBACK_QUEUE, control_routing_key]
                     ]
                 )
                 consumer.run()

awx/main/models/unified_jobs.py

@@ -1319,7 +1319,7 @@ class UnifiedJob(PolymorphicModel, PasswordFieldsModel, CommonModelNameNotUnique
             timeout = 5
             try:
                 running = self.celery_task_id in ControlDispatcher(
-                    'dispatcher', self.execution_node
+                    'dispatcher', queuename=self.execution_node
                 ).running(timeout=timeout)
             except socket.timeout:
                 logger.error('could not reach dispatcher on {} within {}s'.format(