Fix PEX_ROOT leak run and repl goals. (pantsbuild#12066)

Previously we'd leak `.cache/pex_root` to the workspace. Beyond making it
easier to (mistakenly) corrupt the Pex cache for interactive Pex processes
it also could degrade performance of those processes since this was not
a symlink to the `pex_root` named cache used by sandboxed processes.

Force user's of PexEnvironment methods that support setting up a process
to declare what type of process they are setting up by moving those
methods to CompletePexEnvironment with one concrete implementation for
sandboxed processes and one for workspace processes.

Fixes pantsbuild#12055

(cherry picked from commit 0a01eaf)

# Rust tests and lints will be skipped. Delete if not intended.
[ci skip-rust]

# Building wheels and fs_util will be skipped. Delete if not intended.
[ci skip-build-wheels]

src/python/pants/backend/codegen/protobuf/python/rules.py

@@ -21,7 +21,7 @@
     VenvPex,
     VenvPexRequest,
 )
-from pants.backend.python.util_rules.pex_environment import PexEnvironment
+from pants.backend.python.util_rules.pex_environment import SandboxPexEnvironment
 from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
 from pants.core.util_rules.source_files import SourceFilesRequest
 from pants.core.util_rules.stripped_source_files import StrippedSourceFiles
@@ -61,7 +61,7 @@ async def generate_python_from_protobuf(
     grpc_python_plugin: GrpcPythonPlugin,
     python_protobuf_subsystem: PythonProtobufSubsystem,
     python_protobuf_mypy_plugin: PythonProtobufMypyPlugin,
-    pex_environment: PexEnvironment,
+    pex_environment: SandboxPexEnvironment,
 ) -> GeneratedSources:
     download_protoc_request = Get(
         DownloadedExternalTool, ExternalToolRequest, protoc.get_request(Platform.current)
@@ -193,7 +193,7 @@ async def generate_python_from_protobuf(
             description=f"Generating Python sources from {request.protocol_target.address}.",
             level=LogLevel.DEBUG,
             output_directories=(output_dir,),
-            append_only_caches=pex_environment.append_only_caches(),
+            append_only_caches=pex_environment.append_only_caches,
         ),
     )
 

src/python/pants/backend/python/goals/repl.py

@@ -3,7 +3,7 @@
 
 from pants.backend.python.subsystems.ipython import IPython
 from pants.backend.python.util_rules.pex import Pex, PexRequest, PexRequirements
-from pants.backend.python.util_rules.pex_environment import PexEnvironment
+from pants.backend.python.util_rules.pex_environment import WorkspacePexEnvironment
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFiles,
@@ -21,7 +21,9 @@ class PythonRepl(ReplImplementation):
 
 
 @rule(level=LogLevel.DEBUG)
-async def create_python_repl_request(repl: PythonRepl, pex_env: PexEnvironment) -> ReplRequest:
+async def create_python_repl_request(
+    repl: PythonRepl, pex_env: WorkspacePexEnvironment
+) -> ReplRequest:
     requirements_request = Get(
         Pex,
         PexFromTargetsRequest,
@@ -56,7 +58,7 @@ class IPythonRepl(ReplImplementation):
 
 @rule(level=LogLevel.DEBUG)
 async def create_ipython_repl_request(
-    repl: IPythonRepl, ipython: IPython, pex_env: PexEnvironment
+    repl: IPythonRepl, ipython: IPython, pex_env: WorkspacePexEnvironment
 ) -> ReplRequest:
     # Note that we get an intermediate PexRequest here (instead of going straight to a Pex)
     # so that we can get the interpreter constraints for use in ipython_request.

src/python/pants/backend/python/goals/run_pex_binary.py

@@ -10,7 +10,7 @@
     ResolvePexEntryPointRequest,
 )
 from pants.backend.python.util_rules.pex import Pex, PexRequest
