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