Replace custom blob gRPC API with ByteStream

appengine/swarming/server/bot_code.py

 # Copyright 2014 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.
 
 """Swarming bot code. Includes bootstrap and swarming_bot.zip.
 
 It includes everything that is AppEngine specific. The non-GAE code is in
 bot_archive.py.
 """
 
@@ skipping 10 matching lines @@
 from components import auth
 from components import config
 from components import datastore_utils
 from components import utils
 from server import bot_archive
 from server import config as local_config
 
 
 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
+MAX_MEMCACHED_SIZE_BYTES = 1000000

[Vadim Sh., 2017/06/23 18:26:50]

how many parts do we get?

[aludwin, 2017/06/26 17:12:56]

Two currently - though in the unit test, we get 11 since I reduce this by a
factor of 10.

 
 ### Models.
 
 
 File = collections.namedtuple('File', ('content', 'who', 'when', 'version'))
 
 
 class VersionedFile(ndb.Model):
   """Versionned entity.
 
@@ skipping 174 matching lines @@
   return version, additionals
 
 
 def get_swarming_bot_zip(host):
   """Returns a zipped file of all the files a bot needs to run.
 
   Returns:
     A string representing the zipped file's contents.
   """
   version, additionals = get_bot_version(host)
-  content = memcache.get('code-' + version, namespace='bot_code')
+  content = get_cached_swarming_bot_zip(version)
   if content:
     logging.debug('memcached bot code %s; %d bytes', version, len(content))
     return content
 
   # Get the start bot script from the database, if present. Pass an empty
   # file if the files isn't present.
   additionals = additionals or {
     'config/bot_config.py': get_bot_config().content,
   }
   bot_dir = os.path.join(ROOT_DIR, 'swarming_bot')
   content, version = bot_archive.get_swarming_bot_zip(
       bot_dir, host, utils.get_app_version(), additionals,
       local_config.settings().enable_ts_monitoring)
-  # This is immutable so not no need to set expiration time.
-  memcache.set('code-' + version, content, namespace='bot_code')
   logging.info('generated bot code %s; %d bytes', version, len(content))
+  cache_swarming_bot_zip(version, content)
   return content
 

[Vadim Sh., 2017/06/23 18:26:50]

nit: \n (two blank lines between top-level code units)

[aludwin, 2017/06/26 17:12:56]

Done.

+class MemcacheMissingException(Exception):
+  pass
+

[Vadim Sh., 2017/06/23 18:26:50]

nit: \n

[aludwin, 2017/06/26 17:12:56]

Done.

+def get_cached_swarming_bot_zip(version):
+  """Returns the bot contents if its been cached, or None if missing."""
+  try:
+    num_parts = get_cached_bot_entry(version, 'parts')

[Vadim Sh., 2017/06/23 18:26:50]

'parts' and 'signature' should be part of one memcache entry, they are always
fetched and stored together, no need to do multiple RPCs

[aludwin, 2017/06/26 17:12:56]

Done.

[aludwin, 2017/06/26 17:12:56]

Done.

+    true_sig = get_cached_bot_entry(version, 'signature')
+    content = ''
+    for p in range(0, num_parts):
+      content += get_cached_bot_entry(version, 'content', p)

[Vadim Sh., 2017/06/23 18:26:50]

please use asynchronous memcache API to fetch parts in parallel

something similar to:
ctx = ndb.get_context()
futures = [ctx.memcache_get(...) for p in range(num_parts)]
content = ''.join(f.get_result() for f in futures)

(or you can convert get_cached_swarming_bot_zip and cache_swarming_bot_zip to
tasklets and use 'yield' inside)

[aludwin, 2017/06/26 17:12:56]

Done.

+    h = hashlib.sha256()
+    h.update(content)
+    if h.hexdigest() != true_sig:
+      logging.error('bot %s had signature %s instead of expected %s', version,
+                    h.hexdigest(), true_sig)
+      return None
+    return content
+  except MemcacheMissingException:
+    return None
+
+
+def cache_swarming_bot_zip(version, content):
+  """Caches the bot code to memcache."""
+  h = hashlib.sha256()
+  h.update(content)
+  p = 0
+  while len(content) > 0:

[Vadim Sh., 2017/06/23 18:26:50]

same here, it should be done in parallel

[aludwin, 2017/06/26 17:12:56]

Done.

+    chunk_size = min(MAX_MEMCACHED_SIZE_BYTES, len(content))
+    set_cached_bot_entry(content[0:chunk_size], version, 'content', p)
+    content=content[chunk_size:]
+    p += 1
+  set_cached_bot_entry(h.hexdigest(), version, 'signature')
+  set_cached_bot_entry(p, version, 'parts')
+  logging.info('bot %s with sig %s saved in memcached in %d chunks',
+               version, h.hexdigest(), p)
+
+
+def get_cached_bot_entry(version, desc, part=None):
+  """Gets a bot entry from memcached.
+
+  Raise MemcacheMissingException if the entry is None.
+  """
+  res = memcache.get(cached_bot_key(version, desc, part), namespace='bot_code')
+  if res is None:
+    raise MemcacheMissingException()
+  return res
+
+
+def set_cached_bot_entry(content, version, desc, part=None):
+  """Sets a bot entry in memcached."""
+  memcache.set(cached_bot_key(version, desc, part), content,
+               namespace='bot_code')
+
+
+def cached_bot_key(version, desc, part):
+  """Returns a memcache key for bot entries."""
+  key = 'code-%s-%s' % (version, desc)
+  if part is not None:
+    key = '%s-%d' % (key, part)
+  return key
 

[Vadim Sh., 2017/06/23 18:26:50]

nit: \n

[aludwin, 2017/06/26 17:12:56]

Missed this one but checked in locally for next upload.

 ### Bootstrap token.
 
 
 class BootstrapToken(auth.TokenKind):
   expiration_sec = 3600
   secret_key = auth.SecretKey('bot_bootstrap_token')
   version = 1
 
 
 def generate_bootstrap_token():
@@ skipping 51 matching lines @@
 
 ## Config validators
 
 
 @config.validation.self_rule('regex:scripts/.+\\.py')
 def _validate_scripts(content, ctx):
   try:
     ast.parse(content)
   except (SyntaxError, TypeError) as e:
     ctx.error('invalid %s: %s' % (ctx.path, e))