def test_set_has_capacity(self):
d = {u'pool': [u'default'], u'os': [u'Ubuntu-16.04']}
# By default, nothing has capacity. None means no data.
now = utils.utcnow()
self.mock_now(now, 0)
self.assertEqual(None, task_queues.probably_has_capacity(d))
# Keep the value for 2 seconds, exclusive.
task_queues.set_has_capacity(d, 2)
self.assertEqual(True, task_queues.probably_has_capacity(d))
self.mock_now(now, 1)
self.assertEqual(True, task_queues.probably_has_capacity(d))
# The value expired.
self.mock_now(now, 2)
self.assertEqual(None, task_queues.probably_has_capacity(d))
def has_capacity(dimensions):
"""Returns True if there's a reasonable chance for this task request
dimensions set to be serviced by a bot alive.
First look at the task queues, then look into the datastore to figure this
out.
"""
assert not ndb.in_transaction()
# Look at the fast path.
cap = task_queues.probably_has_capacity(dimensions)
if cap is not None:
return cap
# Do a query. That's slower and it's eventually consistent.
q = BotInfo.query()
flat = task_queues.dimensions_to_flat(dimensions)
for f in flat:
q = q.filter(BotInfo.dimensions_flat == f)
if q.count(limit=1):
logging.info('Found capacity via BotInfo: %s', flat)
# Add it to the quick cache to improve performance.
task_queues.set_has_capacity(dimensions)
return True
logging.error('HAS NO CAPACITY: %s', flat)
# TODO(maruel): https://crbug.com/839173
return _FAKE_CAPACITY
def test_probably_has_capacity_get_queues(self):
d = {u'pool': [u'default'], u'os': [u'Ubuntu-16.04']}
# Capacity registers there only once there's a request enqueued and
# get_queues() is called.
_assert_bot()
request = _gen_request(properties=_gen_properties(dimensions=d))
task_queues.assert_task_async(request).get_result()
self.assertEqual(1, self.execute_tasks())
self.assertEqual(None, task_queues.probably_has_capacity(d))
# It get sets only once get_queues() is called.
bot_root_key = bot_management.get_root_key(u'bot1')
task_queues.get_queues(bot_root_key)
self.assertEqual(True, task_queues.probably_has_capacity(d))
self.assertEqual([1843498234],
memcache.get('bot1', namespace='task_queues'))
def test_probably_has_capacity(self):
d = {u'pool': [u'default'], u'os': [u'Ubuntu-16.04']}
# A bot coming online doesn't register capacity automatically.
_assert_bot()
self.assertEqual([], memcache.get('bot1', namespace='task_queues'))
self.assertEqual(1, bot_management.BotInfo.query().count())
self.assertEqual(None, task_queues.probably_has_capacity(d))
def has_capacity(dimensions):
"""Returns True if there's a reasonable chance for this task request
dimensions set to be serviced by a bot alive.
First look at the task queues, then look into the datastore to figure this
out.
"""
assert not ndb.in_transaction()
# Look at the fast path.
cap = task_queues.probably_has_capacity(dimensions)
if cap is not None:
return cap
# Add it to the 'quick cache' to improve performance. This cache is kept for
# the same duration as how long bots are considered still alive without a
# ping. Useful if there's a single bot in the fleet for these dimensions and
# it takes a long time to reboot. This is the case with Android with slow
# initialization and some baremetal bots (thanks SCSI firmware!).
seconds = config.settings().bot_death_timeout_secs
@ndb.tasklet
def run_query(flat):
# Do a query. That's slower and it's eventually consistent.
q = BotInfo.query()
for f in flat:
q = q.filter(BotInfo.dimensions_flat == f)
num = yield q.count_async(limit=1)
if num:
logging.info('Found capacity via BotInfo: %s', flat)
raise ndb.Return(True)
# Search a bit harder. In this case, we're looking for BotEvent which would
# be a bot that used to exist recently.
cutoff = utils.utcnow() - datetime.timedelta(seconds=seconds)
q = BotEvent.query(BotEvent.ts > cutoff)
for f in flat:
q = q.filter(BotEvent.dimensions_flat == f)
num = yield q.count_async(limit=1)
if num:
logging.info('Found capacity via BotEvent: %s', flat)
raise ndb.Return(True)
raise ndb.Return(False)
futures = [
run_query(f)
for f in task_queues.expand_dimensions_to_flats(dimensions)
]
ndb.tasklets.Future.wait_all(futures)
if any(f.get_result() for f in futures):
task_queues.set_has_capacity(dimensions, seconds)
return True
logging.warning('HAS NO CAPACITY: %s', dimensions)
return False
def test_probably_has_capacity_empty(self):
# The bot can service this dimensions.
d = {u'pool': [u'default'], u'os': [u'Ubuntu-16.04']}
# By default, nothing has capacity.
self.assertEqual(None, task_queues.probably_has_capacity(d))
self.assertEqual(None, memcache.get('bot1', namespace='task_queues'))
def test_set_has_capacity(self):
d = {u'pool': [u'default'], u'os': [u'Ubuntu-16.04']}
# By default, nothing has capacity. None means no data.
self.assertEqual(None, task_queues.probably_has_capacity(d))
task_queues.set_has_capacity(d)
self.assertEqual(True, task_queues.probably_has_capacity(d))
