src/tools/telemetry/telemetry/core/local_server.py

   1 # Copyright 2014 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 import collections
   6 import json
   7 import os
   8 import re
   9 import subprocess
  10 import sys
  11
  12 from telemetry.core import forwarders
  13
  14 NamedPort = collections.namedtuple('NamedPort', ['name', 'port'])
  15
  16
  17 class LocalServerBackend(object):
  18   def __init__(self):
  19     pass
  20
  21   def StartAndGetNamedPorts(self, args):
  22     """Starts the actual server and obtains any sockets on which it
  23     should listen.
  24
  25     Returns a list of NamedPort on which this backend is listening.
  26     """
  27     raise NotImplementedError()
  28
  29   def ServeForever(self):
  30     raise NotImplementedError()
  31
  32
  33 class LocalServer(object):
  34   def __init__(self, server_backend_class):
  35     assert LocalServerBackend in server_backend_class.__bases__
  36     server_module_name = server_backend_class.__module__
  37     assert server_module_name in sys.modules, \
  38             'The server class\' module must be findable via sys.modules'
  39     assert getattr(sys.modules[server_module_name],
  40                    server_backend_class.__name__), \
  41       'The server class must getattrable from its __module__ by its __name__'
  42
  43     self._server_backend_class = server_backend_class
  44     self._subprocess = None
  45     self._devnull = None
  46     self._local_server_controller = None
  47     self.forwarder = None
  48     self.host_ip = None
  49
  50   def Start(self, local_server_controller):
  51     assert self._subprocess == None
  52     self._local_server_controller = local_server_controller
  53
  54     self.host_ip = local_server_controller.host_ip
  55
  56     server_args = self.GetBackendStartupArgs()
  57     server_args_as_json = json.dumps(server_args)
  58     server_module_name = self._server_backend_class.__module__
  59
  60     self._devnull = open(os.devnull, 'w')
  61     cmd = [sys.executable, '-m', __name__]
  62     cmd.extend(["run_backend"])
  63     cmd.extend([server_module_name, self._server_backend_class.__name__,
  64                 server_args_as_json])
  65
  66     env = os.environ.copy()
  67     env['PYTHONPATH'] = os.pathsep.join(sys.path)
  68
  69     cwd = os.path.abspath(
  70         os.path.join(os.path.dirname(__file__), '..', '..'))
  71
  72     self._subprocess = subprocess.Popen(
  73         cmd, cwd=cwd, env=env, stdout=subprocess.PIPE)
  74
  75     named_ports = self._GetNamedPortsFromBackend()
  76     named_port_pair_map = {'http': None, 'https': None, 'dns': None}
  77     for name, port in named_ports:
  78       assert name in named_port_pair_map, '%s forwarding is unsupported' % name
  79       named_port_pair_map[name] = (
  80           forwarders.PortPair(port,
  81                               local_server_controller.GetRemotePort(port)))
  82     self.forwarder = local_server_controller.CreateForwarder(
  83         forwarders.PortPairs(**named_port_pair_map))
  84
  85   def _GetNamedPortsFromBackend(self):
  86     named_ports_json = None
  87     named_ports_re = re.compile('LocalServerBackend started: (?P<port>.+)')
  88     # TODO: This will hang if the subprocess doesn't print the correct output.
  89     while self._subprocess.poll() == None:
  90       m = named_ports_re.match(self._subprocess.stdout.readline())
  91       if m:
  92         named_ports_json = m.group('port')
  93         break
  94
  95     if not named_ports_json:
  96       raise Exception('Server process died prematurely ' +
  97                       'without giving us port pairs.')
  98     return [NamedPort(**pair) for pair in json.loads(named_ports_json.lower())]
  99
 100   @property
 101   def is_running(self):
 102     return self._subprocess != None
 103
 104   def __enter__(self):
 105     return self
 106
 107   def __exit__(self, *args):
 108     self.Close()
 109
 110   def __del__(self):
 111     self.Close()
 112
 113   def Close(self):
 114     if self.forwarder:
 115       self.forwarder.Close()
 116       self.forwarder = None
 117     if self._subprocess:
 118       # TODO(tonyg): Should this block until it goes away?
 119       self._subprocess.kill()
 120       self._subprocess = None
 121     if self._devnull:
 122       self._devnull.close()
 123       self._devnull = None
 124     if self._local_server_controller:
 125       self._local_server_controller.ServerDidClose(self)
 126       self._local_server_controller = None
 127
 128   def GetBackendStartupArgs(self):
 129     """Returns whatever arguments are required to start up the backend"""
 130     raise NotImplementedError()
 131
 132
 133 class LocalServerController():
 134   """Manages the list of running servers
 135
 136   This class manages the running servers, but also provides an isolation layer
 137   to prevent LocalServer subclasses from accessing the browser backend directly.
 138
 139   """
 140   def __init__(self, browser_backend):
 141     self._browser_backend = browser_backend
 142     self._local_servers_by_class = {}
 143     self.host_ip = self._browser_backend.forwarder_factory.host_ip
 144
 145   def StartServer(self, server):
 146     assert not server.is_running, 'Server already started'
 147     assert isinstance(server, LocalServer)
 148     if server.__class__ in self._local_servers_by_class:
 149       raise Exception(
 150           'Canont have two servers of the same class running at once. ' +
 151           'Locate the existing one and use it, or call Close() on it.')
 152
 153     server.Start(self)
 154     self._local_servers_by_class[server.__class__] = server
 155
 156   def GetRunningServer(self, server_class, default_value):
 157     return self._local_servers_by_class.get(server_class, default_value)
 158
 159   @property
 160   def local_servers(self):
 161     return self._local_servers_by_class.values()
 162
 163   def Close(self):
 164     while len(self._local_servers_by_class):
 165       server = self._local_servers_by_class.itervalues().next()
 166       try:
 167         server.Close()
 168       except Exception:
 169         import traceback
 170         traceback.print_exc()
 171
 172   def CreateForwarder(self, port_pairs):
 173     return self._browser_backend.forwarder_factory.Create(port_pairs)
 174
 175   def GetRemotePort(self, port):
 176     return self._browser_backend.GetRemotePort(port)
 177
 178   def ServerDidClose(self, server):
 179     del self._local_servers_by_class[server.__class__]
 180
 181
 182 def _LocalServerBackendMain(args):
 183   assert len(args) == 4
 184   (cmd, server_module_name,
 185    server_backend_class_name, server_args_as_json) = args[:4]
 186   assert cmd == 'run_backend'
 187   server_module = __import__(server_module_name, fromlist=[True])
 188   server_backend_class = getattr(server_module, server_backend_class_name)
 189   server = server_backend_class()
 190
 191   server_args = json.loads(server_args_as_json)
 192
 193   named_ports = server.StartAndGetNamedPorts(server_args)
 194   assert isinstance(named_ports, list)
 195   for named_port in named_ports:
 196     assert isinstance(named_port, NamedPort)
 197
 198   # Note: This message is scraped by the parent process'
 199   # _GetNamedPortsFromBackend(). Do **not** change it.
 200   print 'LocalServerBackend started: %s' % json.dumps(
 201       [pair._asdict() for pair in named_ports]) # pylint: disable=W0212
 202   sys.stdout.flush()
 203
 204   return server.ServeForever()
 205
 206
 207 if __name__ == '__main__':
 208   # This trick is needed because local_server.NamedPort is not the
 209   # same as sys.modules['__main__'].NamedPort. The module itself is loaded
 210   # twice, basically.
 211   from telemetry.core import local_server # pylint: disable=W0406
 212   sys.exit(local_server._LocalServerBackendMain( # pylint: disable=W0212
 213       sys.argv[1:]))