1#!/usr/bin/env python
2# Copyright 2016 The Chromium Authors. All rights reserved.
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 logging
11import os
12import signal
13import subprocess
14import sys
15
16_SRC_PATH = os.path.abspath(os.path.join(
17    os.path.dirname(__file__), '..', '..', '..'))
18sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'catapult', 'devil'))
19sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'catapult', 'common',
20                             'py_utils'))
21
22from devil.utils import signal_handler
23from devil.utils import timeout_retry
24from py_utils import tempfile_ext
25
26PROJECT = 'chromium'
27OUTPUT = 'logdog'
28COORDINATOR_HOST = 'luci-logdog.appspot.com'
29SERVICE_ACCOUNT_JSON = ('/creds/service_accounts'
30                        '/service-account-luci-logdog-publisher.json')
31LOGDOG_TERMINATION_TIMEOUT = 30
32
33
34def CommandParser():
35  # Parses the command line arguments being passed in
36  parser = argparse.ArgumentParser()
37  parser.add_argument(
38      '--target',
39      help='The test target to be run. If not set, any extra '
40      'args passed to this script are assumed to be the '
41      'full test command to run.')
42  parser.add_argument('--logdog-bin-cmd', required=True,
43                      help='The logdog bin cmd.')
44  return parser
45
46
47def CreateStopTestsMethod(proc):
48  def StopTests(signum, _frame):
49    logging.error('Forwarding signal %s to test process', str(signum))
50    proc.send_signal(signum)
51  return StopTests
52
53
54@contextlib.contextmanager
55def NoLeakingProcesses(popen):
56  try:
57    yield popen
58  finally:
59    if popen is not None:
60      try:
61        if popen.poll() is None:
62          popen.kill()
63      except OSError:
64        logging.warning('Failed to kill %s. Process may be leaked.',
65                        str(popen.pid))
66
67
68def main():
69  parser = CommandParser()
70  args, extra_cmd_args = parser.parse_known_args(sys.argv[1:])
71
72  logging.basicConfig(level=logging.INFO)
73  if args.target:
74    test_cmd = [os.path.join('bin', 'run_%s' % args.target), '-v']
75    test_cmd += extra_cmd_args
76  else:
77    test_cmd = extra_cmd_args
78
79  test_env = dict(os.environ)
80  logdog_cmd = []
81
82  with tempfile_ext.NamedTemporaryDirectory(
83      prefix='tmp_android_logdog_wrapper') as temp_directory:
84    if not os.path.exists(args.logdog_bin_cmd):
85      logging.error(
86          'Logdog binary %s unavailable. Unable to create logdog client',
87          args.logdog_bin_cmd)
88    else:
89      streamserver_uri = 'unix:%s' % os.path.join(temp_directory,
90                                                  'butler.sock')
91      prefix = os.path.join('android', 'swarming', 'logcats',
92                            os.environ.get('SWARMING_TASK_ID'))
93
94      logdog_cmd = [
95          args.logdog_bin_cmd,
96          '-project', PROJECT,
97          '-output', OUTPUT,
98          '-prefix', prefix,
99          '--service-account-json', SERVICE_ACCOUNT_JSON,
100          '-coordinator-host', COORDINATOR_HOST,
101          'serve',
102          '-streamserver-uri', streamserver_uri]
103      test_env.update({
104          'LOGDOG_STREAM_PROJECT': PROJECT,
105          'LOGDOG_STREAM_PREFIX': prefix,
106          'LOGDOG_STREAM_SERVER_PATH': streamserver_uri,
107          'LOGDOG_COORDINATOR_HOST': COORDINATOR_HOST,
108      })
109
110    logdog_proc = None
111    if logdog_cmd:
112      logdog_proc = subprocess.Popen(logdog_cmd)
113
114    with NoLeakingProcesses(logdog_proc):
115      with NoLeakingProcesses(
116          subprocess.Popen(test_cmd, env=test_env)) as test_proc:
117        with signal_handler.SignalHandler(signal.SIGTERM,
118                                          CreateStopTestsMethod(test_proc)):
119          result = test_proc.wait()
120          if logdog_proc:
121            def logdog_stopped():
122              return logdog_proc.poll() is not None
123
124            logdog_proc.terminate()
125            timeout_retry.WaitFor(logdog_stopped, wait_period=1,
126                                  max_tries=LOGDOG_TERMINATION_TIMEOUT)
127
128            # If logdog_proc hasn't finished by this point, allow
129            # NoLeakingProcesses to kill it.
130
131
132  return result
133
134
135if __name__ == '__main__':
136  sys.exit(main())
137