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