testing/legion/task_registration_server.py

   1 # Copyright 2015 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """The registration server used to register tasks.
   6
   7 The registration server is started by the test controller and allows the tasks
   8 to register themselves when they start. Authentication of the tasks controllers
   9 is based on an OTP passed to the run_task binary on startup.
  10 """
  11
  12 import logging
  13 import threading
  14 import xmlrpclib
  15 import SimpleXMLRPCServer
  16
  17 #pylint: disable=relative-import
  18 import common_lib
  19
  20
  21 class TaskRegistrationServer(object):
  22   """Discovery server run on the host."""
  23
  24   def __init__(self):
  25     self._expected_tasks = {}
  26     self._rpc_server = None
  27     self._thread = None
  28
  29   def _RegisterTaskRPC(self, otp, ip):
  30     """The RPC used by a task to register with the registration server."""
  31     assert otp in self._expected_tasks
  32     cb = self._expected_tasks.pop(otp)
  33     cb(ip)
  34
  35   def RegisterTaskCallback(self, otp, callback):
  36     """Registers a callback associated with an OTP."""
  37     assert callable(callback)
  38     self._expected_tasks[otp] = callback
  39
  40   def Start(self):
  41     """Starts the registration server."""
  42     logging.info('Starting task registration server')
  43     self._rpc_server = SimpleXMLRPCServer.SimpleXMLRPCServer(
  44         (common_lib.SERVER_ADDRESS, common_lib.SERVER_PORT),
  45         allow_none=True, logRequests=False)
  46     self._rpc_server.register_function(
  47         self._RegisterTaskRPC, 'RegisterTask')
  48     self._thread = threading.Thread(target=self._rpc_server.serve_forever)
  49     self._thread.start()
  50
  51   def Shutdown(self):
  52     """Shuts the discovery server down."""
  53     if self._thread and self._thread.is_alive():
  54       logging.info('Shutting down task registration server')
  55       self._rpc_server.shutdown()