• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env vpython3
2# Copyright 2016 The Chromium Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Wrapper for adding logdog streaming support to swarming tasks."""
7
8import argparse
9import contextlib
10import json
11import logging
12import os
13import signal
14import subprocess
15import sys
16
17_SRC_PATH = os.path.abspath(os.path.join(
18    os.path.dirname(__file__), '..', '..', '..'))
19sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'catapult', 'devil'))
20sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'catapult', 'common',
21                             'py_utils'))
22
23from devil.utils import signal_handler
24from devil.utils import timeout_retry
25from py_utils import tempfile_ext
26
27OUTPUT = 'logdog'
28COORDINATOR_HOST = 'luci-logdog.appspot.com'
29LOGDOG_TERMINATION_TIMEOUT = 30
30
31
32def CommandParser():
33  # Parses the command line arguments being passed in
34  parser = argparse.ArgumentParser(allow_abbrev=False)
35  wrapped = parser.add_mutually_exclusive_group()
36  wrapped.add_argument(
37      '--target',
38      help='The test target to be run. If neither target nor script are set,'
39      ' any extra args passed to this script are assumed to be the'
40      ' full test command to run.')
41  wrapped.add_argument(
42      '--script',
43      help='The script target to be run. If neither target nor script are set,'
44      ' any extra args passed to this script are assumed to be the'
45      ' full test command to run.')
46  parser.add_argument('--logdog-bin-cmd', required=True,
47                      help='The logdog bin cmd.')
48  return parser
49
50
51def CreateStopTestsMethod(proc):
52  def StopTests(signum, _frame):
53    logging.error('Forwarding signal %s to test process', str(signum))
54    proc.send_signal(signum)
55  return StopTests
56
57
58@contextlib.contextmanager
59def NoLeakingProcesses(popen):
60  try:
61    yield popen
62  finally:
63    if popen is not None:
64      try:
65        if popen.poll() is None:
66          popen.kill()
67      except OSError:
68        logging.warning('Failed to kill %s. Process may be leaked.',
69                        str(popen.pid))
70
71
72def GetProjectFromLuciContext():
73  """Return the "project" from LUCI_CONTEXT.
74
75  LUCI_CONTEXT contains a section "realm.name" whose value follows the format
76  "<project>:<realm>". This method parses and return the "project" part.
77
78  Fallback to "chromium" if realm name is None
79  """
80  project = 'chromium'
81  ctx_path = os.environ.get('LUCI_CONTEXT')
82  if ctx_path:
83    try:
84      with open(ctx_path) as f:
85        luci_ctx = json.load(f)
86        realm_name = luci_ctx.get('realm', {}).get('name')
87        if realm_name:
88          project = realm_name.split(':')[0]
89    except (OSError, IOError, ValueError):
90      pass
91  return project
92
93
94def main():
95  parser = CommandParser()
96  args, extra_cmd_args = parser.parse_known_args(sys.argv[1:])
97
98  logging.basicConfig(level=logging.INFO)
99  if args.target:
100    test_cmd = [os.path.join('bin', 'run_%s' % args.target), '-v']
101    test_cmd += extra_cmd_args
102  elif args.script:
103    test_cmd = [args.script]
104    test_cmd += extra_cmd_args
105  else:
106    test_cmd = extra_cmd_args
107
108  test_env = dict(os.environ)
109  logdog_cmd = []
110
111  with tempfile_ext.NamedTemporaryDirectory(
112      prefix='tmp_android_logdog_wrapper') as temp_directory:
113    if not os.path.exists(args.logdog_bin_cmd):
114      logging.error(
115          'Logdog binary %s unavailable. Unable to create logdog client',
116          args.logdog_bin_cmd)
117    else:
118      streamserver_uri = 'unix:%s' % os.path.join(temp_directory,
119                                                  'butler.sock')
120      prefix = os.path.join('android', 'swarming', 'logcats',
121                            os.environ.get('SWARMING_TASK_ID'))
122      project = GetProjectFromLuciContext()
123
124      logdog_cmd = [
125          args.logdog_bin_cmd,
126          '-project', project,
127          '-output', OUTPUT,
128          '-prefix', prefix,
129          '-coordinator-host', COORDINATOR_HOST,
130          'serve',
131          '-streamserver-uri', streamserver_uri]
132      test_env.update({
133          'LOGDOG_STREAM_PROJECT': project,
134          'LOGDOG_STREAM_PREFIX': prefix,
135          'LOGDOG_STREAM_SERVER_PATH': streamserver_uri,
136          'LOGDOG_COORDINATOR_HOST': COORDINATOR_HOST,
137      })
138
139    logdog_proc = None
140    if logdog_cmd:
141      logdog_proc = subprocess.Popen(logdog_cmd)
142
143    with NoLeakingProcesses(logdog_proc):
144      with NoLeakingProcesses(
145          subprocess.Popen(test_cmd, env=test_env)) as test_proc:
146        with signal_handler.SignalHandler(signal.SIGTERM,
147                                          CreateStopTestsMethod(test_proc)):
148          result = test_proc.wait()
149          if logdog_proc:
150            def logdog_stopped():
151              return logdog_proc.poll() is not None
152
153            logdog_proc.terminate()
154            timeout_retry.WaitFor(logdog_stopped, wait_period=1,
155                                  max_tries=LOGDOG_TERMINATION_TIMEOUT)
156
157            # If logdog_proc hasn't finished by this point, allow
158            # NoLeakingProcesses to kill it.
159
160
161  return result
162
163
164if __name__ == '__main__':
165  sys.exit(main())
166