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