Allow overriding the hostname and instance name in bytestream:// URIs

src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java

@@ -63,16 +63,11 @@ class ByteStreamBuildEventArtifactUploader extends AbstractReferenceCounted
   ByteStreamBuildEventArtifactUploader(
       ByteStreamUploader uploader,
       MissingDigestsFinder missingDigestsFinder,
-      String remoteServerName,
+      String remoteServerInstanceName,
       String buildRequestId,
       String commandId,
-      @Nullable String remoteInstanceName,
       int maxUploadThreads) {
     this.uploader = Preconditions.checkNotNull(uploader);
-    String remoteServerInstanceName = Preconditions.checkNotNull(remoteServerName);
-    if (!Strings.isNullOrEmpty(remoteInstanceName)) {
-      remoteServerInstanceName += "/" + remoteInstanceName;
-    }
     this.buildRequestId = buildRequestId;
     this.commandId = commandId;
     this.remoteServerInstanceName = remoteServerInstanceName;

src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploaderFactory.java

@@ -27,36 +27,32 @@ class ByteStreamBuildEventArtifactUploaderFactory implements
     BuildEventArtifactUploaderFactory {
 
   private final ByteStreamUploader uploader;
-  private final String remoteServerName;
+  private final String remoteServerInstanceName;
   private final String buildRequestId;
   private final String commandId;
   private final MissingDigestsFinder missingDigestsFinder;
-  @Nullable private final String remoteInstanceName;
 
   ByteStreamBuildEventArtifactUploaderFactory(
       ByteStreamUploader uploader,
       MissingDigestsFinder missingDigestsFinder,
-      String remoteServerName,
+      String remoteServerInstanceName,
       String buildRequestId,
-      String commandId,
-      @Nullable String remoteInstanceName) {
+      String commandId) {
     this.uploader = uploader;
     this.missingDigestsFinder = missingDigestsFinder;
-    this.remoteServerName = remoteServerName;
+    this.remoteServerInstanceName = remoteServerInstanceName;
     this.buildRequestId = buildRequestId;
     this.commandId = commandId;
-    this.remoteInstanceName = remoteInstanceName;
   }
 
   @Override
   public BuildEventArtifactUploader create(CommandEnvironment env) {
     return new ByteStreamBuildEventArtifactUploader(
         uploader.retain(),
         missingDigestsFinder,
-        remoteServerName,
+        remoteServerInstanceName,
         buildRequestId,
         commandId,
-        remoteInstanceName,
         env.getOptions().getOptions(RemoteOptions.class).buildEventUploadMaxThreads);
   }
 }

src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java

@@ -482,6 +482,14 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
       }
     }
 
+    String remoteBytestreamUriPrefix = remoteOptions.remoteBytestreamUriPrefix;
+    if (Strings.isNullOrEmpty(remoteBytestreamUriPrefix)) {
+      remoteBytestreamUriPrefix = cacheChannel.authority();
+      if (!Strings.isNullOrEmpty(remoteOptions.remoteInstanceName)) {
+        remoteBytestreamUriPrefix += "/" + remoteOptions.remoteInstanceName;
+      }
+    }
+
     ByteStreamUploader uploader =
         new ByteStreamUploader(
             remoteOptions.remoteInstanceName,
@@ -504,10 +512,9 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
         new ByteStreamBuildEventArtifactUploaderFactory(
             uploader,
             cacheClient,
-            cacheChannel.authority(),
+            remoteBytestreamUriPrefix,
             buildRequestId,
-            invocationId,
-            remoteOptions.remoteInstanceName));
+            invocationId));
 
     if (enableRemoteExecution) {
       RemoteExecutionClient remoteExecutor;

src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java

@@ -181,6 +181,19 @@ public final class RemoteOptions extends OptionsBase {
               + " the unit is omitted, the value is interpreted as seconds.")
   public Duration remoteTimeout;
 
+  @Option(
+      name = "remote_bytestream_uri_prefix",
+      defaultValue = "null",
+      documentationCategory = OptionDocumentationCategory.REMOTE,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help =
+          "The hostname and instance name to be used in bytestream:// URIs that are written into "
+              + "build event streams. This option can be set when builds are performed using a "
+              + "proxy, which causes the values of --remote_executor and --remote_instance_name "
+              + "to no longer correspond to the canonical name of the remote execution service. "
+              + "When not set, it will default to \"${hostname}/${instance_name}\".")
+  public String remoteBytestreamUriPrefix;
+
   /** Returns the specified duration. Assumes seconds if unitless. */
   public static class RemoteTimeoutConverter implements Converter<Duration> {
     private static final Pattern UNITLESS_REGEX = Pattern.compile("^[0-9]+$");

src/test/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploaderTest.java

@@ -361,10 +361,9 @@ private ByteStreamBuildEventArtifactUploader newArtifactUploader(
     return new ByteStreamBuildEventArtifactUploader(
         uploader,
         missingDigestsFinder,
-        "localhost",
+        "localhost/instance",
         "none",
         "none",
-        "instance",
         /* maxUploadThreads= */ 100);
   }
 

src/test/shell/bazel/remote/remote_execution_test.sh

@@ -1440,6 +1440,41 @@ EOF
   expect_log "uri:.*bytestream://localhost"
 }
 
+function test_bytestream_uri_prefix() {
+  # Test that when --remote_bytestream_uri_prefix is set, bytestream://
+  # URIs do not contain the hostname that's part of --remote_executor.
+  # They should use a fixed value instead.
+  mkdir -p a
+  cat > a/success.sh <<'EOF'
+#!/bin/sh
+exit 0
+EOF
+  chmod 755 a/success.sh
+  cat > a/BUILD <<'EOF'
+sh_test(
+  name = "success_test",
+  srcs = ["success.sh"],
+)
+
+genrule(
+  name = "foo",
+  srcs = [],
+  outs = ["foo.txt"],
+  cmd = "echo \"foo\" > \"$@\"",
+)
+EOF
+
+  bazel test \
+    --remote_executor=grpc://localhost:${worker_port} \
+    --remote_download_minimal \
+    --remote_bytestream_uri_prefix=example.com/my-instance-name \
+    --build_event_text_file=$TEST_log \
+    //a:foo //a:success_test || fail "Failed to test //a:foo //a:success_test"
+
+  expect_not_log 'uri:.*file://'
+  expect_log "uri:.*bytestream://example.com/my-instance-name/blobs"
+}
+
 # This test is derivative of test_bep_output_groups in
 # build_event_stream_test.sh, which verifies that successful output groups'
 # artifacts appear in BEP when a top-level target fails to build.