Implement Nailgun in V2 hermetic execution (pantsbuild#8371)

### Problem

Local hermetic execution of jvm compiles cannot take advantage of Nailgun, missing out on possible performance benefits.

### Solution

Following the discussion from pantsbuild#8387, we have come to the following UX design for now:
- We add a field, `nailgunnable`, to `ExecuteProcessRequest`s (EPR). This bit is off in all cases except for compiling with zinc and rsc.
- We create a new data structure, `NailgunPool`, that will expose a method `connect(req: ExecuteProcessRequest...)`. This is in charge of maintaining the various nailgun server instances alive. When someone calls `connect()`, it will start a server if it needs to, and then store it for all other users.
- We create a new command runner, `NailgunCommandRunner`, that will hold an instance of `NailgunPool` and wrap the `local::CommandRunner`. When asked to run a local EPR, if the EPR is marked _not nailgunnable_, it will forward it to the wrapped `local::CommandRunner`. If **nailgunnable=True**, then:
-   It will parse the CLI args of the request to figure out the startup arguments of the nailgun server. This will turn into an EPR.
-  With that EPR, it will call `NailgunPool::connect` to (maybe start and) retrieve the port of a running nailgun instance.
-  It will modify the original request to be able to connect to nailgun, instead of using the JDK.
-  It will forward this new and modified client request to the `local::CommandRunner`, to run normally.

### Result

Hermetic execution should be able to use nailgun where desired. Fixes pantsbuild#8311, and opens a bunch of followups.

src/python/pants/backend/jvm/tasks/jvm_compile/rsc/rsc_compile.py

@@ -806,6 +806,7 @@ def _runtool_hermetic(self, main, tool_name, distribution, input_digest, ctx):
       # Since this is always hermetic, we need to use `underlying.home` because
       # ExecuteProcessRequest requires an existing, local jdk location.
       jdk_home=distribution.underlying_home,
+      is_nailgunnable=True,
     )
     res = self.context.execute_process_synchronously_without_raising(
       epr,

src/python/pants/backend/jvm/tasks/jvm_compile/zinc/zinc_compile.py

@@ -17,12 +17,15 @@
 from pants.backend.jvm.subsystems.zinc import Zinc
 from pants.backend.jvm.targets.jvm_target import JvmTarget
 from pants.backend.jvm.targets.scalac_plugin import ScalacPlugin
+from pants.backend.jvm.tasks.classpath_entry import ClasspathEntry
 from pants.backend.jvm.tasks.classpath_util import ClasspathUtil
 from pants.backend.jvm.tasks.jvm_compile.jvm_compile import JvmCompile
 from pants.base.build_environment import get_buildroot
 from pants.base.exceptions import TaskError
 from pants.base.hash_utils import hash_file
 from pants.base.workunit import WorkUnitLabel
+from pants.binaries.binary_tool import Script
+from pants.binaries.binary_util import BinaryToolUrlGenerator
 from pants.engine.fs import (
   EMPTY_DIRECTORY_DIGEST,
   DirectoryToMaterialize,
@@ -44,6 +47,18 @@
 logger = logging.getLogger(__name__)
 
 
+class NailgunClientBinary(Script):
+  options_scope = 'nailgun-client'
+  name = 'ng'
+  default_version = '1.0.0'
+
+  def get_external_url_generator(self):
+    class NailgunClientBinaryUrlGenerator(BinaryToolUrlGenerator):
+      def generate_urls(self, version, host_platform):
+        return ["https://github.com/facebook/nailgun/releases/download/nailgun-all-v1.0.0/ng.py"]
+    return NailgunClientBinaryUrlGenerator()
+
+
 class BaseZincCompile(JvmCompile):
   """An abstract base class for zinc compilation tasks."""
 
@@ -154,7 +169,7 @@ def register_options(cls, register):
 
   @classmethod
   def subsystem_dependencies(cls):
-    return super().subsystem_dependencies() + (Zinc.Factory, JvmPlatform,)
+    return super().subsystem_dependencies() + (Zinc.Factory, JvmPlatform, NailgunClientBinary,)
 
   @classmethod
   def prepare(cls, options, round_manager):
@@ -391,6 +406,19 @@ def _compile_nonhermetic(self, jvm_options, ctx, classes_directory):
     if exit_code != 0:
       raise self.ZincCompileError('Zinc compile failed.', exit_code=exit_code)
 
+  # Snapshot the nailgun-server jar, to use it to start nailguns in the hermetic case.
+  # TODO(#8480): Make this jar natively accessible to the engine,
+  #              because it will help when moving the JVM pipeline to v2.
+  @memoized_method
+  def _nailgun_server_classpath_entry(self):
+    nailgun_jar = self.tool_jar('nailgun-server')
+    nailgun_jar_snapshot, = self.context._scheduler.capture_snapshots((PathGlobsAndRoot(
+      PathGlobs((fast_relpath(nailgun_jar, get_buildroot()),)),
+      get_buildroot()
+    ),))
+    nailgun_jar_digest = nailgun_jar_snapshot.directory_digest
+    return ClasspathEntry(nailgun_jar, nailgun_jar_digest)
+
   def _compile_hermetic(self, jvm_options, ctx, classes_dir, jar_file,
                         compiler_bridge_classpath_entry, dependency_classpath,
                         scalac_classpath_entries):
@@ -411,7 +439,10 @@ def _compile_hermetic(self, jvm_options, ctx, classes_dir, jar_file,
     snapshots.extend(java_sources_snapshots)
 
     # Ensure the dependencies and compiler bridge jars are available in the execution sandbox.
-    relevant_classpath_entries = dependency_classpath + [compiler_bridge_classpath_entry]
+    relevant_classpath_entries = (dependency_classpath + [
+      compiler_bridge_classpath_entry,
+      self._nailgun_server_classpath_entry(), # We include nailgun-server, to use it to start servers when needed from the hermetic execution case.
+    ])
     directory_digests = [
       entry.directory_digest for entry in relevant_classpath_entries if entry.directory_digest
     ]
@@ -426,6 +457,9 @@ def _compile_hermetic(self, jvm_options, ctx, classes_dir, jar_file,
       classpath_entry.directory_digest for classpath_entry in scalac_classpath_entries
     )
 
+    _, nailgun_client_snapshot = NailgunClientBinary.global_instance().hackily_snapshot(self.context)
+    directory_digests.append(nailgun_client_snapshot.directory_digest)
+
     if self._zinc.use_native_image:
       if jvm_options:
         raise ValueError(
@@ -493,6 +527,7 @@ def _compile_hermetic(self, jvm_options, ctx, classes_dir, jar_file,
       description="zinc compile for {}".format(ctx.target.address.spec),
       unsafe_local_only_files_because_we_favor_speed_over_correctness_for_this_rule=merged_local_only_scratch_inputs,
       jdk_home=self._zinc.underlying_dist.home,
+      is_nailgunnable=True,
     )
     res = self.context.execute_process_synchronously_or_raise(
       req, self.name(), [WorkUnitLabel.COMPILER])

src/python/pants/engine/isolated_process.py

@@ -37,6 +37,7 @@ class ExecuteProcessRequest:
   timeout_seconds: Union[int, float]
   unsafe_local_only_files_because_we_favor_speed_over_correctness_for_this_rule: Digest
   jdk_home: Optional[str]
+  is_nailgunnable: bool
 
   def __init__(
     self,
@@ -49,7 +50,8 @@ def __init__(
     output_directories: Optional[Tuple[str, ...]] = None,
     timeout_seconds: Union[int, float] = _default_timeout_seconds,
     unsafe_local_only_files_because_we_favor_speed_over_correctness_for_this_rule: Digest = EMPTY_DIRECTORY_DIGEST,
-    jdk_home: Optional[str] = None
+    jdk_home: Optional[str] = None,
+    is_nailgunnable: bool = False,
   ) -> None:
     self.argv = argv
     self.input_files = input_files
@@ -60,6 +62,7 @@ def __init__(
     self.timeout_seconds = timeout_seconds
     self.unsafe_local_only_files_because_we_favor_speed_over_correctness_for_this_rule = unsafe_local_only_files_because_we_favor_speed_over_correctness_for_this_rule
     self.jdk_home = jdk_home
+    self.is_nailgunnable = is_nailgunnable
 
 
 @frozen_after_init

src/python/pants/engine/native.py

@@ -904,6 +904,8 @@ def ti(type_obj):
         self.context.utf8_buf(execution_options.process_execution_speculation_strategy),
         execution_options.process_execution_use_local_cache,
         self.context.utf8_dict(execution_options.remote_execution_headers),
+        self.context.utf8_buf(sys.executable),
+        execution_options.process_execution_local_enable_nailgun,
       )
     if scheduler_result.is_throw:
       value = self.context.from_value(scheduler_result.throw_handle)

src/python/pants/option/global_options.py

@@ -61,6 +61,7 @@ class ExecutionOptions:
   remote_oauth_bearer_token_path: Any
   remote_execution_extra_platform_properties: Any
   remote_execution_headers: Any
+  process_execution_local_enable_nailgun: bool
 
   @classmethod
   def from_bootstrap_options(cls, bootstrap_options):
@@ -85,6 +86,7 @@ def from_bootstrap_options(cls, bootstrap_options):
       remote_oauth_bearer_token_path=bootstrap_options.remote_oauth_bearer_token_path,
       remote_execution_extra_platform_properties=bootstrap_options.remote_execution_extra_platform_properties,
       remote_execution_headers=bootstrap_options.remote_execution_headers,
+      process_execution_local_enable_nailgun=bootstrap_options.process_execution_local_enable_nailgun,
     )
 
 
@@ -109,6 +111,7 @@ def from_bootstrap_options(cls, bootstrap_options):
     remote_oauth_bearer_token_path=None,
     remote_execution_extra_platform_properties=[],
     remote_execution_headers={},
+    process_execution_local_enable_nailgun=False,
   )
 
 
@@ -438,6 +441,9 @@ def register_bootstrap_options(cls, register):
              advanced=True)
     register('--process-execution-use-local-cache', type=bool, default=True, advanced=True,
              help='Whether to keep process executions in a local cache persisted to disk.')
+    register('--process-execution-local-enable-nailgun', type=bool, default=DEFAULT_EXECUTION_OPTIONS.process_execution_local_enable_nailgun,
+             help='Whether or not to use nailgun to run the requests that are marked as nailgunnable.',
+             advanced=True)
 
   @classmethod
   def register_options(cls, register):