-from pants.backend.python.util_rules.pex_environment import PexEnvironment
+from pants.backend.python.util_rules.pex_environment import WorkspacePexEnvironment
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFiles,
@@ -28,7 +28,7 @@
 async def create_pex_binary_run_request(
     field_set: PexBinaryFieldSet,
     pex_binary_defaults: PexBinaryDefaults,
-    pex_env: PexEnvironment,
+    pex_env: WorkspacePexEnvironment,
 ) -> RunRequest:
     entry_point, transitive_targets = await MultiGet(
         Get(

src/python/pants/backend/python/goals/run_pex_binary_integration_test.py

@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import json
+import os
 from textwrap import dedent
 from typing import Optional
 
@@ -124,3 +125,36 @@ def test_no_strip_pex_env_issues_12057() -> None:
         ]
         result = run_pants(args)
         assert result.exit_code == 42, result.stderr
+
+
+def test_no_leak_pex_root_issues_12055() -> None:
+    read_config_result = run_pants(["help-all"])
+    read_config_result.assert_success()
+    config_data = json.loads(read_config_result.stdout)
+    global_advanced_options = {
+        option["config_key"]: [
+            ranked_value["value"] for ranked_value in option["value_history"]["ranked_values"]
+        ][-1]
+        for option in config_data["scope_to_help_info"][""]["advanced"]
+    }
+    named_caches_dir = global_advanced_options["named_caches_dir"]
+
+    sources = {
+        "src/app.py": "import os; print(os.environ['PEX_ROOT'])",
+        "src/BUILD": dedent(
+            """\
+            python_library(name="lib")
+            pex_binary(entry_point="app.py")
+            """
+        ),
+    }
+    with setup_tmpdir(sources) as tmpdir:
+        args = [
+            "--backend-packages=pants.backend.python",
+            f"--source-root-patterns=['/{tmpdir}/src']",
+            "run",
+            f"{tmpdir}/src/app.py",
+        ]
+        result = run_pants(args)
+        result.assert_success()
+        assert os.path.join(named_caches_dir, "pex_root") == result.stdout.strip()

src/python/pants/backend/python/util_rules/pex.py

@@ -29,9 +29,9 @@
 from pants.backend.python.util_rules import pex_cli
 from pants.backend.python.util_rules.pex_cli import PexCliProcess, PexPEX
 from pants.backend.python.util_rules.pex_environment import (
-    PexEnvironment,
     PexRuntimeEnvironment,
     PythonExecutable,
+    SandboxPexEnvironment,
 )
 from pants.engine.addresses import Address, UnparsedAddressInputs
 from pants.engine.collection import Collection, DeduplicatedCollection
@@ -728,7 +728,7 @@ class VenvScriptWriter:
     def _create_venv_script(
         self,
         bash: BashBinary,
-        pex_environment: PexEnvironment,
+        pex_environment: SandboxPexEnvironment,
         *,
         script_path: PurePath,
         venv_executable: PurePath,
@@ -772,14 +772,16 @@ def _create_venv_script(
             content=FileContent(path=str(script_path), content=script.encode(), is_executable=True),
         )
 
-    def exe(self, bash: BashBinary, pex_environment: PexEnvironment) -> VenvScript:
+    def exe(self, bash: BashBinary, pex_environment: SandboxPexEnvironment) -> VenvScript:
         """Writes a safe shim for the venv's executable `pex` script."""
         script_path = PurePath(f"{self.pex.name}_pex_shim.sh")
         return self._create_venv_script(
             bash, pex_environment, script_path=script_path, venv_executable=self.venv_dir / "pex"
         )
 
-    def bin(self, bash: BashBinary, pex_environment: PexEnvironment, name: str) -> VenvScript:
+    def bin(
+        self, bash: BashBinary, pex_environment: SandboxPexEnvironment, name: str
+    ) -> VenvScript:
         """Writes a safe shim for an executable or script in the venv's `bin` directory."""
         script_path = PurePath(f"{self.pex.name}_bin_{name}_shim.sh")
         return self._create_venv_script(
@@ -789,7 +791,7 @@ def bin(self, bash: BashBinary, pex_environment: PexEnvironment, name: str) -> V
             venv_executable=self.venv_dir / "bin" / name,
         )
 
-    def python(self, bash: BashBinary, pex_environment: PexEnvironment) -> VenvScript:
+    def python(self, bash: BashBinary, pex_environment: SandboxPexEnvironment) -> VenvScript:
         """Writes a safe shim for the venv's python binary."""
         return self.bin(bash, pex_environment, "python")
 
@@ -827,7 +829,7 @@ def wrap_venv_prex_request(pex_request: PexRequest) -> VenvPexRequest:
 
 @rule
 async def create_venv_pex(
-    request: VenvPexRequest, bash: BashBinary, pex_environment: PexEnvironment
+    request: VenvPexRequest, bash: BashBinary, pex_environment: SandboxPexEnvironment
 ) -> VenvPex:
     # VenvPex is motivated by improving performance of Python tools by eliminating traditional PEX
     # file startup overhead.
@@ -976,7 +978,7 @@ def __init__(
 
 
 @rule
-async def setup_pex_process(request: PexProcess, pex_environment: PexEnvironment) -> Process:
+async def setup_pex_process(request: PexProcess, pex_environment: SandboxPexEnvironment) -> Process:
     pex = request.pex
     argv = pex_environment.create_argv(pex.name, *request.argv, python=pex.python)
     env = {
@@ -996,7 +998,7 @@ async def setup_pex_process(request: PexProcess, pex_environment: PexEnvironment
         env=env,
         output_files=request.output_files,
         output_directories=request.output_directories,
-        append_only_caches=pex_environment.append_only_caches(),
+        append_only_caches=pex_environment.append_only_caches,
         timeout_seconds=request.timeout_seconds,
         execution_slot_variable=request.execution_slot_variable,
         cache_scope=request.cache_scope,
@@ -1048,7 +1050,7 @@ def __init__(
 
 @rule
 async def setup_venv_pex_process(
-    request: VenvPexProcess, pex_environment: PexEnvironment
+    request: VenvPexProcess, pex_environment: SandboxPexEnvironment
 ) -> Process:
     venv_pex = request.venv_pex
     argv = (venv_pex.pex.argv0, *request.argv)
@@ -1065,7 +1067,7 @@ async def setup_venv_pex_process(
         env=request.extra_env,
         output_files=request.output_files,
         output_directories=request.output_directories,
-        append_only_caches=pex_environment.append_only_caches(),
+        append_only_caches=pex_environment.append_only_caches,
         timeout_seconds=request.timeout_seconds,
         execution_slot_variable=request.execution_slot_variable,
         cache_scope=request.cache_scope,

src/python/pants/backend/python/util_rules/pex_cli.py

@@ -10,9 +10,9 @@
 from pants.backend.python.subsystems.python_native_code import PythonNativeCode
 from pants.backend.python.util_rules import pex_environment
 from pants.backend.python.util_rules.pex_environment import (
-    PexEnvironment,
     PexRuntimeEnvironment,
     PythonExecutable,
+    SandboxPexEnvironment,
 )
 from pants.core.util_rules import external_tool
 from pants.core.util_rules.external_tool import (
@@ -113,7 +113,7 @@ async def download_pex_pex(pex_binary: PexBinary) -> PexPEX:
 async def setup_pex_cli_process(
     request: PexCliProcess,
     pex_binary: PexPEX,
-    pex_env: PexEnvironment,
+    pex_env: SandboxPexEnvironment,
     python_native_code: PythonNativeCode,
     global_options: GlobalOptions,
     pex_runtime_env: PexRuntimeEnvironment,
@@ -179,7 +179,7 @@ async def setup_pex_cli_process(
         env=env,
         output_files=request.output_files,
         output_directories=request.output_directories,
-        append_only_caches=pex_env.append_only_caches(),
+        append_only_caches=pex_env.append_only_caches,
         level=request.level,
         cache_scope=request.cache_scope,
     )

src/python/pants/backend/python/util_rules/pex_environment.py

@@ -1,9 +1,11 @@
 # Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from __future__ import annotations
+
 import os
 from dataclasses import dataclass
-from pathlib import PurePath
+from pathlib import Path, PurePath
 from textwrap import dedent
 from typing import Mapping, Optional, Tuple, cast
 
@@ -106,45 +108,9 @@ class PexEnvironment(EngineAwareReturnType):
     path: Tuple[str, ...]
     interpreter_search_paths: Tuple[str, ...]
     subprocess_environment_dict: FrozenDict[str, str]
+    named_caches_dir: PurePath
     bootstrap_python: Optional[PythonExecutable] = None
 
-    def create_argv(
-        self, pex_filepath: str, *args: str, python: Optional[PythonExecutable] = None
-    ) -> Tuple[str, ...]:
-        python = python or self.bootstrap_python
-        if python:
-            return (python.path, pex_filepath, *args)
-        if os.path.basename(pex_filepath) == pex_filepath:
-            return (f"./{pex_filepath}", *args)
-        return (pex_filepath, *args)
-
-    def environment_dict(self, *, python_configured: bool) -> Mapping[str, str]:
-        """The environment to use for running anything with PEX.
-
-        If the Process is run with a pre-selected Python interpreter, set `python_configured=True`
-        to avoid PEX from trying to find a new interpreter.
-        """
-        d = dict(
-            PATH=create_path_env_var(self.path),
-            PEX_INHERIT_PATH="false",
-            PEX_IGNORE_RCFILES="true",
-            PEX_ROOT=str(self.pex_root),
-            **self.subprocess_environment_dict,
-        )
-        # NB: We only set `PEX_PYTHON_PATH` if the Python interpreter has not already been
-        # pre-selected by Pants. Otherwise, Pex would inadvertently try to find another interpreter
-        # when running PEXes. (Creating a PEX will ignore this env var in favor of `--python-path`.)
-        if not python_configured:
-            d["PEX_PYTHON_PATH"] = create_path_env_var(self.interpreter_search_paths)
-        return d
-
-    @property
-    def pex_root(self) -> PurePath:
-        return PurePath(".cache/pex_root")
-
-    def append_only_caches(self) -> Mapping[str, str]:
-        return {"pex_root": str(self.pex_root)}
-
     def level(self) -> LogLevel:
         return LogLevel.DEBUG if self.bootstrap_python else LogLevel.WARN
 
@@ -233,9 +199,90 @@ def first_python_binary() -> Optional[PythonExecutable]:
             python_setup.interpreter_search_paths(pex_relevant_environment)
         ),
         subprocess_environment_dict=subprocess_env_vars.vars,
+        # TODO: This path normalization is duplicated with `engine_initializer.py`. How can we do
+        #  the normalization only once, via the options system?
+        named_caches_dir=Path(global_options.options.named_caches_dir).resolve(),
         bootstrap_python=first_python_binary(),
     )
 
 
+@dataclass(frozen=True)
+class CompletePexEnvironment:
+    _pex_environment: PexEnvironment
+    pex_root: PurePath
+
+    _PEX_ROOT_DIRNAME = "pex_root"
+
+    @property
+    def interpreter_search_paths(self) -> Tuple[str, ...]:
+        return self._pex_environment.interpreter_search_paths
+
+    def create_argv(
+        self, pex_filepath: str, *args: str, python: Optional[PythonExecutable] = None
+    ) -> Tuple[str, ...]:
+        python = python or self._pex_environment.bootstrap_python
+        if python:
+            return (python.path, pex_filepath, *args)
+        if os.path.basename(pex_filepath) == pex_filepath:
+            return (f"./{pex_filepath}", *args)
+        return (pex_filepath, *args)
+
+    def environment_dict(self, *, python_configured: bool) -> Mapping[str, str]:
+        """The environment to use for running anything with PEX.
+
+        If the Process is run with a pre-selected Python interpreter, set `python_configured=True`
+        to avoid PEX from trying to find a new interpreter.
+        """
+        d = dict(
+            PATH=create_path_env_var(self._pex_environment.path),
+            PEX_INHERIT_PATH="false",
+            PEX_IGNORE_RCFILES="true",
+            PEX_ROOT=str(self.pex_root),
+            **self._pex_environment.subprocess_environment_dict,
+        )
+        # NB: We only set `PEX_PYTHON_PATH` if the Python interpreter has not already been
+        # pre-selected by Pants. Otherwise, Pex would inadvertently try to find another interpreter
+        # when running PEXes. (Creating a PEX will ignore this env var in favor of `--python-path`.)
+        if not python_configured:
+            d["PEX_PYTHON_PATH"] = create_path_env_var(self.interpreter_search_paths)
+        return d
+
+
+@dataclass(frozen=True)
+class SandboxPexEnvironment(CompletePexEnvironment):
+    append_only_caches: FrozenDict[str, str]
+
+    @classmethod
+    def create(cls, pex_environment) -> SandboxPexEnvironment:
+        pex_root = PurePath(".cache") / cls._PEX_ROOT_DIRNAME
+        return cls(
+            _pex_environment=pex_environment,
+            pex_root=pex_root,
+            append_only_caches=FrozenDict({cls._PEX_ROOT_DIRNAME: str(pex_root)}),
+        )
+
+
+class WorkspacePexEnvironment(CompletePexEnvironment):
+    @classmethod
+    def create(cls, pex_environment: PexEnvironment) -> WorkspacePexEnvironment:
+        # N.B.: When running in the workspace the engine doesn't offer an append_only_caches
+        # service to setup a symlink to our named cache for us. As such, we point the PEX_ROOT
+        # directly at the underlying append only cache in that case to re-use results there and
+        # to keep the workspace from being dirtied by the creation of a new Pex cache rooted
+        # there.
+        pex_root = pex_environment.named_caches_dir / cls._PEX_ROOT_DIRNAME
+        return cls(_pex_environment=pex_environment, pex_root=pex_root)
+
+
+@rule
+def sandbox_pex_environment(pex_environment: PexEnvironment) -> SandboxPexEnvironment:
+    return SandboxPexEnvironment.create(pex_environment)
+
+
+@rule
+def workspace_pex_environment(pex_environment: PexEnvironment) -> WorkspacePexEnvironment:
+    return WorkspacePexEnvironment.create(pex_environment)
+
+
 def rules():
     return [*collect_rules(), *process.rules(), *subprocess_environment.rules()]