Complete CIPD integration

appengine/swarming/handlers_bot.py

 # Copyright 2015 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.
 
 """Internal bot API handlers."""
 
 import base64
 import json
 import logging
 import re
@@ skipping 486 matching lines @@
     self.send_response({})
 
 
 class BotTaskUpdateHandler(_BotApiHandler):
   """Receives updates from a Bot for a task.
 
   The handler verifies packets are processed in order and will refuse
   out-of-order packets.
   """
   ACCEPTED_KEYS = {
-    u'bot_overhead', u'cost_usd', u'duration', u'exit_code',
+    u'bot_overhead', u'cipd_stats', u'cost_usd', u'duration', u'exit_code',
     u'hard_timeout', u'id', u'io_timeout', u'isolated_stats', u'output',
     u'output_chunk_start', u'outputs_ref', u'task_id',
   }
   REQUIRED_KEYS = {u'id', u'task_id'}
 
   @auth.require(acl.is_bot)
   def post(self, task_id=None):
     # Unlike handshake and poll, we do not accept invalid keys here. This code
     # path is much more strict.
     request = self.parse_body()
     msg = log_unexpected_subset_keys(
         self.ACCEPTED_KEYS, self.REQUIRED_KEYS, request, self.request, 'bot',
         'keys')
     if msg:
       self.abort_with_error(400, error=msg)
 
     bot_id = request['id']
     cost_usd = request['cost_usd']
     task_id = request['task_id']
 
     # Make sure bot self-reported ID matches the authentication token.
     bot_auth.validate_bot_id(bot_id)
 
     bot_overhead = request.get('bot_overhead')
     duration = request.get('duration')
     exit_code = request.get('exit_code')
     hard_timeout = request.get('hard_timeout')
     io_timeout = request.get('io_timeout')
     isolated_stats = request.get('isolated_stats')
+    cipd_stats = request.get('cipd_stats')
     output = request.get('output')
     output_chunk_start = request.get('output_chunk_start')
     outputs_ref = request.get('outputs_ref')
 
-    if isolated_stats and bot_overhead is None:
+    if (isolated_stats or cipd_stats) and bot_overhead is None:
       ereporter2.log_request(
           request=self.request,
           source='server',
           category='task_failure',
           message='Failed to update task: %s' % task_id)
       self.abort_with_error(
           400,
-          error='isolated_stats requires bot_overhead to be set'
-                '\nbot_overhead: %s\nisolated_stats: %s' %
+          error='isolated_stats and cipd_stats require bot_overhead to be set'
+                '\nbot_overhead: %s\nisolate_stats: %s' %
                 (bot_overhead, isolated_stats))
 
     run_result_key = task_pack.unpack_run_result_key(task_id)
     performance_stats = None
     if bot_overhead:
       performance_stats = task_result.PerformanceStats(
           bot_overhead=bot_overhead)
       if isolated_stats:
         download = isolated_stats.get('download') or {}
         upload = isolated_stats.get('upload') or {}
         def unpack_base64(d, k):
           x = d.get(k)
           if x:
             return base64.b64decode(x)
         performance_stats.isolated_download = task_result.OperationStats(
             duration=download.get('duration'),
             initial_number_items=download.get('initial_number_items'),
             initial_size=download.get('initial_size'),
             items_cold=unpack_base64(download, 'items_cold'),
             items_hot=unpack_base64(download, 'items_hot'))
         performance_stats.isolated_upload = task_result.OperationStats(
             duration=upload.get('duration'),
             items_cold=unpack_base64(upload, 'items_cold'),
             items_hot=unpack_base64(upload, 'items_hot'))
+      if cipd_stats:
+        performance_stats.package_installation = task_result.OperationStats(
+            duration=cipd_stats.get('duration'))
 
     if output is not None:
       try:
         output = base64.b64decode(output)
       except UnicodeEncodeError as e:
         logging.error('Failed to decode output\n%s\n%r', e, output)
         output = output.encode('ascii', 'replace')
       except TypeError as e:
         # Save the output as-is instead. The error will be logged in ereporter2
         # and returning a HTTP 500 would only force the bot to stay in a retry
@@ skipping 120 matching lines @@
       ('/swarming/api/v1/bot/poll', BotPollHandler),
       ('/swarming/api/v1/bot/server_ping', ServerPingHandler),
       ('/swarming/api/v1/bot/task_update', BotTaskUpdateHandler),
       ('/swarming/api/v1/bot/task_update/<task_id:[a-f0-9]+>',
           BotTaskUpdateHandler),
       ('/swarming/api/v1/bot/task_error', BotTaskErrorHandler),
       ('/swarming/api/v1/bot/task_error/<task_id:[a-f0-9]+>',
           BotTaskErrorHandler),
   ]
   return [webapp2.Route(*i) for i in routes]