Migrates generate-lockfiles to USES_ENVIRONMENTS. (#17300)

This is a least-effort attempt at getting `generate-lockfiles` into a position where individual backend implementations can select an environment in which to generate the lockfiles. 

The request generation phase runs in the default local environment, and is responsible for adding an `environments` field to the `GenerateLockfile` request object. 

If a request does not specify environments, the `GenerateLockfile` request is run in the default local environment. If a single environment is specified, that environment is to run the request. If more than one environment is specified, a warning is displayed, and the request runs in either the default local environment (if present in the request), or an arbitrary environment. 

This approach will give us an avenue in the future to run the request in _all_ the specified environments, and add behaviour to somehow merge the resulting lockfiles if different, since the eventual act of writing the lockfiles to disk happens in the rule goal code.

Addresses #17129

src/python/pants/core/goals/generate_lockfiles.py

@@ -11,7 +11,7 @@
 from typing import Callable, ClassVar, Iterable, Sequence
 
 from pants.engine.collection import Collection
-from pants.engine.environment import EnvironmentName
+from pants.engine.environment import ChosenLocalEnvironmentName, EnvironmentName
 from pants.engine.fs import Digest, MergeDigests, Workspace
 from pants.engine.goal import Goal, GoalSubsystem
 from pants.engine.internals.selectors import Get, MultiGet
@@ -51,6 +51,14 @@ class GenerateLockfile:
     lockfile_dest: str
 
 
+@dataclass(frozen=True)
+class GenerateLockfileWithEnvironments(GenerateLockfile):
+    """Allows a `GenerateLockfile` subclass to specify which environments the request is compatible
+    with, if the relevant backend supports environments."""
+
+    environments: tuple[EnvironmentName, ...]
+
+
 @dataclass(frozen=True)
 class WrappedGenerateLockfile:
     request: GenerateLockfile
@@ -359,14 +367,15 @@ def activated(cls, union_membership: UnionMembership) -> bool:
 
 class GenerateLockfilesGoal(Goal):
     subsystem_cls = GenerateLockfilesSubsystem
-    environment_behavior = Goal.EnvironmentBehavior.LOCAL_ONLY  # TODO(#17129) — Migrate this.
+    environment_behavior = Goal.EnvironmentBehavior.USES_ENVIRONMENTS
 
 
 @goal_rule
 async def generate_lockfiles_goal(
     workspace: Workspace,
     union_membership: UnionMembership,
     generate_lockfiles_subsystem: GenerateLockfilesSubsystem,
+    local_environment: ChosenLocalEnvironmentName,
 ) -> GenerateLockfilesGoal:
     known_user_resolve_names = await MultiGet(
         Get(KnownUserResolveNames, KnownUserResolveNamesRequest, request())
@@ -378,27 +387,45 @@ async def generate_lockfiles_goal(
         set(generate_lockfiles_subsystem.resolve),
     )
 
+    # This is the "planning" phase of lockfile generation. Currently this is all done in the local
+    # environment, since there's not currently a clear mechanism to prescribe an environment.
     all_specified_user_requests = await MultiGet(
-        Get(UserGenerateLockfiles, RequestedUserResolveNames, resolve_names)
+        Get(
+            UserGenerateLockfiles,
+            {resolve_names: RequestedUserResolveNames, local_environment.val: EnvironmentName},
+        )
         for resolve_names in requested_user_resolve_names
     )
     specified_tool_requests = await MultiGet(
-        Get(WrappedGenerateLockfile, GenerateToolLockfileSentinel, sentinel())
+        Get(
+            WrappedGenerateLockfile,
+            {sentinel(): GenerateToolLockfileSentinel, local_environment.val: EnvironmentName},
+        )
         for sentinel in requested_tool_sentinels
     )
     applicable_tool_requests = filter_tool_lockfile_requests(
         specified_tool_requests,
         resolve_specified=bool(generate_lockfiles_subsystem.resolve),
     )
 
+    # Execute the actual lockfile generation in each request's environment.
+    # Currently, since resolves specify a single filename for output, we pick a resonable
+    # environment to execute the request in. Currently we warn if multiple environments are
+    # specified.
+    all_requests = itertools.chain(*all_specified_user_requests, applicable_tool_requests)
     results = await MultiGet(
-        Get(GenerateLockfileResult, GenerateLockfile, req)
-        for req in (
-            *(req for reqs in all_specified_user_requests for req in reqs),
-            *applicable_tool_requests,
+        Get(
+            GenerateLockfileResult,
+            {
+                req: GenerateLockfile,
+                _preferred_environment(req, local_environment.val): EnvironmentName,
+            },
         )
+        for req in all_requests
     )
 
+    # Lockfiles are actually written here. This would be an acceptable place to handle conflict
+    # resolution behaviour if we start executing requests in multiple environments.
     merged_digest = await Get(Digest, MergeDigests(res.digest for res in results))
     workspace.write_digest(merged_digest)
     for result in results:
@@ -407,6 +434,25 @@ async def generate_lockfiles_goal(
     return GenerateLockfilesGoal(exit_code=0)
 
 
+def _preferred_environment(request: GenerateLockfile, default: EnvironmentName) -> EnvironmentName:
+
+    if not isinstance(request, GenerateLockfileWithEnvironments):
+        return default  # This request has not been migrated to use environments.
+
+    if len(request.environments) == 1:
+        return request.environments[0]
+
+    ret = default if default in request.environments else request.environments[0]
+
+    logger.warning(
+        f"The `{request.__class__.__name__}` for resolve `{request.resolve_name}` specifies more "
+        "than one environment. Pants will generate the lockfile using only the environment "
+        f"`{ret.val}`, which may have unintended effects when executing in the other environments."
+    )
+
+    return ret
+
+
 # -----------------------------------------------------------------------------------------------
 # Helpers for determining the resolve
 # -----------------------------------------------------------------------------------------------

src/python/pants/core/goals/generate_lockfiles_test.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from typing import Iterable
+
 import pytest
 
 from pants.backend.python.subsystems.setup import PythonSetup
@@ -19,15 +21,18 @@
     NO_TOOL_LOCKFILE,
     AmbiguousResolveNamesError,
     GenerateLockfile,
+    GenerateLockfileWithEnvironments,
     GenerateToolLockfileSentinel,
     KnownUserResolveNames,
     NoCompatibleResolveException,
     RequestedUserResolveNames,
     UnrecognizedResolveNamesError,
     WrappedGenerateLockfile,
+    _preferred_environment,
     determine_resolves_to_generate,
     filter_tool_lockfile_requests,
 )
+from pants.engine.environment import EnvironmentName
 from pants.engine.target import Dependencies, Target
 from pants.testutil.option_util import create_subsystem
 from pants.util.strutil import softwrap
@@ -266,3 +271,41 @@ def maybe_get_resolve(t: Target) -> str | None:
             """
         )
     )
+
+
+_default = "_default"
+
+
+@pytest.mark.parametrize(
+    ("env_names", "expected", "in_output", "not_in_output"),
+    (
+        ((), _default, None, "anything"),
+        (("jeremy", "derek"), "jeremy", "`jeremy`, which may have ", None),
+        (("jeremy", "_default"), "_default", "`_default`, which may have ", None),
+        (("_default", "jeremy"), "_default", "`_default`, which may have ", None),
+        (("jeremy",), "jeremy", None, "anything"),
+    ),
+)
+def test_preferred_environment(
+    env_names: Iterable[str],
+    expected: str,
+    in_output: str | None,
+    not_in_output: str | None,
+    caplog,
+):
+
+    resolve_name = "boop"
+    resolve_dest = "beep"
+    if not env_names:
+        request = GenerateLockfile(resolve_name, resolve_dest)
+    else:
+        envs = tuple(EnvironmentName(name) for name in env_names)
+        request = GenerateLockfileWithEnvironments(resolve_name, resolve_dest, envs)
+
+    default = EnvironmentName(_default)
+
+    preferred_env = _preferred_environment(request, default)
+
+    assert preferred_env.val == expected
+    assert in_output is None or in_output in caplog.text
+    assert not_in_output is None or not_in_output not in caplog.text