Internal change

PiperOrigin-RevId: 675522967
Change-Id: I178fe567711ad8f551afd3b95567d4296f00fca0

src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java

@@ -372,9 +372,12 @@ ProfilerStartedEvent initProfiler(
                     env.getCommandId().toString(),
                     commandOptions.profilesToRetain);
             profile =
-                instrumentationOutputFactory.createLocalInstrumentationOutput(
+                instrumentationOutputFactory.createInstrumentationOutput(
                     profileName,
                     profilePath,
+                    options,
+                    commandOptions.redirectLocalInstrumentationOutputWrites,
+                    eventHandler,
                     /* convenienceName= */ profileName,
                     /* append= */ null,
                     /* internal= */ null);
@@ -386,11 +389,14 @@ ProfilerStartedEvent initProfiler(
                   : Format.JSON_TRACE_FILE_FORMAT;
           var profilePath = workspace.getWorkspace().getRelative(commandOptions.profilePath);
           profile =
-              instrumentationOutputFactory.createLocalInstrumentationOutput(
+              instrumentationOutputFactory.createInstrumentationOutput(
                   (format == Format.JSON_TRACE_FILE_COMPRESSED_FORMAT)
                       ? "command.profile.gz"
                       : "command.profile.json",
                   profilePath,
+                  options,
+                  commandOptions.redirectLocalInstrumentationOutputWrites,
+                  eventHandler,
                   /* convenienceName= */ null,
                   /* append= */ false,
                   /* internal= */ true);
@@ -1561,6 +1567,13 @@ public BlazeRuntime build() throws AbruptExitException {
       }
       ServerBuilder serverBuilder = new ServerBuilder();
       serverBuilder.addQueryOutputFormatters(OutputFormatters.getDefaultFormatters());
+      serverBuilder
+          .getInstrumentationOutputFactoryBuilder()
+          .setLocalInstrumentationOutputBuilderSupplier(LocalInstrumentationOutput.Builder::new);
+      serverBuilder
+          .getInstrumentationOutputFactoryBuilder()
+          .setBuildEventArtifactInstrumentationOutputBuilderSupplier(
+              BuildEventArtifactInstrumentationOutput.Builder::new);
       for (BlazeModule module : blazeModules) {
         module.serverInit(startupOptionsProvider, serverBuilder);
       }
@@ -1632,12 +1645,7 @@ public BlazeRuntime build() throws AbruptExitException {
               serverBuilder.getBuildEventArtifactUploaderMap(),
               serverBuilder.getAuthHeadersProvidersMap(),
               serverBuilder.getRepositoryRemoteExecutorFactory(),
-              new InstrumentationOutputFactory.Builder()
-                  .setLocalInstrumentationOutputBuilderSupplier(
-                      LocalInstrumentationOutput.Builder::new)
-                  .setBuildEventArtifactInstrumentationOutputBuilderSupplier(
-                      BuildEventArtifactInstrumentationOutput.Builder::new)
-                  .build());
+              serverBuilder.createInstrumentationOutputFactory());
       AutoProfiler.setClock(runtime.getClock());
       BugReport.setRuntime(runtime);
       return runtime;

src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java

@@ -596,4 +596,14 @@ public ProfilerTaskConverter() {
       super(ProfilerTask.class, "profiler task");
     }
   }
+
+  @Option(
+      name = "redirect_local_instrumentation_output_writes",
+      defaultValue = "false",
+      documentationCategory = OptionDocumentationCategory.LOGGING,
+      effectTags = {OptionEffectTag.BAZEL_MONITORING},
+      help =
+          "If true and supported, instrumentation output is redirected to be written locally on a"
+              + " different machine than where bazel is running on.")
+  public boolean redirectLocalInstrumentationOutputWrites;
 }

src/main/java/com/google/devtools/build/lib/runtime/InstrumentationOutputBuilder.java

@@ -13,6 +13,8 @@
 // limitations under the License.
 package com.google.devtools.build.lib.runtime;
 
+import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.common.options.OptionsProvider;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 
 /** Builds different {@link InstrumentationOutput} objects with correct input parameters. */
