def babysit():
"""
Babysit the workers, updating our tables with information about their queues.
"""
# Inspect the available workers to build our state variables
active_queues = controller.inspect().active_queues()
# Now we need the entire list of AvailableQueues from the database, though we only need their
# _id and missing_since attributes. This is preferrable to using a Map/Reduce operation to get
# Mongo to tell us which workers Celery knows about that aren't found in the database.
all_queues_criteria = Criteria(filters={}, fields=('_id', 'missing_since'))
all_queues = list(resources.filter_available_queues(all_queues_criteria))
all_queues_set = set([q.name for q in all_queues])
active_queues_set = set()
for worker, queues in active_queues.items():
# If this worker is a reserved task worker, let's make sure we know about it in our
# available_queues collection, and make sure it is processing a queue with its own name
if re.match('^%s' % RESERVED_WORKER_NAME_PREFIX, worker):
# Make sure that this worker is subscribed to a queue of his own name. If not,
# subscribe him to one
if not worker in [queue['name'] for queue in queues]:
controller.add_consumer(queue=worker, destination=(worker,))
active_queues_set.add(worker)
# Determine which queues are in active_queues_set that aren't in all_queues_set. These are new
# workers and we need to add them to the database.
for worker in (active_queues_set - all_queues_set):
resources.get_or_create_available_queue(worker)
# If there are any AvalailableQueues that have their missing_since attribute set and they are
# present now, let's set their missing_since attribute back to None.
missing_since_queues = set([q.name for q in all_queues if q.missing_since is not None])
for queue in (active_queues_set & missing_since_queues):
active_queue = resources.get_or_create_available_queue(queue)
active_queue.missing_since = None
active_queue.save()
# Now we must delete queues for workers that don't exist anymore, but only if they've been
# missing for at least five minutes.
for queue in (all_queues_set - active_queues_set):
active_queue = list(resources.filter_available_queues(Criteria(filters={'_id': queue})))[0]
# We only want to delete this queue if it has been missing for at least 5 minutes. If this
# AvailableQueue doesn't have a missing_since attribute, that means it has just now gone
# missing. Let's mark its missing_since attribute and continue.
if active_queue.missing_since is None:
active_queue.missing_since = datetime.utcnow()
active_queue.save()
continue
# This queue has been missing for some time. Let's check to see if it's been 5 minutes yet,
# and if it has, let's delete it.
if active_queue.missing_since < datetime.utcnow() - timedelta(minutes=5):
_delete_queue.apply_async(args=(queue,), queue=RESOURCE_MANAGER_QUEUE)
def test_create(self):
"""
Test for the case when the requested queue does not exist.
"""
# Let's add an AvailableQueue just to make sure that it doesn't return any existing queue.
aq_1 = AvailableQueue('queue_1')
aq_1.save()
aq_2 = resources.get_or_create_available_queue('queue_2')
# Assert that the returned instance is correct
self.assertEqual(type(aq_2), AvailableQueue)
self.assertEqual(aq_2.name, 'queue_2')
# A new queue should default to 0 reservations
self.assertEqual(aq_2.num_reservations, 0)
# missing_since defaults to None
self.assertEqual(aq_2.missing_since, None)
# Now we need to assert that it made it to the database as well
aqc = aq_2.get_collection()
aq_bson = aqc.find_one({'_id': 'queue_2'})
# Make sure the default values were set correctly
self.assertEqual(aq_bson['num_reservations'], 0)
self.assertEqual(aq_bson['missing_since'], None)
def test_get(self):
"""
Test for the case when the requested queue does exist.
"""
# Let's add two AvailableQueues just to make sure that it doesn't return the wrong queue.
aq_1 = AvailableQueue('queue_1')
aq_1.save()
missing_since = datetime(2013, 12, 16)
aq_2 = AvailableQueue('queue_2', 7, missing_since)
aq_2.save()
aq_2 = resources.get_or_create_available_queue('queue_2')
# Assert that the returned instance is correct
self.assertEqual(type(aq_2), AvailableQueue)
self.assertEqual(aq_2.name, 'queue_2')
# Make sure the instance attributes are correct
self.assertEqual(aq_2.num_reservations, 7)
self.assertEqual(aq_2.missing_since, missing_since)
# Now we need to assert that the DB is still correct
aqc = aq_2.get_collection()
aq_bson = aqc.find_one({'_id': 'queue_2'})
self.assertEqual(aq_bson['num_reservations'], 7)
self.assertEqual(aq_bson['missing_since'], missing_since)
