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
+BOT_CODE_NS = 'bot_code'
 
 ### 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
 
 
+def get_cached_swarming_bot_zip(version):
+  """Returns the bot contents if its been cached, or None if missing."""
+  # see cache_swarming_bot_zip for how the "meta" entry is set
+  meta = bot_memcache_get(version, 'meta').get_result()
+  if meta is None:
+    logging.info('memcache did not include metadata for version %s', version)
+    return None
+  num_parts, true_sig = meta.split(':')
+
+  # Get everything asynchronously. If something's missing, the hash will be
+  # wrong so no need to check that we got something from each call.
+  futures = [bot_memcache_get(version, 'content', p)
+             for p in range(int(num_parts))]
+  content = ''
+  for f in futures:
+    chunk = f.get_result()
+    if chunk is None:
+      logging.error('bot %s was missing some of its contents', version)

[Vadim Sh., 2017/06/27 21:14:29]

nit: 'bot code ver %s was missing ...'

[aludwin, 2017/06/28 13:23:54]

Done.

+      return None
+    content += chunk
+  h = hashlib.sha256()
+  h.update(content)
+  if h.hexdigest() != true_sig:
+    logging.error('bot %s had signature %s instead of expected %s', version,

[Vadim Sh., 2017/06/27 21:14:29]

same here actually

[aludwin, 2017/06/28 13:23:54]

Done.

+                  h.hexdigest(), true_sig)
+    return None
+  return content
+
+
+def cache_swarming_bot_zip(version, content):
+  """Caches the bot code to memcache."""
+  h = hashlib.sha256()
+  h.update(content)
+  p = 0
+  futures = []
+  while len(content) > 0:
+    chunk_size = min(MAX_MEMCACHED_SIZE_BYTES, len(content))
+    futures.append(bot_memcache_set(content[0:chunk_size],
+                                          version, 'content', p))
+    content = content[chunk_size:]
+    p += 1
+  meta = "%s:%s" % (p, h.hexdigest())
+  for f in futures:
+    f.check_success()
+  bot_memcache_set(meta, version, 'meta').check_success()
+  logging.info('bot %s with sig %s saved in memcached in %d chunks',
+               version, h.hexdigest(), p)
+
+
+def bot_memcache_get(version, desc, part=None):
+  """Mockable async memcache getter."""
+  return ndb.get_context().memcache_get(bot_key(version, desc, part),
+                                        namespace=BOT_CODE_NS)
+
+
+def bot_memcache_set(value, version, desc, part=None):
+  """Mockable async memcache setter."""
+  return ndb.get_context().memcache_set(bot_key(version, desc, part),
+                                        value, namespace=BOT_CODE_NS)
+
+
+def bot_key(version, desc, part=None):
+  """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
+
+
 ### 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))