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