autoninja: Implement build_telemetry utils

This will replace the user consent logics in ninjalog_uploader_wrapper.py and reclient_metrics.py

See also https://crrev.com/c/5669094 for how this will be used.

Bug: 345113094
Change-Id: Iffc4975c152ba63f7577c09e5254c499fe8973c0
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/5671052
Reviewed-by: Takuto Ikuta <tikuta@chromium.org>
Reviewed-by: Josip Sokcevic <sokcevic@chromium.org>
Reviewed-by: Fumitoshi Ukai <ukai@google.com>
Auto-Submit: Junji Watanabe <jwata@google.com>
Commit-Queue: Josip Sokcevic <sokcevic@chromium.org>

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

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" "$@"

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" "%*"

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 <user>@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())

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

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()