OLD | NEW |
1 # Copyright 2014 The LUCI Authors. All rights reserved. | 1 # Copyright 2014 The LUCI Authors. All rights reserved. |
2 # Use of this source code is governed under the Apache License, Version 2.0 | 2 # Use of this source code is governed under the Apache License, Version 2.0 |
3 # that can be found in the LICENSE file. | 3 # that can be found in the LICENSE file. |
4 | 4 |
5 """Swarming bot code. Includes bootstrap and swarming_bot.zip. | 5 """Swarming bot code. Includes bootstrap and swarming_bot.zip. |
6 | 6 |
7 It includes everything that is AppEngine specific. The non-GAE code is in | 7 It includes everything that is AppEngine specific. The non-GAE code is in |
8 bot_archive.py. | 8 bot_archive.py. |
9 """ | 9 """ |
10 | 10 |
(...skipping 10 matching lines...) Expand all Loading... |
21 from components import auth | 21 from components import auth |
22 from components import config | 22 from components import config |
23 from components import datastore_utils | 23 from components import datastore_utils |
24 from components import utils | 24 from components import utils |
25 from server import bot_archive | 25 from server import bot_archive |
26 from server import config as local_config | 26 from server import config as local_config |
27 | 27 |
28 | 28 |
29 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | 29 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
30 | 30 |
| 31 MAX_MEMCACHED_SIZE_BYTES = 1000000 |
| 32 BOT_CODE_NS = 'bot_code' |
31 | 33 |
32 ### Models. | 34 ### Models. |
33 | 35 |
34 | 36 |
35 File = collections.namedtuple('File', ('content', 'who', 'when', 'version')) | 37 File = collections.namedtuple('File', ('content', 'who', 'when', 'version')) |
36 | 38 |
37 | 39 |
38 class VersionedFile(ndb.Model): | 40 class VersionedFile(ndb.Model): |
39 """Versionned entity. | 41 """Versionned entity. |
40 | 42 |
(...skipping 174 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
215 return version, additionals | 217 return version, additionals |
216 | 218 |
217 | 219 |
218 def get_swarming_bot_zip(host): | 220 def get_swarming_bot_zip(host): |
219 """Returns a zipped file of all the files a bot needs to run. | 221 """Returns a zipped file of all the files a bot needs to run. |
220 | 222 |
221 Returns: | 223 Returns: |
222 A string representing the zipped file's contents. | 224 A string representing the zipped file's contents. |
223 """ | 225 """ |
224 version, additionals = get_bot_version(host) | 226 version, additionals = get_bot_version(host) |
225 content = memcache.get('code-' + version, namespace='bot_code') | 227 content = get_cached_swarming_bot_zip(version) |
226 if content: | 228 if content: |
227 logging.debug('memcached bot code %s; %d bytes', version, len(content)) | 229 logging.debug('memcached bot code %s; %d bytes', version, len(content)) |
228 return content | 230 return content |
229 | 231 |
230 # Get the start bot script from the database, if present. Pass an empty | 232 # Get the start bot script from the database, if present. Pass an empty |
231 # file if the files isn't present. | 233 # file if the files isn't present. |
232 additionals = additionals or { | 234 additionals = additionals or { |
233 'config/bot_config.py': get_bot_config().content, | 235 'config/bot_config.py': get_bot_config().content, |
234 } | 236 } |
235 bot_dir = os.path.join(ROOT_DIR, 'swarming_bot') | 237 bot_dir = os.path.join(ROOT_DIR, 'swarming_bot') |
236 content, version = bot_archive.get_swarming_bot_zip( | 238 content, version = bot_archive.get_swarming_bot_zip( |
237 bot_dir, host, utils.get_app_version(), additionals, | 239 bot_dir, host, utils.get_app_version(), additionals, |
238 local_config.settings().enable_ts_monitoring) | 240 local_config.settings().enable_ts_monitoring) |
239 # This is immutable so not no need to set expiration time. | |
240 memcache.set('code-' + version, content, namespace='bot_code') | |
241 logging.info('generated bot code %s; %d bytes', version, len(content)) | 241 logging.info('generated bot code %s; %d bytes', version, len(content)) |
| 242 cache_swarming_bot_zip(version, content) |
242 return content | 243 return content |
243 | 244 |
244 | 245 |
| 246 def get_cached_swarming_bot_zip(version): |
| 247 """Returns the bot contents if its been cached, or None if missing.""" |
| 248 # see cache_swarming_bot_zip for how the "meta" entry is set |
| 249 meta = bot_memcache_get(version, 'meta').get_result() |
| 250 if meta is None: |
| 251 logging.info('memcache did not include metadata for version %s', version) |
| 252 return None |
| 253 num_parts, true_sig = meta.split(':') |
| 254 |
| 255 # Get everything asynchronously. If something's missing, the hash will be |
| 256 # wrong so no need to check that we got something from each call. |
| 257 futures = [bot_memcache_get(version, 'content', p) |
| 258 for p in range(int(num_parts))] |
| 259 content = '' |
| 260 for f in futures: |
| 261 chunk = f.get_result() |
| 262 if chunk is None: |
| 263 logging.error('bot code %s was missing some of its contents', version) |
| 264 return None |
| 265 content += chunk |
| 266 h = hashlib.sha256() |
| 267 h.update(content) |
| 268 if h.hexdigest() != true_sig: |
| 269 logging.error('bot code %s had signature %s instead of expected %s', |
| 270 version, h.hexdigest(), true_sig) |
| 271 return None |
| 272 return content |
| 273 |
| 274 |
| 275 def cache_swarming_bot_zip(version, content): |
| 276 """Caches the bot code to memcache.""" |
| 277 h = hashlib.sha256() |
| 278 h.update(content) |
| 279 p = 0 |
| 280 futures = [] |
| 281 while len(content) > 0: |
| 282 chunk_size = min(MAX_MEMCACHED_SIZE_BYTES, len(content)) |
| 283 futures.append(bot_memcache_set(content[0:chunk_size], |
| 284 version, 'content', p)) |
| 285 content = content[chunk_size:] |
| 286 p += 1 |
| 287 meta = "%s:%s" % (p, h.hexdigest()) |
| 288 for f in futures: |
| 289 f.check_success() |
| 290 bot_memcache_set(meta, version, 'meta').check_success() |
| 291 logging.info('bot %s with sig %s saved in memcached in %d chunks', |
| 292 version, h.hexdigest(), p) |
| 293 |
| 294 |
| 295 def bot_memcache_get(version, desc, part=None): |
| 296 """Mockable async memcache getter.""" |
| 297 return ndb.get_context().memcache_get(bot_key(version, desc, part), |
| 298 namespace=BOT_CODE_NS) |
| 299 |
| 300 |
| 301 def bot_memcache_set(value, version, desc, part=None): |
| 302 """Mockable async memcache setter.""" |
| 303 return ndb.get_context().memcache_set(bot_key(version, desc, part), |
| 304 value, namespace=BOT_CODE_NS) |
| 305 |
| 306 |
| 307 def bot_key(version, desc, part=None): |
| 308 """Returns a memcache key for bot entries.""" |
| 309 key = 'code-%s-%s' % (version, desc) |
| 310 if part is not None: |
| 311 key = '%s-%d' % (key, part) |
| 312 return key |
| 313 |
| 314 |
245 ### Bootstrap token. | 315 ### Bootstrap token. |
246 | 316 |
247 | 317 |
248 class BootstrapToken(auth.TokenKind): | 318 class BootstrapToken(auth.TokenKind): |
249 expiration_sec = 3600 | 319 expiration_sec = 3600 |
250 secret_key = auth.SecretKey('bot_bootstrap_token') | 320 secret_key = auth.SecretKey('bot_bootstrap_token') |
251 version = 1 | 321 version = 1 |
252 | 322 |
253 | 323 |
254 def generate_bootstrap_token(): | 324 def generate_bootstrap_token(): |
(...skipping 51 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
306 | 376 |
307 ## Config validators | 377 ## Config validators |
308 | 378 |
309 | 379 |
310 @config.validation.self_rule('regex:scripts/.+\\.py') | 380 @config.validation.self_rule('regex:scripts/.+\\.py') |
311 def _validate_scripts(content, ctx): | 381 def _validate_scripts(content, ctx): |
312 try: | 382 try: |
313 ast.parse(content) | 383 ast.parse(content) |
314 except (SyntaxError, TypeError) as e: | 384 except (SyntaxError, TypeError) as e: |
315 ctx.error('invalid %s: %s' % (ctx.path, e)) | 385 ctx.error('invalid %s: %s' % (ctx.path, e)) |
OLD | NEW |