From 54b3bbd1ce66e7ee5fbe85f6368b33e98b9f7b7d Mon Sep 17 00:00:00 2001
From: Ted Mielczarek <ted@mielczarek.org>
Date: Sat, 10 Nov 2018 19:04:30 +0000
Subject: [PATCH] bug 1505205 - don't write telemetry for recursive mach
 command invocations. r=firefox-build-system-reviewers,chmanchester

This change tries to ensure that we don't write telemetry data for mach
commands invoked recursively as part of other mach commands. The intent of
build system telemetry is to only collect data about commands that users are
invoking directly.

There are two ways that we found mach commands can be recursively invoked:
* By running a python subprocess to recursively invoke mach (used in
  `mach bootstrap` to call `mach artifact toolchain`)
* By using `Registrar.dispatch` to delegate to a sub-command (used by many
  build system commands to invoke `mach build`).

The subprocess case is handled here by having mach set a `MACH_MAIN_PID`
environment variable whose value is the current process' pid on startup if it
does not already exist in the environment. Telemetry code then checks that the
value of that variable matches the current pid and skips writing telemetry data
if not.

The dispatch case is handled by making `MachRegistrar` store the current depth
of the command stack and pass it to the `post_dispatch_handler` which will skip
writing telemetry data if depth != 1.

Additionally the `should_skip_dispatch` function in mach_bootstrap is renamed
to `should_skip_telemetry_submission`, which was its original intent. The
combination of checks added in this change should be sufficient for deciding
when to write telemetry data, and we were not collecting telemetry for the set
of mach commands in that function (which included `mach bootstrap`).

In order to facilitate writing a test for the dispatch case this change adds a
`mach python --exec-file` option to execute Python code directly in the context
of the `mach python` command.

Differential Revision: https://phabricator.services.mozilla.com/D11207

--HG--
extra : moz-landing-system : lando
---
 build/mach_bootstrap.py                      | 21 +++++++++++-----
 python/mach/mach/registrar.py                |  5 +++-
 python/mach/mach/test/invoke_mach_command.py |  4 ++++
 python/mach/mach/test/registrar_dispatch.py  |  3 +++
 python/mach/mach/test/test_telemetry.py      | 25 +++++++++++++++++++-
 python/mach_commands.py                      |  9 ++++++-
 6 files changed, 58 insertions(+), 9 deletions(-)
 create mode 100644 python/mach/mach/test/invoke_mach_command.py
 create mode 100644 python/mach/mach/test/registrar_dispatch.py

diff --git a/build/mach_bootstrap.py b/build/mach_bootstrap.py
index fe12c99a68a6..d9efd5075e43 100644
--- a/build/mach_bootstrap.py
+++ b/build/mach_bootstrap.py
@@ -201,7 +201,7 @@ def bootstrap(topsrcdir, mozilla_dir=None):
                 mozversioncontrol.MissingVCSTool):
             return None
 
-    def should_skip_dispatch(context, handler):
+    def should_skip_telemetry_submission(handler):
         # The user is performing a maintenance command.
         if handler.name in ('bootstrap', 'doctor', 'mach-commands', 'vcs-setup',
                             # We call mach environment in client.mk which would cause the
@@ -209,17 +209,22 @@ def bootstrap(topsrcdir, mozilla_dir=None):
                             'environment'):
             return True
 
+        # Never submit data when running in automation or when running tests.
+        if any(e in os.environ for e in ('MOZ_AUTOMATION', 'TASK_ID', 'MACH_TELEMETRY_NO_SUBMIT')):
+            return True
+
         return False
 
     def post_dispatch_handler(context, handler, instance, result,
-                              start_time, end_time, args):
+                              start_time, end_time, depth, args):
         """Perform global operations after command dispatch.
 
 
         For now,  we will use this to handle build system telemetry.
         """
-        # Don't do anything when...
-        if should_skip_dispatch(context, handler):
+        # Don't write telemetry data if this mach command was invoked as part of another
+        # mach command.
+        if depth != 1 or os.environ.get('MACH_MAIN_PID') != str(os.getpid()):
             return
 
         # We have not opted-in to telemetry
@@ -270,8 +275,7 @@ def bootstrap(topsrcdir, mozilla_dir=None):
                       'w') as f:
                 json.dump(data, f, sort_keys=True)
 
-        # Never submit data when running in automation.
-        if 'MOZ_AUTOMATION' in os.environ or 'TASK_ID' in os.environ:
+        if should_skip_telemetry_submission(handler):
             return True
 
         # But only submit about every n-th operation
@@ -320,6 +324,11 @@ def bootstrap(topsrcdir, mozilla_dir=None):
 
         raise AttributeError(key)
 
+    # Note which process is top-level so that recursive mach invocations can avoid writing
+    # telemetry data.
+    if 'MACH_MAIN_PID' not in os.environ:
+        os.environ[b'MACH_MAIN_PID'] = str(os.getpid()).encode('ascii')
+
     driver = mach.main.Mach(os.getcwd())
     driver.populate_context_handler = populate_context
 
