diff --git a/OWNERS b/OWNERS index 4a8e78afe..e1a54b4a4 100644 --- a/OWNERS +++ b/OWNERS @@ -35,6 +35,9 @@ per-file siso*=file://BUILD_OWNERS per-file reclient*=file://BUILD_OWNERS per-file reclient*=file://RECLIENT_OWNERS +# Build telemetry +per-file build_telemetry*=file://BUILD_OWNERS + # Bazel per-file bazel*=file://CROS_OWNERS per-file bazel*=file://BUILD_OWNERS diff --git a/build_telemetry b/build_telemetry new file mode 100755 index 000000000..52f781490 --- /dev/null +++ b/build_telemetry @@ -0,0 +1,8 @@ +#!/usr/bin/env bash + +# Copyright 2024 The Chromium Authors +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +base_dir=$(dirname "$0") +PYTHONDONTWRITEBYTECODE=1 exec python3 "$base_dir/build_telemetry.py" "$@" diff --git a/build_telemetry.bat b/build_telemetry.bat new file mode 100755 index 000000000..4661119da --- /dev/null +++ b/build_telemetry.bat @@ -0,0 +1,12 @@ +@echo off +:: Copyright 2024 The Chromium Authors +:: Use of this source code is governed by a BSD-style license that can be +:: found in the LICENSE file. +setlocal + +:: Ensure that "depot_tools" is somewhere in PATH so this tool can be used +:: standalone, but allow other PATH manipulations to take priority. +set PATH=%PATH%;%~dp0 + +:: Defer control. +python3 "%~dp0\build_telemetry.py" "%*" diff --git a/build_telemetry.py b/build_telemetry.py new file mode 100755 index 000000000..2659d74b0 --- /dev/null +++ b/build_telemetry.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +# Copyright 2024 The Chromium Authors. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +import argparse +import json +import os +import subprocess +import sys +import textwrap + +import utils + +_DEFAULT_CONFIG_PATH = utils.depot_tools_config_path("build_telemetry.cfg") + +_DEFAULT_COUNTDOWN = 10 + +VERSION = 1 + + +class Config: + + def __init__(self, config_path, countdown): + self._config_path = config_path + self._config = None + self._notice_displayed = False + self._countdown = countdown + + def load(self): + """Loads the build telemetry config.""" + if self._config: + return + + config = {} + if os.path.isfile(self._config_path): + with open(self._config_path) as f: + try: + config = json.load(f) + except Exception: + pass + if config.get("version") != VERSION: + config = None # Reset the state for version change. + + if not config: + config = { + "is_googler": is_googler(), + "status": None, + "countdown": self._countdown, + "version": VERSION, + } + + self._config = config + + def save(self): + with open(self._config_path, "w") as f: + json.dump(self._config, f) + + @property + def path(self): + return self._config_path + + @property + def is_googler(self): + if not self._config: + return + return self._config.get("is_googler") == True + + @property + def countdown(self): + if not self._config: + return + return self._config.get("countdown") + + @property + def version(self): + if not self._config: + return + return self._config.get("version") + + def enabled(self): + if not self._config: + print("WARNING: depot_tools.build_telemetry: %s is not loaded." % + self._config_path, + file=sys.stderr) + return False + if not self._config.get("is_googler"): + return False + if self._config.get("status") == "opt-out": + return False + + if self._should_show_notice(): + remaining = max(0, self._config["countdown"] - 1) + self._show_notice(remaining) + self._notice_displayed = True + self._config["countdown"] = remaining + self.save() + + # Telemetry collection will happen. + return True + + def _should_show_notice(self): + if self._notice_displayed: + return False + if self._config.get("countdown") == 0: + return False + if self._config.get("status") == "opt-in": + return False + return True + + def _show_notice(self, remaining): + """Dispalys notice when necessary.""" + print( + textwrap.dedent(f"""\ + *** NOTICE *** + Google-internal telemetry (including build logs, username, and hostname) is collected on corp machines to diagnose performance and fix build issues. This reminder will be shown {remaining} more times. See http://go/chrome-build-telemetry for details. Hide this notice or opt out by running: build_telemetry [opt-in] [opt-out] + *** END NOTICE *** + """)) + + def opt_in(self): + self._config["status"] = "opt-in" + self.save() + print("build telemetry collection is opted in") + + def opt_out(self): + self._config["status"] = "opt-out" + self.save() + print("build telemetry collection is opted out") + + +def load_config(cfg_path=_DEFAULT_CONFIG_PATH, countdown=_DEFAULT_COUNTDOWN): + """Loads the config from the default location.""" + cfg = Config(cfg_path, countdown) + cfg.load() + return cfg + + +def is_googler(): + """Checks whether this user is Googler or not.""" + p = subprocess.run( + "cipd auth-info", + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + shell=True, + ) + if p.returncode != 0: + return False + lines = p.stdout.splitlines() + if len(lines) == 0: + return False + l = lines[0] + # |l| will be like 'Logged in as @google.com.' for googler using + # reclient. + return l.startswith("Logged in as ") and l.endswith("@google.com.") + + +def enabled(): + """Checks whether the build can upload build telemetry.""" + cfg = load_config() + return cfg.enabled() + + +def main(): + parser = argparse.ArgumentParser(prog="Build Telemetry util") + parser.add_argument('status', nargs=1, choices=['opt-in', 'opt-out']) + args = parser.parse_args() + + cfg = load_config() + + if not cfg.is_googler: + cfg.save() + return + + if args.status == "opt-in": + cfg.opt_in() + return + if args.status == "opt-out": + cfg.opt_out() + return + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/OWNERS b/tests/OWNERS index 72386f9b4..120404cee 100644 --- a/tests/OWNERS +++ b/tests/OWNERS @@ -1,6 +1,7 @@ per-file autoninja_test.py=brucedawson@chromium.org per-file autoninja_test.py=file://BUILD_OWNERS per-file bazel_test.py=file://CROS_OWNERS +per-file build_telemetry_test.py=file://BUILD_OWNERS per-file gn_helper_test.py=file://BUILD_OWNERS per-file ninjalog_uploader_test.py=tikuta@chromium.org per-file reclient*=file://BUILD_OWNERS diff --git a/tests/build_telemetry_test.py b/tests/build_telemetry_test.py new file mode 100755 index 000000000..f71c6bd26 --- /dev/null +++ b/tests/build_telemetry_test.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +# Copyright (c) 2024 The Chromium Authors. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. + +import os +import sys +import tempfile +import unittest +import unittest.mock + +ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT_DIR) + +import build_telemetry + + +class BuildTelemetryTest(unittest.TestCase): + + def test_is_googler(self): + with unittest.mock.patch('subprocess.run') as run_mock: + run_mock.return_value.returncode = 0 + run_mock.return_value.stdout = 'Logged in as foo@google.com.\n' + self.assertTrue(build_telemetry.is_googler()) + + with unittest.mock.patch('subprocess.run') as run_mock: + run_mock.return_value.returncode = 1 + self.assertFalse(build_telemetry.is_googler()) + + with unittest.mock.patch('subprocess.run') as run_mock: + run_mock.return_value.returncode = 0 + run_mock.return_value.stdout = '' + self.assertFalse(build_telemetry.is_googler()) + + with unittest.mock.patch('subprocess.run') as run_mock: + run_mock.return_value.returncode = 0 + run_mock.return_value.stdout = 'Logged in as foo@example.com.\n' + self.assertFalse(build_telemetry.is_googler()) + + def test_load_and_save_config(self): + test_countdown = 2 + with tempfile.TemporaryDirectory() as tmpdir: + cfg_path = os.path.join(tmpdir, "build_telemetry.cfg") + with unittest.mock.patch( + 'build_telemetry.is_googler') as is_googler: + is_googler.return_value = True + + # Initial config load + cfg = build_telemetry.load_config(cfg_path, test_countdown) + self.assertEqual(cfg.path, cfg_path) + self.assertTrue(cfg.is_googler) + self.assertEqual(cfg.countdown, test_countdown) + self.assertEqual(cfg.version, build_telemetry.VERSION) + + cfg.save() + + # 2nd config load + cfg = build_telemetry.load_config(cfg_path, test_countdown) + self.assertEqual(cfg.path, cfg_path) + self.assertTrue(cfg.is_googler) + self.assertEqual(cfg.countdown, test_countdown) + self.assertEqual(cfg.version, build_telemetry.VERSION) + + # build_telemetry.is_googler() is an expensive call. + # The cached result should be reused. + is_googler.assert_called_once() + + def test_enabled(self): + test_countdown = 2 + + # Googler auto opt-in. + with tempfile.TemporaryDirectory() as tmpdir: + cfg_path = os.path.join(tmpdir, "build_telemetry.cfg") + with unittest.mock.patch( + 'build_telemetry.is_googler') as is_googler: + is_googler.return_value = True + + # Initial config load + cfg = build_telemetry.load_config(cfg_path, test_countdown) + cfg._show_notice = unittest.mock.MagicMock() + self.assertEqual(cfg.countdown, test_countdown) + + # 1st enabled() call should print the notice and + # change the countdown. + self.assertTrue(cfg.enabled()) + self.assertEqual(cfg.countdown, test_countdown - 1) + cfg._show_notice.assert_called_once() + cfg._show_notice.reset_mock() + + # 2nd enabled() call shouldn't print the notice and + # change the countdown. + self.assertTrue(cfg.enabled()) + self.assertEqual(cfg.countdown, test_countdown - 1) + cfg._show_notice.assert_not_called() + + cfg.save() + + # 2nd config load + cfg = build_telemetry.load_config(cfg_path) + cfg._show_notice = unittest.mock.MagicMock() + self.assertTrue(cfg.enabled()) + self.assertEqual(cfg.countdown, test_countdown - 2) + cfg._show_notice.assert_called_once() + + cfg.save() + + # 3rd config load + cfg = build_telemetry.load_config(cfg_path) + cfg._show_notice = unittest.mock.MagicMock() + self.assertTrue(cfg.enabled()) + self.assertEqual(cfg.countdown, 0) + cfg._show_notice.assert_not_called() + + # Googler opt-in/opt-out. + with tempfile.TemporaryDirectory() as tmpdir: + cfg_path = os.path.join(tmpdir, "build_telemetry.cfg") + with unittest.mock.patch( + 'build_telemetry.is_googler') as is_googler: + is_googler.return_value = True + # After opt-out, it should not display the notice and + # change the countdown. + cfg = build_telemetry.load_config(cfg_path, test_countdown) + cfg.opt_out() + + cfg = build_telemetry.load_config(cfg_path, test_countdown) + cfg._show_notice = unittest.mock.MagicMock() + self.assertFalse(cfg.enabled()) + self.assertEqual(cfg.countdown, test_countdown) + cfg._show_notice.assert_not_called() + + # After opt-in, it should not display the notice and + # change the countdown. + cfg = build_telemetry.load_config(cfg_path, test_countdown) + cfg.opt_in() + + cfg = build_telemetry.load_config(cfg_path, test_countdown) + cfg._show_notice = unittest.mock.MagicMock() + self.assertTrue(cfg.enabled()) + self.assertEqual(cfg.countdown, test_countdown) + cfg._show_notice.assert_not_called() + + # Non-Googler + with tempfile.TemporaryDirectory() as tmpdir: + cfg_path = os.path.join(tmpdir, "build_telemetry.cfg") + with unittest.mock.patch( + 'build_telemetry.is_googler') as is_googler: + is_googler.return_value = False + cfg = build_telemetry.load_config(cfg_path) + self.assertFalse(cfg.enabled()) + + +if __name__ == '__main__': + unittest.main()