logdog_wrapper.py (5768B)
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 8 import argparse 9 import contextlib 10 import json 11 import logging 12 import os 13 import shutil 14 import signal 15 import subprocess 16 import sys 17 18 _SRC_PATH = os.path.abspath(os.path.join( 19 os.path.dirname(__file__), '..', '..', '..')) 20 sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'catapult', 'devil')) 21 sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'catapult', 'common', 22 'py_utils')) 23 24 from devil.utils import signal_handler 25 from devil.utils import timeout_retry 26 from py_utils import tempfile_ext 27 28 OUTPUT = 'logdog' 29 COORDINATOR_HOST = 'luci-logdog.appspot.com' 30 LOGDOG_TERMINATION_TIMEOUT = 30 31 32 33 def 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 55 def 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 63 def 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 76 def 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 98 def 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_logdog') as temp_directory: 123 if logdog_butler_bin: 124 if os.path.exists(logdog_butler_bin): 125 streamserver_uri = 'unix:%s' % os.path.join(temp_directory, 126 'butler.sock') 127 prefix = os.path.join('android', 'swarming', 'logcats', 128 os.environ.get('SWARMING_TASK_ID', "")) 129 project = GetProjectFromLuciContext() 130 131 logdog_cmd = [ 132 logdog_butler_bin, '-project', project, '-output', OUTPUT, 133 '-prefix', prefix, '-coordinator-host', COORDINATOR_HOST, 'serve', 134 '-streamserver-uri', streamserver_uri 135 ] 136 test_env.update({ 137 'LOGDOG_STREAM_PROJECT': project, 138 'LOGDOG_STREAM_PREFIX': prefix, 139 'LOGDOG_STREAM_SERVER_PATH': streamserver_uri, 140 'LOGDOG_COORDINATOR_HOST': COORDINATOR_HOST, 141 }) 142 else: 143 logging.warning('--logdog-bin-cmd specified, but binary was not found.' 144 ' Will not be using logdog butler server.') 145 146 logdog_proc = None 147 if logdog_cmd: 148 logdog_proc = subprocess.Popen(logdog_cmd) 149 150 with NoLeakingProcesses(logdog_proc): 151 with NoLeakingProcesses( 152 subprocess.Popen(test_cmd, env=test_env)) as test_proc: 153 with signal_handler.SignalHandler(signal.SIGTERM, 154 CreateStopTestsMethod(test_proc)): 155 result = test_proc.wait() 156 if logdog_proc: 157 def logdog_stopped(): 158 return logdog_proc.poll() is not None 159 160 logdog_proc.terminate() 161 timeout_retry.WaitFor(logdog_stopped, wait_period=1, 162 max_tries=LOGDOG_TERMINATION_TIMEOUT) 163 164 # If logdog_proc hasn't finished by this point, allow 165 # NoLeakingProcesses to kill it. 166 167 168 return result 169 170 171 if __name__ == '__main__': 172 sys.exit(main())