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
  15 #pylint: disable=relative-import
  16 import common_lib
  17 import SimpleJSONRPCServer
  18
  19
  20 class TaskRegistrationServer(object):
  21   """Discovery server run on the host."""
  22
  23   def __init__(self):
  24     self._expected_tasks = {}
  25     self._rpc_server = None
  26     self._thread = None
  27
  28   def _RegisterTaskRPC(self, otp, ip):
  29     """The RPC used by a task to register with the registration server."""
  30     assert otp in self._expected_tasks
  31     cb = self._expected_tasks.pop(otp)
  32     cb(ip)
  33
  34   def RegisterTaskCallback(self, otp, callback):
  35     """Registers a callback associated with an OTP."""
  36     assert callable(callback)
  37     self._expected_tasks[otp] = callback
  38
  39   def Start(self):
  40     """Starts the registration server."""
  41     logging.info('Starting task registration server')
  42     self._rpc_server = SimpleJSONRPCServer.SimpleJSONRPCServer(
  43         (common_lib.SERVER_ADDRESS, common_lib.SERVER_PORT),
  44         allow_none=True, logRequests=False)
  45     self._rpc_server.register_function(
  46         self._RegisterTaskRPC, 'RegisterTask')
  47     self._thread = threading.Thread(target=self._rpc_server.serve_forever)
  48     self._thread.start()
  49
  50   def Shutdown(self):
  51     """Shuts the discovery server down."""
  52     if self._thread and self._thread.is_alive():
  53       logging.info('Shutting down task registration server')
  54       self._rpc_server.shutdown()