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