Replace custom blob gRPC API with ByteStream

appengine/swarming/server/bot_code.py

Index: appengine/swarming/server/bot_code.py
diff --git a/appengine/swarming/server/bot_code.py b/appengine/swarming/server/bot_code.py
index 81a56e0d7ad39fbc51469174e872f37d4705ad7a..a346c1ae200829cf309a5c8443837d6b6699e6af 100644
--- a/appengine/swarming/server/bot_code.py
+++ b/appengine/swarming/server/bot_code.py
@@ -28,6 +28,7 @@ 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.
 
@@ -222,7 +223,7 @@ def get_swarming_bot_zip(host):
     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
@@ -236,11 +237,71 @@ def get_swarming_bot_zip(host):
   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.