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
|
### 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 un
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 memc
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 part
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 up
|
### Bootstrap token. |