Add support for repeated keys in TaskRequest.

appengine/swarming/server/task_queues.py

 # Copyright 2017 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Ambient task queues generated from the actual load.
 
 This means that the task queues are deduced by the actual load, they are never
 explicitly defined. They are eventually deleted by a cron job once no incoming
 task with the exact set of dimensions is triggered anymore.
 
@@ skipping 220 matching lines @@
       s.valid_until_ts = valid_until_ts
       self.sets = [s for s in self.sets if s.valid_until_ts >= now]
       return True
     # It was updated already, skip storing again.
     old = len(self.sets)
     self.sets = [s for s in self.sets if s.valid_until_ts >= now]
     return len(self.sets) != old
 
   def match_request(self, dimensions):
     """Confirms that this instance actually stores this set."""
-    return self._match_request_flat(
-        u'%s:%s' % (k, v) for k, v in dimensions.iteritems())
+    assert isinstance(dimensions, list), repr(dimensions)
+    return self._match_request_flat(u'%s:%s' % (k, v) for k, v in dimensions)
 
   def match_bot(self, bot_dimensions):
     """Returns the TaskDimensionsSet that matches this bot_dimensions, if any.
     """
     for s in self.sets:
       if s.match_bot(bot_dimensions):
         return s
 
   def _match_request_flat(self, dimensions_flat):
     d = frozenset(dimensions_flat)
@@ skipping 145 matching lines @@
         id=1, parent=bot_root_key,
         dimensions_flat=_flatten_bot_dimensions(bot_dimensions)).put()
     memcache.set(bot_id, sorted(matches), namespace='task_queues')
   finally:
     logging.debug(
         '_rebuild_bot_cache(%s) in %.3fs. Registered for %d queues; cleaned %d',
         bot_id, (utils.utcnow()-now).total_seconds(), len(matches), cleaned)
 
 
 def _get_task_dims_key(dimensions_hash, dimensions):
-  if u'id' in dimensions:
-    return ndb.Key(
-        TaskDimensionsRoot, u'id:%s' % dimensions['id'],
-        TaskDimensions, dimensions_hash)
-  return ndb.Key(
-      TaskDimensionsRoot, u'pool:%s' % dimensions['pool'],
-      TaskDimensions, dimensions_hash)
+  """Return the ndb.Key() for the request dimensions.
+
+  The dimensions may have at most one 'id' key but may have multiple 'pool'. In
+  the case of 'pool', return the first value in sorted order. This means that
+  for a request with both 'pool:bar' and 'pool:foo', the request is rooted at
+  'pool:bar'.
+
+  Arguments:
+  - dimensions: sorted list(tuple(unicode, unicode))
+  """
+  assert isinstance(dimensions, list), repr(dimensions)
+  # Search for ID.
+  for key, value in dimensions:
+    if key == u'id':
+      return ndb.Key(
+          TaskDimensionsRoot, u'id:%s' % value, TaskDimensions, dimensions_hash)
+    # 'pool' comes after 'id'.
+    if key == u'pool':
+      return ndb.Key(
+          TaskDimensionsRoot, u'pool:%s' % value,
+          TaskDimensions, dimensions_hash)
+  raise datastore_errors.BadValueError(u'invalid request')
 
 
 def _hash_data(data):
   """Returns a 32 bits non-zero hash."""
   assert isinstance(data, str), repr(data)
   digest = hashlib.md5(data).digest()
   # Note that 'L' means C++ unsigned long which is (usually) 32 bits and
   # python's int is 64 bits.
   return int(struct.unpack('<L', digest[:4])[0]) or 1
 
@@ skipping 75 matching lines @@
   raise ndb.Return(need_db_store)
 
 
 ### Public APIs.
 
 
 def hash_dimensions(dimensions):
   """Returns a 32 bits int that is a hash of the request dimensions specified.
 
   Arguments:
-    dimensions: dict(str, str)
+    dimensions: sorted list(tuple(unicode, unicode))
 
   The return value is guaranteed to be non-zero so it can be used as a key id in
   a ndb.Key.
   """
+  assert isinstance(dimensions, list), repr(dimensions)
   # This horrible code is the product of micro benchmarks.
   data = ''
-  for k, v in sorted(dimensions.items()):
+  for k, v in dimensions:
     data += k.encode('utf8')
     data += '\000'
     data += v.encode('utf8')
     data += '\000'
   return _hash_data(data)
 
 
 def assert_bot(bot_dimensions):
   """Prepares BotTaskDimensions entities as needed.
 
@@ skipping 78 matching lines @@
   data = {
     u'dimensions': request.properties.dimensions,
     u'dimensions_hash': str(dimensions_hash),
     u'valid_until_ts': request.expiration_ts + _ADVANCE,
   }
   payload = utils.encode_to_json(data)
 
   # If this task specifies an 'id' value, updates the cache inline since we know
   # there's only one bot that can run it, so it won't take long. This permits
   # tasks like 'terminate' tasks to execute faster.
-  if request.properties.dimensions.get(u'id'):
+  if any(k == u'id' for k, _ in request.properties.dimensions):
     rebuild_task_cache(payload)
     return
 
   # We can't use the request ID since the request was not stored yet, so embed
   # all the necessary information.
   url = '/internal/taskqueue/rebuild-task-cache'
   if not utils.enqueue_task(
       url, queue_name='rebuild-task-cache', payload=payload):
     logging.error('Failed to enqueue TaskDimensions update %x', dimensions_hash)
     # Technically we'd want to raise a endpoints.InternalServerErrorException.
@@ skipping 54 matching lines @@
     - 'dimensions': dict of task dimensions to refresh
     - 'dimensions_hash': precalculated hash for dimensions
     - 'valid_until_ts': expiration_ts + _ADVANCE for how long this cache is
       valid
   """
   data = json.loads(payload)
   logging.debug('rebuild_task_cache(%s)', data)
   dimensions = data[u'dimensions']
   dimensions_hash = int(data[u'dimensions_hash'])
   valid_until_ts = utils.parse_datetime(data[u'valid_until_ts'])
-  dimensions_flat = sorted(u'%s:%s' % (k, v) for k, v in dimensions.iteritems())
+  dimensions_flat = sorted(u'%s:%s' % (k, v) for k, v in dimensions)
 
   now = utils.utcnow()
   updated = 0
   viable = 0
   try:
     pending = set()
     for bot_task_key in _yield_BotTaskDimensions_keys(
         dimensions_hash, dimensions_flat):
       viable += 1
       future = _refresh_BotTaskDimensions(
@@ skipping 56 matching lines @@
         btd_deleted += 1
         bot_id = key.parent().id()
         memcache.delete(bot_id, namespace='task_queues')
         logging.debug('- %d for bot %s', key.id(), bot_id)
   finally:
     logging.info(
         'tidy_stale() in %.3fs; TaskDimensions: found %d, deleted %d; '
         'BotTaskDimensions: found %d, deleted %d',
         (utils.utcnow() - now).total_seconds(),
         td_found, td_deleted, btd_found, btd_deleted)