swarming: Add server-side implementation for supplemental bot_config

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 419 matching lines @@
 
     data = {
       'bot_version': bot_code.get_bot_version(self.get_bot_contact_server())[0],
       'server_version': utils.get_app_version(),
       'bot_group_cfg_version': res.bot_group_cfg.version,
       'bot_group_cfg': {
         # Let the bot know its server-side dimensions (from bots.cfg file).
         'dimensions': res.bot_group_cfg.dimensions,
       },
     }
+    if res.bot_group_cfg.bot_config_script_content:
+      logging.info(
+          'Injecting %s: %d bytes',
+          res.bot_group_cfg.bot_config_script,
+          len(res.bot_group_cfg.bot_config_script_content))
+      data['bot_config'] = res.bot_group_cfg.bot_config_script_content
     self.send_response(data)
 
 
 class BotPollHandler(_BotBaseHandler):
   """The bot polls for a task; returns either a task, update command or sleep.
 
   In case of exception on the bot, this is enough to get it just far enough to
   eventually self-update to a working version. This is to ensure that coding
   errors in bot code doesn't kill all the fleet at once, they should still be up
   just enough to be able to self-update again even if they don't get task
@@ skipping 48 matching lines @@
       self._cmd_sleep(sleep_streak, quarantined)
       return
 
     # If the server-side per-bot config for the bot has changed, we need
     # to restart this particular bot, so it picks up new config in /handshake.
     # Do this check only for bots that know about server-side per-bot configs
     # already (such bots send 'bot_group_cfg_version' state attribute).
     cur_bot_cfg_ver = res.state.get('bot_group_cfg_version')
     if cur_bot_cfg_ver and cur_bot_cfg_ver != res.bot_group_cfg.version:
       bot_event('request_restart')
-      self._cmd_restart('Restarting to pick up new bots.cfg config')
+      self._cmd_bot_restart('Restarting to pick up new bots.cfg config')
       return
 
     #
     # At that point, the bot should be in relatively good shape since it's
     # running the right version. It is still possible that invalid code was
     # pushed to the server, so be diligent about it.
     #
 
+    # TODO(maruel): Remove this and migrate all use cases in bot_config.py
+    # on_bot_idle().
     # Bot may need a reboot if it is running for too long. We do not reboot
     # quarantined bots.
     needs_restart, restart_message = bot_management.should_restart_bot(
         res.bot_id, res.state)
     if needs_restart:
       bot_event('request_restart')
-      self._cmd_restart(restart_message)
+      self._cmd_host_reboot(restart_message)
       return
 
     # The bot is in good shape. Try to grab a task.
     try:
       # This is a fairly complex function call, exceptions are expected.
       request, secret_bytes, run_result = task_scheduler.bot_reap_task(
           res.dimensions, res.bot_id, res.version,
           res.state.get('lease_expiration_ts'))
       if not request:
         # No task found, tell it to sleep a bit.
@@ skipping 83 matching lines @@
     }
     self.send_response(out)
 
   def _cmd_update(self, expected_version):
     out = {
       'cmd': 'update',
       'version': expected_version,
     }
     self.send_response(out)
 
-  def _cmd_restart(self, message):
-    logging.info('Rebooting bot: %s', message)
+  def _cmd_host_reboot(self, message):
+    logging.info('Rebooting host: %s', message)
     out = {
-      'cmd': 'restart',
+      'cmd': 'host_reboot',
       'message': message,
     }
     self.send_response(out)
+
+  def _cmd_bot_restart(self, message):
+    logging.info('Restarting bot: %s', message)
+    out = {
+      'cmd': 'bot_restart',
+      'message': message,
+    }
+    self.send_response(out)
 
 
 class BotEventHandler(_BotBaseHandler):
   """On signal that a bot had an event worth logging."""
 
   EXPECTED_KEYS = _BotBaseHandler.EXPECTED_KEYS | {u'event', u'message'}
 
   ALLOWED_EVENTS = ('bot_error', 'bot_log', 'bot_rebooting', 'bot_shutdown')
 
   @auth.public  # auth happens in self._process()
@@ skipping 264 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]