@@ -21,6 +23,18 @@ public interface InstrumentationOutputBuilder {
   @CanIgnoreReturnValue
   InstrumentationOutputBuilder setName(String name);
 
+  /** Sets the path to write the {@link InstrumentationOutput}. */
+  @CanIgnoreReturnValue
+  default InstrumentationOutputBuilder setPath(Path path) {
+    return this;
+  }
+
+  /** Provides the options necessary for building the {@link InstrumentationOutput}. */
+  @CanIgnoreReturnValue
+  default InstrumentationOutputBuilder setOptions(OptionsProvider options) {
+    return this;
+  }
+
   /** Builds the {@link InstrumentationOutput} object. */
   InstrumentationOutput build();
 }

src/main/java/com/google/devtools/build/lib/runtime/InstrumentationOutputFactory.java

@@ -16,7 +16,10 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
+import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.common.options.OptionsProvider;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
@@ -29,35 +32,69 @@ public final class InstrumentationOutputFactory {
   private final Supplier<BuildEventArtifactInstrumentationOutput.Builder>
       buildEventArtifactInstrumentationOutputBuilderSupplier;
 
+  @Nullable
+  final Supplier<InstrumentationOutputBuilder> redirectInstrumentationOutputBuilderSupplier;
+
   private InstrumentationOutputFactory(
       Supplier<LocalInstrumentationOutput.Builder> localInstrumentationOutputBuilderSupplier,
       Supplier<BuildEventArtifactInstrumentationOutput.Builder>
-          buildEventArtifactInstrumentationOutputBuilderSupplier) {
+          buildEventArtifactInstrumentationOutputBuilderSupplier,
+      @Nullable
+          Supplier<InstrumentationOutputBuilder> redirectInstrumentationOutputBuilderSupplier) {
     this.localInstrumentationOutputBuilderSupplier = localInstrumentationOutputBuilderSupplier;
     this.buildEventArtifactInstrumentationOutputBuilderSupplier =
         buildEventArtifactInstrumentationOutputBuilderSupplier;
+    this.redirectInstrumentationOutputBuilderSupplier =
+        redirectInstrumentationOutputBuilderSupplier;
   }
 
   /**
-   * Creates {@link LocalInstrumentationOutput} with given parameters.
+   * Creates {@link LocalInstrumentationOutput} or an {@link InstrumentationOutput} object
+   * redirecting outputs to be written on a different machine.
+   *
+   * <p>If {@link #redirectInstrumentationOutputBuilderSupplier} is not provided but {@code
+   * isRedirect} is {@code true}, this method will default to return {@link
+   * LocalInstrumentationOutput}.
    *
    * <p>{@code name} and {@code path} are required since they indicate what the output is and where
-   * it is stored.
+   * it is stored. {@code options} might also be necessary for the redirect {@link
+   * InstrumentationOutput}.
    *
-   * <p>When the name of the instrumentation output is complicated, an optional {@code
-   * convenienceName} parameter can be passed in so that a symlink pointing to the output with such
-   * a simpler name is created. See {@link LocalInstrumentationOutput.Builder#setConvenienceName}.
+   * <p>For {@link LocalInstrumentationOutput}, there are two additional considerations:
    *
-   * <p>User can also pass in the optional {@code append} and {@code internal} {@code Boolean}s to
-   * control how {@code path} creates the {@link OutputStream}. See {@link
-   * LocalInstrumentationOutput#createOutputStream()} for more details.
+   * <ul>
+   *   <li>When the name of the instrumentation output is complicated, an optional {@code
+   *       convenienceName} parameter can be passed in so that a symlink pointing to the output with
+   *       such a simpler name is created. See {@link
+   *       LocalInstrumentationOutput.Builder#setConvenienceName}.
+   *   <li>User can also pass in the optional {@code append} and {@code internal} {@code Boolean}s
+   *       to control how {@code path} creates the {@link OutputStream}. See {@link
+   *       LocalInstrumentationOutput#createOutputStream} for more details.
+   * </ul>
    */
-  public LocalInstrumentationOutput createLocalInstrumentationOutput(
+  public InstrumentationOutput createInstrumentationOutput(
       String name,
       Path path,
+      OptionsProvider options,
+      boolean isRedirect,
+      EventHandler eventHandler,
       @Nullable String convenienceName,
       @Nullable Boolean append,
       @Nullable Boolean internal) {
+    if (isRedirect) {
+      if (redirectInstrumentationOutputBuilderSupplier != null) {
+        return redirectInstrumentationOutputBuilderSupplier
+            .get()
+            .setName(name)
+            .setPath(path)
+            .setOptions(options)
+            .build();
+      }
+      eventHandler.handle(
+          Event.warn(
+              "Redirecting to write Instrumentation Output on a different machine is not"
+                  + " supported. Defaulting to writing output locally."));
+    }
     return localInstrumentationOutputBuilderSupplier
         .get()
         .setName(name)
@@ -86,6 +123,9 @@ public static class Builder {
     private Supplier<BuildEventArtifactInstrumentationOutput.Builder>
         buildEventArtifactInstrumentationOutputBuilderSupplier;
 
+    @Nullable
+    private Supplier<InstrumentationOutputBuilder> redirectInstrumentationOutputBuilderSupplier;
+
     @CanIgnoreReturnValue
     public Builder setLocalInstrumentationOutputBuilderSupplier(
         Supplier<LocalInstrumentationOutput.Builder> localInstrumentationOutputBuilderSupplier) {
@@ -102,14 +142,23 @@ public Builder setBuildEventArtifactInstrumentationOutputBuilderSupplier(
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder setRedirectInstrumentationOutputBuilderSupplier(
+        Supplier<InstrumentationOutputBuilder> redirectInstrumentationOutputBuilderSupplier) {
+      this.redirectInstrumentationOutputBuilderSupplier =
+          redirectInstrumentationOutputBuilderSupplier;
+      return this;
+    }
+
     public InstrumentationOutputFactory build() {
       return new InstrumentationOutputFactory(
           checkNotNull(
               localInstrumentationOutputBuilderSupplier,
               "Cannot create InstrumentationOutputFactory without localOutputBuilderSupplier"),
           checkNotNull(
               buildEventArtifactInstrumentationOutputBuilderSupplier,
-              "Cannot create InstrumentationOutputFactory without bepOutputBuilderSupplier"));
+              "Cannot create InstrumentationOutputFactory without bepOutputBuilderSupplier"),
+          redirectInstrumentationOutputBuilderSupplier);
     }
   }
 }

src/main/java/com/google/devtools/build/lib/runtime/LocalInstrumentationOutput.java

@@ -85,6 +85,7 @@ public Builder setName(String name) {
 
     /** Sets the path to the local {@link InstrumentationOutput}. */
     @CanIgnoreReturnValue
+    @Override
     public Builder setPath(Path path) {
       this.path = path;
       return this;

src/main/java/com/google/devtools/build/lib/runtime/ServerBuilder.java

@@ -40,6 +40,8 @@ public final class ServerBuilder {
   private final ImmutableMap.Builder<String, AuthHeadersProvider> authHeadersProvidersMap =
       ImmutableMap.builder();
   private RepositoryRemoteExecutorFactory repositoryRemoteExecutorFactory;
+  private final InstrumentationOutputFactory.Builder instrumentationOutputFactoryBuilder =
+      new InstrumentationOutputFactory.Builder();
 
   @VisibleForTesting
   public ServerBuilder() {}
@@ -185,4 +187,20 @@ public ServerBuilder addAuthHeadersProvider(
   public ImmutableMap<String, AuthHeadersProvider> getAuthHeadersProvidersMap() {
     return authHeadersProvidersMap.buildOrThrow();
   }
+
+  /**
+   * Returns the builder for {@link InstrumentationOutputFactory} so that suppliers for different
+   * types of {@link InstrumentationOutputBuilder} can be added.
+   */
+  public InstrumentationOutputFactory.Builder getInstrumentationOutputFactoryBuilder() {
+    return instrumentationOutputFactoryBuilder;
+  }
+
+  /**
+   * Creates the {@link InstrumentationOutputFactory} so that user can choose to create the {@link
+   * InstrumentationOutputBuilder} object.
+   */
+  public InstrumentationOutputFactory createInstrumentationOutputFactory() {
+    return instrumentationOutputFactoryBuilder.build();
+  }
 }

src/test/java/com/google/devtools/build/lib/runtime/InstrumentationOutputFactoryTest.java

@@ -18,13 +18,21 @@
 import static org.mockito.Mockito.mock;
 
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
+import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.EventKind;
+import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
+import com.google.devtools.common.options.OptionsProvider;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import com.google.testing.junit.testparameterinjector.TestParameter;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import java.util.ArrayList;
+import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
-@RunWith(JUnit4.class)
+@RunWith(TestParameterInjector.class)
 public class InstrumentationOutputFactoryTest {
   @Test
   public void testInstrumentationOutputFactory_cannotCreateFactorIfLocalSupplierUnset() {
@@ -53,23 +61,81 @@ public void testInstrumentationOutputFactory_cannotCreateFactorIfBepSupplierUnse
   }
 
   @Test
-  public void testInstrumentationOutputFactory_successfulCreateFactory() {
+  public void testInstrumentationOutputFactory_successfulFactoryCreation(
+      @TestParameter boolean injectRedirectOutputBuilderSupplier,
+      @TestParameter boolean createRedirectOutput) {
     InstrumentationOutputFactory.Builder factoryBuilder =
         new InstrumentationOutputFactory.Builder();
     factoryBuilder.setLocalInstrumentationOutputBuilderSupplier(
         LocalInstrumentationOutput.Builder::new);
     factoryBuilder.setBuildEventArtifactInstrumentationOutputBuilderSupplier(
         BuildEventArtifactInstrumentationOutput.Builder::new);
 
+    InstrumentationOutput fakeRedirectInstrumentationOutput = mock(InstrumentationOutput.class);
+    if (injectRedirectOutputBuilderSupplier) {
+      InstrumentationOutputBuilder fakeRedirectInstrumentationBuilder =
+          new InstrumentationOutputBuilder() {
+            @Override
+            @CanIgnoreReturnValue
+            public InstrumentationOutputBuilder setName(String name) {
+              return this;
+            }
+
+            @Override
+            public InstrumentationOutput build() {
+              return fakeRedirectInstrumentationOutput;
+            }
+          };
+
+      factoryBuilder.setRedirectInstrumentationOutputBuilderSupplier(
+          () -> fakeRedirectInstrumentationBuilder);
+    }
+
+    List<Event> warningEvents = new ArrayList<>();
+    ExtendedEventHandler eventHandler =
+        new ExtendedEventHandler() {
+          @Override
+          public void post(Postable obj) {}
+
+          @Override
+          public void handle(Event event) {
+            warningEvents.add(event);
+          }
+        };
+
     InstrumentationOutputFactory outputFactory = factoryBuilder.build();
-    assertThat(
-            outputFactory.createLocalInstrumentationOutput(
-                /* name= */ "local",
-                new InMemoryFileSystem(DigestHashFunction.SHA256).getPath("/file"),
-                /* convenienceName= */ null,
-                /* append= */ null,
-                /* internal= */ null))
-        .isNotNull();
+    var instrumentationOutput =
+        outputFactory.createInstrumentationOutput(
+            /* name= */ "local",
+            new InMemoryFileSystem(DigestHashFunction.SHA256).getPath("/file"),
+            mock(OptionsProvider.class),
+            createRedirectOutput,
+            eventHandler,
+            /* convenienceName= */ null,
+            /* append= */ null,
+            /* internal= */ null);
+
+    // Only when redirectOutputBuilderSupplier is provided to the factory, and we intend to create a
+    // RedirectOutputBuilder object, we expect a non-LocalInstrumentationOutput to be created. In
+    // all other scenarios, a LocalInstrumentationOutput is returned.
+    if (createRedirectOutput && injectRedirectOutputBuilderSupplier) {
+      assertThat(instrumentationOutput).isEqualTo(fakeRedirectInstrumentationOutput);
+    } else {
+      assertThat(instrumentationOutput).isInstanceOf(LocalInstrumentationOutput.class);
+    }
+
+    // When user wants to create a redirectOutputBuilder object but its builder supplier is not
+    // provided, eventHandler should post a warning event.
+    if (createRedirectOutput && !injectRedirectOutputBuilderSupplier) {
+      assertThat(warningEvents)
+          .containsExactly(
+              Event.of(
+                  EventKind.WARNING,
+                  "Redirecting to write Instrumentation Output on a different machine is not"
+                      + " supported. Defaulting to writing output locally."));
+    } else {
+      assertThat(warningEvents).isEmpty();
+    }
     assertThat(
             outputFactory.createBuildEventArtifactInstrumentationOutput(
                 /* name= */ "bep", mock(BuildEventArtifactUploader.class)))