go: add support for test coverage (pantsbuild#16550)

Update the Go backend to gather code coverage data for packages under test. The coverage files are written under `dist/coverage/go/IMPORT_PATH` directories (where `IMPORT_PATH` has slashes escaped to underscores). The coverage output is the raw `cover.out` format from the `go` tooling and is not rendered by this PR (which can be done in a follow-on since it is a bit involved).

[ci skip-rust]

[ci skip-build-wheels]

src/python/pants/backend/experimental/go/register.py

@@ -16,6 +16,8 @@
     assembly,
     build_pkg,
     build_pkg_target,
+    coverage,
+    coverage_output,
     first_party_pkg,
     go_bootstrap,
     go_mod,
@@ -39,6 +41,8 @@ def rules():
         *build_pkg.rules(),
         *build_pkg_target.rules(),
         *check.rules(),
+        *coverage.rules(),
+        *coverage_output.rules(),
         *third_party_pkg.rules(),
         *go_bootstrap.rules(),
         *goroot.rules(),

src/python/pants/backend/go/go_sources/generate_testmain/main.go

@@ -80,6 +80,10 @@ type Analysis struct {
 
 	// True if Go 1.18 is in use. This is not set by analyze but rather by generate based on release tags.
 	IsGo1_18 bool
+
+	// True if coverage is enabled. This is not set by `analyze` but rather by `generate` based on an
+	// environment variable set by the invoker.
+	Cover bool
 }
 
 // isTestFunc tells whether fn has the type of a testing function. arg
@@ -323,6 +327,10 @@ func init() {
 }
 
 func main() {
+{{if .Cover}}
+	registerCover()
+{{end}}
+
 {{- if .IsGo1_18 }}
 	m := testing.MainStart(testdeps.TestDeps{}, tests, benchmarks, fuzzTargets, examples)
 {{- else }}
@@ -337,7 +345,7 @@ func main() {
 }
 `
 
-func generate(analysis *Analysis) ([]byte, error) {
+func generate(analysis *Analysis, genCover bool) ([]byte, error) {
 	tmpl, err := template.New("testmain").Parse(testMainTemplate)
 	if err != nil {
 		return nil, err
@@ -351,6 +359,9 @@ func generate(analysis *Analysis) ([]byte, error) {
 		}
 	}
 
+	// Pass through the config to generate the call to the coverage stubs.
+	analysis.Cover = genCover
+
 	var buffer bytes.Buffer
 
 	err = tmpl.Execute(&buffer, analysis)
@@ -368,7 +379,9 @@ func main() {
 		os.Exit(1)
 	}
 
-	testmain, err := generate(analysis)
+	genCover := os.Getenv("GENERATE_COVER") != ""
+
+	testmain, err := generate(analysis, genCover)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Failed to generate _testmain.go: %s\n", err)
 		os.Exit(1)

src/python/pants/backend/go/goals/test.py

@@ -16,10 +16,17 @@
 )
 from pants.backend.go.util_rules.build_pkg import (
     BuildGoPackageRequest,
+    BuiltGoPackage,
     FallibleBuildGoPackageRequest,
     FallibleBuiltGoPackage,
 )
 from pants.backend.go.util_rules.build_pkg_target import BuildGoPackageTargetRequest
+from pants.backend.go.util_rules.coverage import (
+    GenerateCoverageSetupCodeRequest,
+    GenerateCoverageSetupCodeResult,
+    GoCoverageConfig,
+    GoCoverageData,
+)
 from pants.backend.go.util_rules.first_party_pkg import (
     FallibleFirstPartyPkgAnalysis,
     FallibleFirstPartyPkgDigest,
@@ -41,6 +48,7 @@
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.engine.environment import Environment, EnvironmentRequest
 from pants.engine.fs import EMPTY_FILE_DIGEST, AddPrefix, Digest, MergeDigests
+from pants.engine.internals.native_engine import EMPTY_DIGEST
 from pants.engine.process import FallibleProcessResult, Process, ProcessCacheScope
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import Dependencies, DependenciesRequest, SourcesField, Target, Targets
@@ -189,17 +197,18 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
     testmain = await Get(
         GeneratedTestMain,
         GenerateTestMainRequest(
-            pkg_digest.digest,
-            FrozenOrderedSet(
+            digest=pkg_digest.digest,
+            test_paths=FrozenOrderedSet(
                 os.path.join(".", pkg_analysis.dir_path, name)
                 for name in pkg_analysis.test_go_files
             ),
-            FrozenOrderedSet(
+            xtest_paths=FrozenOrderedSet(
                 os.path.join(".", pkg_analysis.dir_path, name)
                 for name in pkg_analysis.xtest_go_files
             ),
-            import_path,
-            field_set.address,
+            import_path=import_path,
+            register_cover=test_subsystem.use_coverage,
+            address=field_set.address,
         ),
     )
 
@@ -210,10 +219,16 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
     if not testmain.has_tests and not testmain.has_xtests:
         return TestResult.skip(field_set.address, output_setting=test_subsystem.output)
 
+    coverage_config: GoCoverageConfig | None = None
+    if test_subsystem.use_coverage:
+        coverage_config = GoCoverageConfig(cover_mode=go_test_subsystem.coverage_mode)
+
     # Construct the build request for the package under test.
     maybe_test_pkg_build_request = await Get(
         FallibleBuildGoPackageRequest,
-        BuildGoPackageTargetRequest(field_set.address, for_tests=True),
+        BuildGoPackageTargetRequest(
+            field_set.address, for_tests=True, coverage_config=coverage_config
+        ),
     )
     if maybe_test_pkg_build_request.request is None:
         assert maybe_test_pkg_build_request.stderr is not None
@@ -222,6 +237,8 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
         )
     test_pkg_build_request = maybe_test_pkg_build_request.request
 
+    # TODO: Eventually support adding coverage to non-test packages. Those other packages will need to be
+    # added to `main_direct_deps` and to the coverage setup in the testmain.
     main_direct_deps = [test_pkg_build_request]
 
     if testmain.has_xtests:
@@ -251,17 +268,48 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
             direct_dependencies=tuple(direct_dependencies),
             minimum_go_version=pkg_analysis.minimum_go_version,
             embed_config=pkg_digest.xtest_embed_config,
+            coverage_config=coverage_config,
         )
         main_direct_deps.append(xtest_pkg_build_request)
 
+    # Generate coverage setup code for the test main if coverage is enabled.
+    #
+    # Note: Go coverage analysis is a form of codegen. It rewrites the Go source code at issue to include explicit
+    # references to "coverage variables" which contain the statement counts for coverage analysis. The test main
+    # generated for a Go test binary has to explicitly reference the coverage variables generated by this codegen and
+    # register them with the coverage runtime.
+    coverage_setup_digest = EMPTY_DIGEST
+    coverage_setup_files = []
+    if coverage_config is not None:
+        # Build the `main_direct_deps` when in coverage mode to obtain the "coverage variables" for those packages.
+        built_main_direct_deps = await MultiGet(
+            Get(BuiltGoPackage, BuildGoPackageRequest, build_req) for build_req in main_direct_deps
+        )
+        coverage_metadata = [
+            pkg.coverage_metadata for pkg in built_main_direct_deps if pkg.coverage_metadata
+        ]
+        coverage_setup_result = await Get(
+            GenerateCoverageSetupCodeResult,
+            GenerateCoverageSetupCodeRequest(
+                packages=FrozenOrderedSet(coverage_metadata),
+                cover_mode=go_test_subsystem.coverage_mode,
+            ),
+        )
+        coverage_setup_digest = coverage_setup_result.digest
+        coverage_setup_files = [GenerateCoverageSetupCodeResult.PATH]
+
+    testmain_input_digest = await Get(
+        Digest, MergeDigests([testmain.digest, coverage_setup_digest])
+    )
+
     # Generate the synthetic main package which imports the test and/or xtest packages.
     maybe_built_main_pkg = await Get(
         FallibleBuiltGoPackage,
         BuildGoPackageRequest(
             import_path="main",
-            digest=testmain.digest,
+            digest=testmain_input_digest,
             dir_path="",
-            go_file_names=(GeneratedTestMain.TEST_MAIN_FILE,),
+            go_file_names=(GeneratedTestMain.TEST_MAIN_FILE, *coverage_setup_files),
             s_file_names=(),
             direct_dependencies=tuple(main_direct_deps),
             minimum_go_version=pkg_analysis.minimum_go_version,
@@ -326,25 +374,48 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
         ProcessCacheScope.PER_SESSION if test_subsystem.force else ProcessCacheScope.SUCCESSFUL
     )
 
+    maybe_cover_args = []
+    maybe_cover_output_file = []
+    if test_subsystem.use_coverage:
+        maybe_cover_args = ["-test.coverprofile=cover.out"]
+        maybe_cover_output_file = ["cover.out"]
+
+    test_run_args = [
+        "./test_runner",
+        *transform_test_args(
+            go_test_subsystem.args,
+            field_set.timeout.calculate_from_global_options(test_subsystem),
+        ),
+        *maybe_cover_args,
+    ]
+
     result = await Get(
         FallibleProcessResult,
         Process(
-            [
-                "./test_runner",
-                *transform_test_args(
-                    go_test_subsystem.args,
-                    field_set.timeout.calculate_from_global_options(test_subsystem),
-                ),
-            ],
+            argv=test_run_args,
             env=extra_env,
             input_digest=test_input_digest,
             description=f"Run Go tests: {field_set.address}",
             cache_scope=cache_scope,
             working_directory=working_dir,
+            output_files=maybe_cover_output_file,
             level=LogLevel.DEBUG,
         ),
     )
-    return TestResult.from_fallible_process_result(result, field_set.address, test_subsystem.output)
+
+    coverage_data: GoCoverageData | None = None
+    if test_subsystem.use_coverage:
+        coverage_data = GoCoverageData(
+            coverage_digest=result.output_digest,
+            import_path=import_path,
+        )
+
+    return TestResult.from_fallible_process_result(
+        process_result=result,
+        address=field_set.address,
+        output_setting=test_subsystem.output,
+        coverage_data=coverage_data,
+    )
 
 
 @rule

src/python/pants/backend/go/subsystems/gotest.py

@@ -3,7 +3,11 @@
 
 from __future__ import annotations
 
-from pants.option.option_types import ArgsListOption
+from pathlib import PurePath
+
+from pants.backend.go.util_rules.coverage import GoCoverMode
+from pants.core.util_rules.distdir import DistDir
+from pants.option.option_types import ArgsListOption, EnumOption, StrOption
 from pants.option.subsystem import Subsystem
 from pants.util.strutil import softwrap
 
@@ -24,3 +28,38 @@ class GoTestSubsystem(Subsystem):
         ),
         passthrough=True,
     )
+
+    coverage_mode = EnumOption(
+        "--cover-mode",
+        default=GoCoverMode.SET,
+        help=softwrap(
+            """\
+            Coverage mode to use when running Go tests with coverage analysis enabled via --test-use-coverage.
+            Valid values are `set`, `count`, and `atomic`:\n
+            * `set`: bool: does this statement run?\n
+            * `count`: int: how many times does this statement run?\n
+            * `atomic`: int: count, but correct in multithreaded tests; significantly more expensive.\n
+            """
+        ),
+    )
+
+    _coverage_output_dir = StrOption(
+        default=str(PurePath("{distdir}", "coverage", "go", "{import_path_escaped}")),
+        advanced=True,
+        help=softwrap(
+            """
+            Path to write the Go coverage reports to. Must be relative to the build root.
+            `{distdir}` is replaced with the Pants `distdir`, and `{import_path_escaped}` is
+            replaced with the applicable package's import path but with slashes converted to
+            underscores.
+            """
+        ),
+    )
+
+    def coverage_output_dir(self, distdir: DistDir, import_path: str) -> PurePath:
+        import_path_escaped = import_path.replace("/", "_")
+        return PurePath(
+            self._coverage_output_dir.format(
+                distdir=distdir.relpath, import_path_escaped=import_path_escaped
+            )
+        )