Add support for path mapping (#3915)

Path mapping improves disk/remote cache hit ratio by automatically
removing the configuration-specific path segments (e.g.
`k8-fastbuild-ST-12345678`) from action command lines before staging
them for execution.

This commit makes the stdlib, compilepkg (assuming no cgo) and
gentestmain actions compatible with path mapping, which is mostly
automatic except that `GOROOT` now needs to be passed in via a flag
rather than an environment variable.

Can be tested via `--experimental_output_paths=strip` with Bazel
last_green.

See https://www.youtube.com/watch?v=Et1rjb7ixUU for more information on
path mapping.

go/private/actions/compilepkg.bzl

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load("//go/private:common.bzl", "GO_TOOLCHAIN_LABEL")
+load("//go/private:common.bzl", "GO_TOOLCHAIN_LABEL", "SUPPORTS_PATH_MAPPING_REQUIREMENT")
 load(
     "//go/private:mode.bzl",
     "link_mode_args",
@@ -86,9 +86,8 @@ def emit_compilepkg(
               [archive.data.export_file for archive in archives] +
               go.sdk.tools + go.sdk.headers + go.stdlib.libs)
     outputs = [out_lib, out_export]
-    env = go.env
 
-    args = go.builder_args(go, "compilepkg")
+    args = go.builder_args(go, "compilepkg", use_path_mapping = True)
     args.add_all(sources, before_each = "-src")
     args.add_all(embedsrcs, before_each = "-embedsrc", expand_directories = False)
     args.add_all(
@@ -155,7 +154,15 @@ def emit_compilepkg(
     args.add("-gcflags", quote_opts(gc_flags))
     args.add("-asmflags", quote_opts(asm_flags))
 
-    env = go.env
+    # cgo and the linker action don't support path mapping yet
+    # TODO: Remove the second condition after https://github.com/bazelbuild/bazel/pull/21921.
+    if cgo or "local" in go._ctx.attr.tags:
+        # cgo doesn't support path mapping yet
+        env = go.env
+        execution_requirements = {}
+    else:
+        env = go.env_for_path_mapping
+        execution_requirements = SUPPORTS_PATH_MAPPING_REQUIREMENT
     if cgo:
         inputs.extend(cgo_inputs.to_list())  # OPT: don't expand depset
         inputs.extend(go.crosstool)
@@ -183,6 +190,7 @@ def emit_compilepkg(
         mnemonic = "GoCompilePkgExternal" if is_external_pkg else "GoCompilePkg",
         executable = go.toolchain._builder,
         arguments = [args],
-        env = go.env,
+        env = env,
         toolchain = GO_TOOLCHAIN_LABEL,
+        execution_requirements = execution_requirements,
     )

go/private/actions/stdlib.bzl

@@ -16,6 +16,7 @@ load(
     "//go/private:common.bzl",
     "COVERAGE_OPTIONS_DENYLIST",
     "GO_TOOLCHAIN_LABEL",
+    "SUPPORTS_PATH_MAPPING_REQUIREMENT",
 )
 load(
     "//go/private:mode.bzl",
@@ -120,10 +121,16 @@ def _sdk_stdlib(go):
         root_file = go.sdk.root_file,
     )
 
+def _dirname(file):
+    return file.dirname
+
 def _build_stdlib(go):
     pkg = go.declare_directory(go, path = "pkg")
-    args = go.builder_args(go, "stdlib")
-    args.add("-out", pkg.dirname)
+    args = go.builder_args(go, "stdlib", use_path_mapping = True)
+
+    # Use a file rather than pkg.dirname as the latter is just a string and thus
+    # not subject to path mapping.
+    args.add_all("-out", [pkg], map_each = _dirname, expand_directories = False)
     if go.mode.race:
         args.add("-race")
     args.add("-package", "std")
@@ -152,6 +159,7 @@ def _build_stdlib(go):
         arguments = [args],
         env = _build_env(go),
         toolchain = GO_TOOLCHAIN_LABEL,
+        execution_requirements = SUPPORTS_PATH_MAPPING_REQUIREMENT,
     )
     return GoStdLib(
         _list_json = _build_stdlib_list_json(go),

go/private/common.bzl

@@ -271,3 +271,7 @@ RULES_GO_STDLIB_PREFIX = RULES_GO_REPO_NAME + "//stdlib:"
 
 # TODO: Remove the "and" once the rules_go repo itself uses Bzlmod.
 RULES_GO_IS_BZLMOD_REPO = _RULES_GO_RAW_REPO_NAME.lstrip("@") != "io_bazel_rules_go" and _RULES_GO_RAW_REPO_NAME.lstrip("@")
+
+# Marks an action as supporting path mapping (--experimental_output_paths=strip).
+# See https://www.youtube.com/watch?v=Et1rjb7ixUU for more details.
+SUPPORTS_PATH_MAPPING_REQUIREMENT = {"supports-path-mapping": "1"}

go/private/context.bzl

@@ -160,13 +160,28 @@ def _new_args(go):
     # TODO(jayconrod): print warning.
     return go.builder_args(go)
 
-def _builder_args(go, command = None):
+def _dirname(file):
+    return file.dirname
+
+def _builder_args(go, command = None, use_path_mapping = False):
     args = go.actions.args()
     args.use_param_file("-param=%s")
     args.set_param_file_format("shell")
     if command:
         args.add(command)
     args.add("-sdk", go.sdk.root_file.dirname)
+
+    # Path mapping can't map the values of environment variables, so we need to pass GOROOT to the
+    # action via an argument instead.
+    if use_path_mapping:
+        if go.stdlib:
+            goroot_file = go.stdlib.root_file
+        else:
+            goroot_file = go.sdk_root
+
+        # Use a file rather than goroot as the latter is just a string and thus
+        # not subject to path mapping.
+        args.add_all("-goroot", [goroot_file], map_each = _dirname, expand_directories = False)
     args.add("-installsuffix", installsuffix(go.mode))
     args.add_joined("-tags", go.tags, join_with = ",")
     return args
@@ -487,6 +502,12 @@ def go_context(ctx, attr = None):
         "GOTOOLCHAIN": "local",
     }
 
+    # Path mapping can't map the values of environment variables, so we pass GOROOT to the action
+    # via an argument instead in builder_args. We need to drop it from the environment to get cache
+    # hits across different configurations since the stdlib path typically contains a Bazel
+    # configuration segment.
+    env_for_path_mapping = {k: v for k, v in env.items() if k != "GOROOT"}
+
     # The level of support is determined by the platform constraints in
     # //go/constraints/amd64.
     # See https://go.dev/wiki/MinimumRequirements#amd64
@@ -572,6 +593,7 @@ def go_context(ctx, attr = None):
         coverage_enabled = ctx.configuration.coverage_enabled,
         coverage_instrumented = ctx.coverage_instrumented(),
         env = env,
+        env_for_path_mapping = env_for_path_mapping,
         tags = tags,
         stamp = mode.stamp,
         label = ctx.label,

go/private/rules/test.bzl

@@ -20,6 +20,7 @@ load(
     "//go/private:common.bzl",
     "GO_TOOLCHAIN",
     "GO_TOOLCHAIN_LABEL",
+    "SUPPORTS_PATH_MAPPING_REQUIREMENT",
     "as_list",
     "asm_exts",
     "cgo_exts",
@@ -94,7 +95,7 @@ def _go_test_impl(ctx):
         run_dir = repo_relative_rundir
 
     main_go = go.declare_file(go, path = "testmain.go")
-    arguments = go.builder_args(go, "gentestmain")
+    arguments = go.builder_args(go, "gentestmain", use_path_mapping = True)
     arguments.add("-output", main_go)
     if go.coverage_enabled:
         if go.mode.race:
@@ -114,14 +115,16 @@ def _go_test_impl(ctx):
     )
     arguments.add("-pkgname", internal_source.library.importpath)
     arguments.add_all(go_srcs, before_each = "-src", format_each = "l=%s")
+
     ctx.actions.run(
         inputs = go_srcs,
         outputs = [main_go],
         mnemonic = "GoTestGenTest",
         executable = go.toolchain._builder,
         arguments = [arguments],
         toolchain = GO_TOOLCHAIN_LABEL,
-        env = go.env,
+        env = go.env_for_path_mapping,
+        execution_requirements = SUPPORTS_PATH_MAPPING_REQUIREMENT,
     )
 
     test_gc_linkopts = gc_linkopts(ctx)

go/tools/builders/compilepkg.go

@@ -88,7 +88,7 @@ func compilePkg(args []string) error {
 	if err := fs.Parse(args); err != nil {
 		return err
 	}
-	if err := goenv.checkFlags(); err != nil {
+	if err := goenv.checkFlagsAndSetGoroot(); err != nil {
 		return err
 	}
 	if importPath == "" {

go/tools/builders/env.go

@@ -48,6 +48,9 @@ type env struct {
 	// platform. This may be different than GOROOT.
 	sdk string
 
+	// goroot is set as the value of GOROOT if non-empty.
+	goroot string
+
 	// installSuffix is the name of the directory below GOROOT/pkg that contains
 	// the .a files for the standard library we should build against.
 	// For example, linux_amd64_race.
@@ -67,19 +70,26 @@ type env struct {
 func envFlags(flags *flag.FlagSet) *env {
 	env := &env{}
 	flags.StringVar(&env.sdk, "sdk", "", "Path to the Go SDK.")
+	flags.StringVar(&env.goroot, "goroot", "", "The value to set for GOROOT.")
 	flags.Var(&tagFlag{}, "tags", "List of build tags considered true.")
 	flags.StringVar(&env.installSuffix, "installsuffix", "", "Standard library under GOROOT/pkg")
 	flags.BoolVar(&env.verbose, "v", false, "Whether subprocess command lines should be printed")
 	flags.BoolVar(&env.shouldPreserveWorkDir, "work", false, "if true, the temporary work directory will be preserved")
 	return env
 }
 
-// checkFlags checks whether env flags were set to valid values. checkFlags
-// should be called after parsing flags.
-func (e *env) checkFlags() error {
+// checkFlagsAndSetGoroot checks whether env flags were set to valid values and sets GOROOT.
+// checkFlagsAndSetGoroot should be called after parsing flags.
+func (e *env) checkFlagsAndSetGoroot() error {
 	if e.sdk == "" {
 		return errors.New("-sdk was not set")
 	}
+	if e.goroot != "" {
+		err := os.Setenv("GOROOT", e.goroot)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 

go/tools/builders/generate_test_main.go

@@ -272,7 +272,7 @@ func genTestMain(args []string) error {
 	if err := flags.Parse(args); err != nil {
 		return err
 	}
-	if err := goenv.checkFlags(); err != nil {
+	if err := goenv.checkFlagsAndSetGoroot(); err != nil {
 		return err
 	}
 	// Process import args

go/tools/builders/info.go

@@ -35,7 +35,7 @@ func run(args []string) error {
 	if err := flags.Parse(args); err != nil {
 		return err
 	}
-	if err := goenv.checkFlags(); err != nil {
+	if err := goenv.checkFlagsAndSetGoroot(); err != nil {
 		return err
 	}
 	os.Setenv("GO111MODULE", "off")

go/tools/builders/link.go

@@ -54,7 +54,7 @@ func link(args []string) error {
 	if err := flags.Parse(builderArgs); err != nil {
 		return err
 	}
-	if err := goenv.checkFlags(); err != nil {
+	if err := goenv.checkFlagsAndSetGoroot(); err != nil {
 		return err
 	}
 

go/tools/builders/stdlib.go

@@ -41,7 +41,7 @@ func stdlib(args []string) error {
 	if err := flags.Parse(args); err != nil {
 		return err
 	}
-	if err := goenv.checkFlags(); err != nil {
+	if err := goenv.checkFlagsAndSetGoroot(); err != nil {
 		return err
 	}
 	goroot := os.Getenv("GOROOT")

go/tools/builders/stdliblist.go

@@ -206,7 +206,7 @@ func stdliblist(args []string) error {
 	if err := flags.Parse(args); err != nil {
 		return err
 	}
-	if err := goenv.checkFlags(); err != nil {
+	if err := goenv.checkFlagsAndSetGoroot(); err != nil {
 		return err
 	}