| # Copyright 2020 Google LLC |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # https://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """Utilities for RBE-enabled builds.""" |
| |
| import os |
| import random |
| import subprocess |
| import tempfile |
| |
| # These are the environment variables that control RBE usage with the |
| # --use_rbe flag. If defined on the environment, the values will be |
| # propagated to the build; otherwise, those defaults will be used. |
| TOOLS_DIR = 'prebuilts/remoteexecution-client/latest' |
| _RBE_ENV = { |
| 'USE_RBE': 'true', |
| 'RBE_DIR': TOOLS_DIR, |
| 'NINJA_REMOTE_NUM_JOBS': '500', |
| 'FLAG_log_dir': 'out', |
| 'FLAG_server_address': 'unix:///tmp/reproxy_%s.sock' % random.randint(0,100000), |
| 'FLAG_exec_root': '/src', |
| 'FLAG_invocation_id': 'treble-%s' % random.randint(0,100000), |
| 'RBE_use_application_default_credentials': 'true', |
| 'RBE_reproxy_wait_seconds': '20', |
| 'RBE_output_dir': 'out', |
| 'RBE_proxy_log_dir': 'out', |
| 'RBE_cpp_dependency_scanner_plugin': os.path.join(TOOLS_DIR, 'dependency_scanner_go_plugin.so'), |
| 'RBE_re_proxy': os.path.join(TOOLS_DIR, 'reproxy'), |
| 'RBE_JAVAC': 'true', |
| 'RBE_D8': 'true', |
| 'RBE_R8': 'true', |
| } |
| |
| |
| def get_nsjail_bin_wrapper(): |
| """Returns the command executed in a closed network namespace.""" |
| return ['netns-exec', 'rbe-closed-ns'] |
| |
| |
| def env_array_to_dict(env_array): |
| """Converts an env var array to a dict. |
| |
| Args: |
| env: An array of environment variables in the `var=val` syntax. |
| |
| Returns: |
| A dict of string values keyed by string names. |
| """ |
| env_dict = {} |
| for var in env_array: |
| var = var.split('=') |
| name = var[0] |
| value = var[1] |
| env_dict[name] = value |
| return env_dict |
| |
| def prepare_env(env): |
| """Prepares an env dict for enabling RBE. |
| |
| Checks that all environment variables required to be set |
| by the user are defined and sets some default |
| values for optional environment variables |
| |
| Args: |
| env: An array of environment variables in the `var=val` syntax. |
| |
| Returns: |
| An array of environment variables in the `var=val` syntax. |
| """ |
| # Start with the default values |
| prepared_env = _RBE_ENV.copy() |
| |
| # Host environment variables take precedence over defaults. |
| for k,v in os.environ.items(): |
| if k.startswith('RBE_'): |
| prepared_env[k] = v |
| |
| # Input parameter variables take precedence over everything else |
| prepared_env.update(env_array_to_dict(env)) |
| |
| if 'RBE_instance' not in prepared_env: |
| raise EnvironmentError('The RBE_instance environment ' |
| 'variables must be defined') |
| |
| if 'RBE_service' not in prepared_env: |
| raise EnvironmentError('The RBE_service environment ' |
| 'variables must be defined') |
| |
| return ['%s=%s' % (k,v) for k,v in prepared_env.items()] |
| |
| |
| def get_readonlybind_mounts(): |
| """Returns a dictionary of readonly bind mounts""" |
| creds_file = '.config/gcloud/application_default_credentials.json' |
| # Bind the gcloud credentials file, if present, to authenticate. |
| source_creds_file = os.path.join(os.getenv('HOME'), creds_file) |
| dest_creds_file = os.path.join('/tmp', creds_file) |
| if not os.path.exists(source_creds_file): |
| raise IOError('Required credentials file not found: ' + source_creds_file) |
| return ['%s:%s' % (source_creds_file, dest_creds_file)] |
| |
| |
| def get_extra_nsjail_args(): |
| """Returns a dictionary of extra nsjail.run arguments for RBE.""" |
| # The nsjail should be invoked in a closed network namespace. |
| return ['--disable_clone_newnet'] |
| |
| |
| def setup(env, build_log=subprocess.DEVNULL): |
| """Prerequisite for having RBE enabled for the build. |
| |
| Calls RBE http proxy in a separate network namespace. |
| |
| Args: |
| env: An array of environment variables in the `var=val` syntax. |
| build_log: a file handle to write executed commands to. |
| |
| Returns: |
| A cleanup function to be called after the build is done. |
| """ |
| env_dict = env_array_to_dict(env) |
| |
| # Create the RBE http proxy allowlist file. |
| if 'RBE_service' in env_dict: |
| rbe_service = env_dict['RBE_service'] |
| else: |
| rbe_service = os.getenv('RBE_service') |
| if not rbe_service: |
| raise EnvironmentError('The RBE_service environment ' |
| 'variables must be defined') |
| if ':' in rbe_service: |
| rbe_service = rbe_service.split(':', 1)[0] |
| rbe_allowlist = [ |
| rbe_service, |
| 'oauth2.googleapis.com', |
| 'accounts.google.com', |
| ] |
| with open('/tmp/rbe_allowlist.txt', 'w+') as t: |
| for w in rbe_allowlist: |
| t.write(w + '\n') |
| |
| # Restart RBE http proxy. |
| script_dir = os.path.dirname(os.path.abspath(__file__)) |
| proxy_kill_command = ['killall', 'tinyproxy'] |
| port = 8000 + random.randint(0,1000) |
| new_conf_contents = '' |
| with open(os.path.join(script_dir, 'rbe_http_proxy.conf'), 'r') as base_conf: |
| new_conf_contents = base_conf.read() |
| with tempfile.NamedTemporaryFile(prefix='rbe_http_proxy_', mode='w', delete=False) as new_conf: |
| new_conf.write(new_conf_contents) |
| new_conf.write('\nPort %i\n' % port) |
| new_conf.close() |
| env.append("RBE_HTTP_PROXY=10.1.2.1:%i" % port) |
| |
| proxy_command = [ |
| 'netns-exec', 'rbe-open-ns', 'tinyproxy', '-c', new_conf.name, '-d'] |
| rbe_proxy_log = tempfile.NamedTemporaryFile(prefix='tinyproxy_', delete=False) |
| if build_log != subprocess.DEVNULL: |
| print('RBE http proxy restart commands:', file=build_log) |
| print(' '.join(proxy_kill_command), file=build_log) |
| print('cd ' + script_dir, file=build_log) |
| print(' '.join(proxy_command) + ' &> ' + rbe_proxy_log.name + ' &', |
| file=build_log) |
| subprocess.call( |
| proxy_kill_command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) |
| rbe_proxy = subprocess.Popen( |
| proxy_command, |
| cwd=script_dir, |
| stdout=rbe_proxy_log, |
| stderr=rbe_proxy_log) |
| |
| def cleanup(): |
| """Should be called after an RBE build is done.""" |
| if build_log != subprocess.DEVNULL: |
| print('RBE http proxy kill command:', file=build_log) |
| print(' '.join(proxy_kill_command), file=build_log) |
| rbe_proxy.terminate() |
| # TODO(diegowilson): Calling wait() sometimes dead locks. |
| # Not sure if it's a tinyproxy bug or the issue described in the wait() documentation |
| # https://docs.python.org/2/library/subprocess.html#subprocess.Popen.wait |
| # rbe_proxy.wait() |
| rbe_proxy_log.close() |
| |
| return cleanup |