Bug 1739067: Scope Mach virtualenv to be checkout-specific r=perftest-reviewers,ahal,sparky

Build and run the Mach virtualenv from a `state_dir` that is
"specific-to-topsrcdir".

As part of this, move `get_state_dir()` to `mach` so that it's usable
before `sys.path` entries are fully set up.

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

build/mach_initialize.py

@@ -167,18 +167,6 @@ install a recent enough Python 3.
 
 
 def _activate_python_environment(topsrcdir, state_dir):
-    # We need the "mach" module to access the logic to activate the top-level
-    # Mach site. Since that depends on "packaging" (and, transitively,
-    # "pyparsing"), we add those to the path too.
-    sys.path[0:0] = [
-        os.path.join(topsrcdir, module)
-        for module in (
-            os.path.join("python", "mach"),
-            os.path.join("third_party", "python", "packaging"),
-            os.path.join("third_party", "python", "pyparsing"),
-        )
-    ]
-
     from mach.site import MachSiteManager
 
     mach_environment = MachSiteManager.from_environment(
@@ -212,13 +200,25 @@ def initialize(topsrcdir):
     if os.path.exists(deleted_dir):
         shutil.rmtree(deleted_dir, ignore_errors=True)
 
+    # We need the "mach" module to access the logic to parse virtualenv
+    # requirements. Since that depends on "packaging" (and, transitively,
+    # "pyparsing"), we add those to the path too.
+    sys.path[0:0] = [
+        os.path.join(topsrcdir, module)
+        for module in (
+            os.path.join("python", "mach"),
+            os.path.join("third_party", "python", "packaging"),
+            os.path.join("third_party", "python", "pyparsing"),
+        )
+    ]
+
+    from mach.util import setenv, get_state_dir
+
     state_dir = _create_state_dir()
-    _activate_python_environment(topsrcdir, state_dir)
+    _activate_python_environment(topsrcdir, get_state_dir(True, topsrcdir=topsrcdir))
 
     import mach.base
     import mach.main
-    from mach.util import setenv
-    from mozboot.util import get_state_dir
 
     # Set a reasonable limit to the number of open files.
     #
@@ -322,7 +322,7 @@ def initialize(topsrcdir):
             return state_dir
 
         if key == "local_state_dir":
-            return get_state_dir(srcdir=True)
+            return get_state_dir(specific_to_topsrcdir=True)
 
         if key == "topdir":
             return topsrcdir

python/l10n/test_fluent_migrations/fmt.py

@@ -8,7 +8,7 @@ import shutil
 import sys
 
 import hglib
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 import mozpack.path as mozpath
 
 from compare_locales.merge import merge_channels

python/mach/mach/sentry.py

@@ -10,7 +10,7 @@ from os.path import expanduser
 from threading import Thread
 
 import sentry_sdk
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mach.telemetry import is_telemetry_enabled
 from mozversioncontrol import (
     get_repository_object,

python/mach/mach/telemetry.py

@@ -17,7 +17,8 @@ import six.moves.urllib.parse as urllib_parse
 
 from mach.config import ConfigSettings
 from mach.telemetry_interface import NoopTelemetry, GleanTelemetry
-from mozboot.util import get_state_dir, get_mach_virtualenv_binary
+from mach.util import get_state_dir
+from mozboot.util import get_mach_virtualenv_binary
 from mozbuild.base import MozbuildObject, BuildEnvironmentNotFoundException
 from mozbuild.settings import TelemetrySettings
 from mozbuild.telemetry import filter_args

python/mach/mach/util.py

@@ -4,11 +4,10 @@
 
 from __future__ import absolute_import, unicode_literals
 
+import hashlib
 import os
 import sys
 
-from six import text_type
-
 
 class UserError(Exception):
     """Represents an error caused by something the user did wrong rather than
@@ -21,6 +20,8 @@ def setenv(key, value):
     """Compatibility shim to ensure the proper string type is used with
     os.environ for the version of Python being used.
     """
+    from six import text_type
+
     encoding = "mbcs" if sys.platform == "win32" else "utf-8"
 
     if sys.version_info[0] == 2:
@@ -35,3 +36,46 @@ def setenv(key, value):
             value = value.decode(encoding)
 
     os.environ[key] = value
+
+
+def get_state_dir(specific_to_topsrcdir=False, topsrcdir=None):
+    """Obtain path to a directory to hold state.
+
+    Args:
+        specific_to_topsrcdir (bool): If True, return a state dir specific to the current
+            srcdir instead of the global state dir (default: False)
+
+    Returns:
+        A path to the state dir (str)
+    """
+    state_dir = os.environ.get("MOZBUILD_STATE_PATH", os.path.expanduser("~/.mozbuild"))
+    if not specific_to_topsrcdir:
+        return state_dir
+
+    if not topsrcdir:
+        # Only import MozbuildObject if topsrcdir isn't provided. This is to cover
+        # the Mach initialization stage, where "mozbuild" isn't in the import scope.
+        from mozbuild.base import MozbuildObject
+
+        topsrcdir = os.path.abspath(
+            MozbuildObject.from_environment(cwd=os.path.dirname(__file__)).topsrcdir
+        )
+    # Shortening to 12 characters makes these directories a bit more manageable
+    # in a terminal and is more than good enough for this purpose.
+    srcdir_hash = hashlib.sha256(topsrcdir.encode("utf-8")).hexdigest()[:12]
+
+    state_dir = os.path.join(
+        state_dir, "srcdirs", "{}-{}".format(os.path.basename(topsrcdir), srcdir_hash)
+    )
+
+    if not os.path.isdir(state_dir):
+        # We create the srcdir here rather than 'mach_initialize.py' so direct
+        # consumers of this function don't create the directory inconsistently.
+        print("Creating local state directory: %s" % state_dir)
+        os.makedirs(state_dir, mode=0o770)
+        # Save the topsrcdir that this state dir corresponds to so we can clean
+        # it up in the event its srcdir was deleted.
+        with open(os.path.join(state_dir, "topsrcdir.txt"), "w") as fh:
+            fh.write(topsrcdir)
+
+    return state_dir

python/mozboot/mozboot/base.py

@@ -401,7 +401,7 @@ class BaseBootstrapper(object):
             raise ValueError(
                 "Need a state directory (e.g. ~/.mozbuild) to download " "artifacts"
             )
-        python_location = get_mach_virtualenv_binary(state_dir=self.state_dir)
+        python_location = get_mach_virtualenv_binary()
         if not os.path.exists(python_location):
             raise ValueError("python not found at %s" % python_location)
 

python/mozboot/mozboot/bootstrap.py

@@ -14,7 +14,7 @@ import subprocess
 import time
 from distutils.version import LooseVersion
 from mozfile import which
-from mach.util import UserError
+from mach.util import get_state_dir, UserError
 from mach.telemetry import initialize_telemetry_setting
 
 from mozboot.base import MODERN_RUST_VERSION
@@ -31,7 +31,6 @@ from mozboot.void import VoidBootstrapper
 from mozboot.windows import WindowsBootstrapper
 from mozboot.mozillabuild import MozillaBuildBootstrapper
 from mozboot.mozconfig import find_mozconfig, MozconfigBuilder
-from mozboot.util import get_state_dir
 
 # Use distro package to retrieve linux platform information
 import distro

python/mozboot/mozboot/util.py

@@ -4,19 +4,16 @@
 
 from __future__ import absolute_import, print_function, unicode_literals
 
-import hashlib
 import os
 import platform
 import subprocess
 from subprocess import CalledProcessError
 
 from mach.site import PythonVirtualenv
+from mach.util import get_state_dir
 from mozfile import which
 
 
-here = os.path.join(os.path.dirname(__file__))
-
-
 MINIMUM_RUST_VERSION = "1.53.0"
 
 
@@ -26,53 +23,14 @@ def get_tools_dir(srcdir=False):
     return get_state_dir(srcdir)
 
 
-def get_state_dir(srcdir=False):
-    """Obtain path to a directory to hold state.
-
-    Args:
-        srcdir (bool): If True, return a state dir specific to the current
-            srcdir instead of the global state dir (default: False)
-
-    Returns:
-        A path to the state dir (str)
-    """
-    state_dir = os.environ.get("MOZBUILD_STATE_PATH", os.path.expanduser("~/.mozbuild"))
-    if not srcdir:
-        return state_dir
-
-    # This function can be called without the build virutualenv, and in that
-    # case srcdir is supposed to be False. Import mozbuild here to avoid
-    # breaking that usage.
-    from mozbuild.base import MozbuildObject
-
-    srcdir = os.path.abspath(MozbuildObject.from_environment(cwd=here).topsrcdir)
-    # Shortening to 12 characters makes these directories a bit more manageable
-    # in a terminal and is more than good enough for this purpose.
-    srcdir_hash = hashlib.sha256(srcdir.encode("utf-8")).hexdigest()[:12]
-
-    state_dir = os.path.join(
-        state_dir, "srcdirs", "{}-{}".format(os.path.basename(srcdir), srcdir_hash)
+def get_mach_virtualenv_root():
+    return os.path.join(
+        get_state_dir(specific_to_topsrcdir=True), "_virtualenvs", "mach"
     )
 
-    if not os.path.isdir(state_dir):
-        # We create the srcdir here rather than 'mach_initialize.py' so direct
-        # consumers of this function don't create the directory inconsistently.
-        print("Creating local state directory: %s" % state_dir)
-        os.makedirs(state_dir, mode=0o770)
-        # Save the topsrcdir that this state dir corresponds to so we can clean
-        # it up in the event its srcdir was deleted.
-        with open(os.path.join(state_dir, "topsrcdir.txt"), "w") as fh:
-            fh.write(srcdir)
 
-    return state_dir
-
-
-def get_mach_virtualenv_root(state_dir=None):
-    return os.path.join(state_dir or get_state_dir(), "_virtualenvs", "mach")
-
-
-def get_mach_virtualenv_binary(state_dir=None):
-    root = get_mach_virtualenv_root(state_dir=state_dir)
+def get_mach_virtualenv_binary():
+    root = get_mach_virtualenv_root()
     return PythonVirtualenv(root).python_path
 
 

python/mozbuild/mozbuild/repackaging/msix.py

@@ -24,7 +24,7 @@ import urllib
 
 from six.moves import shlex_quote
 
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mozbuild.util import ensureParentDir
 from mozfile import which
 from mozpack.copier import FileCopier

python/mozperftest/mozperftest/fzf/fzf.py

@@ -8,7 +8,7 @@ from pathlib import Path
 import json
 
 from mozterm import Terminal
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from distutils.spawn import find_executable
 
 

python/mozperftest/mozperftest/runner.py

@@ -184,8 +184,8 @@ def main(argv=sys.argv[1:]):
     from mozbuild.mozconfig import MozconfigLoader
     from mozbuild.base import MachCommandBase, MozbuildObject
     from mozperftest import PerftestArgumentParser
-    from mozboot.util import get_state_dir
     from mach.logging import LoggingManager
+    from mach.util import get_state_dir
 
     mozconfig = SRC_ROOT / "browser" / "config" / "mozconfig"
     if mozconfig.exists():

testing/raptor/mach_commands.py

@@ -18,7 +18,7 @@ import sys
 
 import mozfile
 from mach.decorators import Command
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mozbuild.base import (
     MozbuildObject,
     BinaryNotFoundException,

testing/web-platform/mach_commands.py

@@ -436,7 +436,7 @@ def create_parser_fission_regressions():
 
 def create_parser_testpaths():
     import argparse
-    from mozboot.util import get_state_dir
+    from mach.util import get_state_dir
 
     parser = argparse.ArgumentParser()
     parser.add_argument(

testing/web-platform/manifestupdate.py

@@ -12,7 +12,7 @@ import sys
 
 from six.moves import configparser
 
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 
 from mozlog.structured import commandline
 

tools/lint/python/l10n_lint.py

@@ -5,7 +5,7 @@
 from datetime import datetime, timedelta
 import os
 
-from mozboot import util as mb_util
+from mach import util as mach_util
 from mozlint import result, pathutils
 from mozpack import path as mozpath
 import mozversioncontrol.repoupdate
@@ -23,7 +23,7 @@ PULL_AFTER = timedelta(days=2)
 
 
 def lint(paths, lintconfig, **lintargs):
-    l10n_base = mb_util.get_state_dir()
+    l10n_base = mach_util.get_state_dir()
     root = lintargs["root"]
     exclude = lintconfig.get("exclude")
     extensions = lintconfig.get("extensions")
@@ -81,7 +81,7 @@ def lint(paths, lintconfig, **lintargs):
 
 
 def gecko_strings_setup(**lint_args):
-    gs = mozpath.join(mb_util.get_state_dir(), LOCALE)
+    gs = mozpath.join(mach_util.get_state_dir(), LOCALE)
     marker = mozpath.join(gs, ".hg", "l10n_pull_marker")
     try:
         last_pull = datetime.fromtimestamp(os.stat(marker).st_mtime)

tools/tryselect/docs/tasks.rst

@@ -59,7 +59,7 @@ You can test that everything is working by running these commands:
 
 .. code-block:: shell
 
-    $ statedir=`mach python -c "from mozboot.util import get_state_dir; print(get_state_dir(srcdir=True))"`
+    $ statedir=`mach python -c "from mach.util import get_state_dir; print(get_state_dir(specific_to_topsrcdir=True))"`
     $ rm -rf $statedir/cache/taskgraph
     $ touch taskcluster/mach_commands.py
     # wait a minute for generation to trigger and finish

tools/tryselect/mach_commands.py

@@ -13,7 +13,7 @@ from mach.decorators import (
     SettingsProvider,
     SubCommand,
 )
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mozbuild.base import BuildEnvironmentNotFoundException
 from mozbuild.util import memoize
 

tools/tryselect/push.py

@@ -9,7 +9,7 @@ import sys
 import traceback
 
 import six
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mozbuild.base import MozbuildObject
 from mozversioncontrol import get_repository_object, MissingVCSExtension
 from .util.manage_estimates import (
@@ -51,7 +51,7 @@ build = MozbuildObject.from_environment(cwd=here)
 vcs = get_repository_object(build.topsrcdir)
 
 history_path = os.path.join(
-    get_state_dir(srcdir=True), "history", "try_task_configs.json"
+    get_state_dir(specific_to_topsrcdir=True), "history", "try_task_configs.json"
 )
 
 
@@ -120,7 +120,9 @@ def display_push_estimates(try_task_config):
     if task_labels is None:
         return
 
-    cache_dir = os.path.join(get_state_dir(srcdir=True), "cache", "taskgraph")
+    cache_dir = os.path.join(
+        get_state_dir(specific_to_topsrcdir=True), "cache", "taskgraph"
+    )
 
     graph_cache = None
     dep_cache = None

tools/tryselect/selectors/coverage.py

@@ -15,7 +15,7 @@ import requests
 import datetime
 
 
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mozbuild.base import MozbuildObject
 from mozpack.files import FileFinder
 from moztest.resolve import TestResolver

tools/tryselect/selectors/fuzzy.py

@@ -12,7 +12,7 @@ from distutils.spawn import find_executable
 from distutils.version import StrictVersion
 
 from mozbuild.base import MozbuildObject
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mozterm import Terminal
 
 from ..cli import BaseTryParser
@@ -348,7 +348,9 @@ def run(
     all_tasks = sorted(tg.tasks.keys())
 
     # graph_Cache created by generate_tasks, recreate the path to that file.
-    cache_dir = os.path.join(get_state_dir(srcdir=True), "cache", "taskgraph")
+    cache_dir = os.path.join(
+        get_state_dir(specific_to_topsrcdir=True), "cache", "taskgraph"
+    )
     if full:
         graph_cache = os.path.join(cache_dir, "full_task_graph")
         dep_cache = os.path.join(cache_dir, "full_task_dependencies")

tools/tryselect/tasks.py

@@ -9,7 +9,7 @@ import re
 import sys
 from collections import defaultdict
 
-from mozboot.util import get_state_dir
+from mach.util import get_state_dir
 from mozbuild.base import MozbuildObject
 from mozpack.files import FileFinder
 from moztest.resolve import TestResolver, TestManifestLoader, get_suite_definition
@@ -85,7 +85,9 @@ def generate_tasks(params=None, full=False, disable_target_task_filter=False):
     gecko_taskgraph.fast = True
     generator = TaskGraphGenerator(root_dir=root, parameters=params)
 
-    cache_dir = os.path.join(get_state_dir(srcdir=True), "cache", "taskgraph")
+    cache_dir = os.path.join(
+        get_state_dir(specific_to_topsrcdir=True), "cache", "taskgraph"
+    )
     key = cache_key(attr, generator.parameters, disable_target_task_filter)
     cache = os.path.join(cache_dir, key)
 

tools/tryselect/test/setup.sh

@@ -12,7 +12,7 @@ cat > $MACHRC << EOF
 default=syntax
 EOF
 
-cmd="$topsrcdir/mach python -c 'from mozboot.util import get_state_dir; print(get_state_dir(srcdir=True))'"
+cmd="$topsrcdir/mach python -c 'from mach.util import get_state_dir; print(get_state_dir(specific_to_topsrcdir=True))'"
 # First run local state dir generation so it doesn't affect test output.
 eval $cmd > /dev/null 2>&1
 # Now run it again to get the actual directory.