diff --git a/python/mach/mach/registrar.py b/python/mach/mach/registrar.py
index 23ac4fd6029c..ff508ed73d84 100644
--- a/python/mach/mach/registrar.py
+++ b/python/mach/mach/registrar.py
@@ -24,6 +24,7 @@ class MachRegistrar(object):
         self.settings_providers = set()
         self.categories = {}
         self.require_conditions = False
+        self.command_depth = 0
 
     def register_command_handler(self, handler):
         name = handler.name
@@ -82,6 +83,7 @@ class MachRegistrar(object):
                 print(self._condition_failed_message(handler.name, fail_conditions))
                 return 1
 
+        self.command_depth += 1
         fn = getattr(instance, handler.method)
 
         start_time = time.time()
@@ -101,7 +103,8 @@ class MachRegistrar(object):
             postrun = getattr(context, 'post_dispatch_handler', None)
             if postrun:
                 postrun(context, handler, instance, result,
-                        start_time, end_time, args=kwargs)
+                        start_time, end_time, self.command_depth, args=kwargs)
+        self.command_depth -= 1
 
         return result
 
diff --git a/python/mach/mach/test/invoke_mach_command.py b/python/mach/mach/test/invoke_mach_command.py
new file mode 100644
index 000000000000..1efa102ef551
--- /dev/null
+++ b/python/mach/mach/test/invoke_mach_command.py
@@ -0,0 +1,4 @@
+import subprocess
+import sys
+
+subprocess.check_call([sys.executable] + sys.argv[1:])
diff --git a/python/mach/mach/test/registrar_dispatch.py b/python/mach/mach/test/registrar_dispatch.py
new file mode 100644
index 000000000000..35f6b508d969
--- /dev/null
+++ b/python/mach/mach/test/registrar_dispatch.py
@@ -0,0 +1,3 @@
+# This code is loaded via `mach python --exec-file`, so it runs in the scope of
+# the `mach python` command.
+self._mach_context.commands.dispatch('uuid', self._mach_context)  # noqa: F821
diff --git a/python/mach/mach/test/test_telemetry.py b/python/mach/mach/test/test_telemetry.py
index 27e59c1c2aca..711650923306 100644
--- a/python/mach/mach/test/test_telemetry.py
+++ b/python/mach/mach/test/test_telemetry.py
@@ -24,7 +24,10 @@ def run_mach(tmpdir):
     # a machrc there.
     update_or_create_build_telemetry_config(unicode(tmpdir.join('machrc')))
     env = dict(os.environ)
-    env['MOZBUILD_STATE_PATH'] = str(tmpdir)
+    env[b'MOZBUILD_STATE_PATH'] = str(tmpdir)
+    env[b'MACH_TELEMETRY_NO_SUBMIT'] = b'1'
+    # Let whatever mach command we invoke from tests believe it's the main command.
+    del env['MACH_MAIN_PID']
     mach = os.path.join(buildconfig.topsrcdir, 'mach')
 
     def run(*args, **kwargs):
@@ -120,5 +123,25 @@ def test_path_filtering_other_cwd(run_mach, tmpdir):
     assert d['argv'] == expected
 
 
+def test_mach_invoke_recursive(run_mach):
+    data = run_mach('python',
+                    os.path.join(os.path.dirname(__file__), 'invoke_mach_command.py'),
+                    os.path.join(buildconfig.topsrcdir, 'mach'),
+                    'uuid')
+    assert len(data) == 1
+    d = data[0]
+    assert d['command'] == 'python'
+
+
+def test_registrar_dispatch(run_mach):
+    # Use --exec-file so this script can use Registrar.dispatch to dispatch a mach command
+    # from within the same interpreter as the `mach python` command.
+    data = run_mach('python', '--exec-file',
+                    os.path.join(os.path.dirname(__file__), 'registrar_dispatch.py'))
+    assert len(data) == 1
+    d = data[0]
+    assert d['command'] == 'python'
+
+
 if __name__ == '__main__':
     mozunit.main()
diff --git a/python/mach_commands.py b/python/mach_commands.py
index ec1163203e66..70cd847abf11 100644
--- a/python/mach_commands.py
+++ b/python/mach_commands.py
@@ -40,8 +40,11 @@ class MachCommands(MachCommandBase):
              description='Run Python.')
     @CommandArgument('--no-virtualenv', action='store_true',
                      help='Do not set up a virtualenv')
+    @CommandArgument('--exec-file',
+                     default=None,
+                     help='Execute this Python file using `execfile`')
     @CommandArgument('args', nargs=argparse.REMAINDER)
-    def python(self, no_virtualenv, args):
+    def python(self, no_virtualenv, exec_file, args):
         # Avoid logging the command
         self.log_manager.terminal_handler.setLevel(logging.CRITICAL)
 
@@ -57,6 +60,10 @@ class MachCommands(MachCommandBase):
             self._activate_virtualenv()
             python_path = self.virtualenv_manager.python_path
 
+        if exec_file:
+            execfile(exec_file)
+            return 0
+
         return self.run_process([python_path] + args,
                                 pass_thru=True,  # Allow user to run Python interactively.
                                 ensure_exit_code=False,  # Don't throw on non-zero exit code.