Async {AbstractFileWrite,TemplateExpansion}Action

While the contents of these files have to be evaluated locally, we have
code to write the files to remote storage, which can block the current
thread for at least one network round-trip time. These blocks show up in
the profiler, even though they don't consume any local CPU, making it
more difficult to interpret the results.

This change may improve build performance in cases where Bazel has to
write a lot of such files and also has different actions which could
execute in parallel. However, I did not measure any improvements in the
builds I tested.

Progress on #6394.

PiperOrigin-RevId: 242104204

src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java

@@ -16,18 +16,25 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.AbstractAction;
+import com.google.devtools.build.lib.actions.ActionContinuationOrResult;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.ActionResult;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
+import com.google.devtools.build.lib.actions.SpawnContinuation;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.OutputStream;
+import javax.annotation.Nullable;
 
 /**
  * Abstract Action to write to a file.
@@ -56,6 +63,80 @@ public boolean makeExecutable() {
     return makeExecutable;
   }
 
+  @Override
+  public final ActionContinuationOrResult beginExecution(
+      ActionExecutionContext actionExecutionContext)
+      throws ActionExecutionException, InterruptedException {
+    try {
+      DeterministicWriter deterministicWriter;
+      try {
+        deterministicWriter = newDeterministicWriter(actionExecutionContext);
+      } catch (IOException e) {
+        // Message is a bit misleading but is good enough for the end user.
+        throw new EnvironmentalExecException(
+            "failed to create file '"
+                + getPrimaryOutput().prettyPrint()
+                + "' due to I/O error: "
+                + e.getMessage(),
+            e);
+      }
+
+      FileWriteActionContext context = getStrategy(actionExecutionContext);
+      SpawnContinuation first =
+          new SpawnContinuation() {
+            @Override
+            public ListenableFuture<?> getFuture() {
+              return null;
+            }
+
+            @Override
+            public SpawnContinuation execute() throws ExecException, InterruptedException {
+              return context.beginWriteOutputToFile(
+                  AbstractFileWriteAction.this,
+                  actionExecutionContext,
+                  deterministicWriter,
+                  makeExecutable,
+                  isRemotable());
+            }
+          };
+
+      return new ActionContinuationOrResult() {
+        private SpawnContinuation spawnContinuation = first;
+
+        @Nullable
+        @Override
+        public ListenableFuture<?> getFuture() {
+          return spawnContinuation.getFuture();
+        }
+
+        @Override
+        public ActionContinuationOrResult execute()
+            throws ActionExecutionException, InterruptedException {
+          SpawnContinuation next;
+          try {
+            next = spawnContinuation.execute();
+            if (!next.isDone()) {
+              spawnContinuation = next;
+              return this;
+            }
+          } catch (ExecException e) {
+            throw e.toActionExecutionException(
+                "Writing file for rule '" + Label.print(getOwner().getLabel()) + "'",
+                actionExecutionContext.getVerboseFailures(),
+                AbstractFileWriteAction.this);
+          }
+          afterWrite(actionExecutionContext);
+          return ActionContinuationOrResult.of(ActionResult.create(next.get()));
+        }
+      }.execute();
+    } catch (ExecException e) {
+      throw e.toActionExecutionException(
+          "Writing file for rule '" + Label.print(getOwner().getLabel()) + "'",
+          actionExecutionContext.getVerboseFailures(),
+          this);
+    }
+  }
+
   @Override
   public final ActionResult execute(ActionExecutionContext actionExecutionContext)
       throws ActionExecutionException, InterruptedException {
@@ -69,15 +150,20 @@ public final ActionResult execute(ActionExecutionContext actionExecutionContext)
         throw new EnvironmentalExecException("failed to create file '"
             + getPrimaryOutput().prettyPrint() + "' due to I/O error: " + e.getMessage(), e);
       }
-      actionResult =
-          ActionResult.create(
-              getStrategy(actionExecutionContext)
-                  .writeOutputToFile(
-                      this,
-                      actionExecutionContext,
-                      deterministicWriter,
-                      makeExecutable,
-                      isRemotable()));
+      FileWriteActionContext context = getStrategy(actionExecutionContext);
+      try (SilentCloseable c =
+          Profiler.instance()
+              .profile(ProfilerTask.INFO, "FileWriteActionContext.writeOutputToFile")) {
+        actionResult =
+            ActionResult.create(
+                SpawnContinuation.completeBlocking(
+                    context.beginWriteOutputToFile(
+                        this,
+                        actionExecutionContext,
+                        deterministicWriter,
+                        makeExecutable,
+                        isRemotable())));
+      }
     } catch (ExecException e) {
       throw e.toActionExecutionException(
           "Writing file for rule '" + Label.print(getOwner().getLabel()) + "'",

src/main/java/com/google/devtools/build/lib/analysis/actions/FileWriteActionContext.java

@@ -17,9 +17,8 @@
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ExecException;
-import com.google.devtools.build.lib.actions.SpawnResult;
+import com.google.devtools.build.lib.actions.SpawnContinuation;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction.DeterministicWriter;
-import java.util.List;
 
 /**
  * The action context for {@link AbstractFileWriteAction} instances (technically instances of
@@ -31,7 +30,7 @@ public interface FileWriteActionContext extends ActionContext {
    * Writes the output created by the {@link DeterministicWriter} to the sole output of the given
    * action.
    */
-  List<SpawnResult> writeOutputToFile(
+  SpawnContinuation beginWriteOutputToFile(
       AbstractAction action,
       ActionExecutionContext actionExecutionContext,
       DeterministicWriter deterministicWriter,

src/main/java/com/google/devtools/build/lib/analysis/actions/LocalTemplateExpansionStrategy.java

@@ -14,19 +14,17 @@
 
 package com.google.devtools.build.lib.analysis.actions;
 
-import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionStrategy;
-import com.google.devtools.build.lib.actions.SpawnResult;
+import com.google.devtools.build.lib.actions.SpawnContinuation;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction.DeterministicWriter;
 import com.google.devtools.build.lib.util.StringUtilities;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
-import java.util.List;
 
 /** Strategy to perform tempate expansion locally */
 @ExecutionStrategy(
@@ -39,7 +37,7 @@ public class LocalTemplateExpansionStrategy implements TemplateExpansionContext
   public static LocalTemplateExpansionStrategy INSTANCE = new LocalTemplateExpansionStrategy();
 
   @Override
-  public List<SpawnResult> expandTemplate(
+  public SpawnContinuation expandTemplate(
       TemplateExpansionAction action, ActionExecutionContext ctx)
       throws ExecException, InterruptedException {
     try {
@@ -51,13 +49,12 @@ public void writeOutputFile(OutputStream out) throws IOException {
               out.write(expandedTemplate.getBytes(StandardCharsets.UTF_8));
             }
           };
-      ctx.getContext(FileWriteActionContext.class)
-          .writeOutputToFile(
+      return ctx.getContext(FileWriteActionContext.class)
+          .beginWriteOutputToFile(
               action, ctx, deterministicWriter, action.makeExecutable(), /*isRemotable=*/ true);
     } catch (IOException e) {
       throw new EnvironmentalExecException("IOException during template expansion", e);
     }
-    return ImmutableList.of();
   }
 
   /**

src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java

@@ -18,7 +18,9 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.AbstractAction;
+import com.google.devtools.build.lib.actions.ActionContinuationOrResult;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
@@ -27,6 +29,7 @@
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.ExecException;
+import com.google.devtools.build.lib.actions.SpawnContinuation;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
@@ -35,6 +38,7 @@
 import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
+import javax.annotation.Nullable;
 
 /** Action to expand a template and write the expanded content to a file. */
 @AutoCodec
@@ -124,13 +128,65 @@ public String getSkylarkContent() throws IOException {
     return getFileContents();
   }
 
+  @Override
+  public final ActionContinuationOrResult beginExecution(
+      ActionExecutionContext actionExecutionContext)
+      throws ActionExecutionException, InterruptedException {
+    SpawnContinuation first =
+        new SpawnContinuation() {
+          @Override
+          public ListenableFuture<?> getFuture() {
+            return null;
+          }
+
+          @Override
+          public SpawnContinuation execute() throws ExecException, InterruptedException {
+            TemplateExpansionContext expansionContext =
+                actionExecutionContext.getContext(TemplateExpansionContext.class);
+            return expansionContext.expandTemplate(
+                TemplateExpansionAction.this, actionExecutionContext);
+          }
+        };
+
+    return new ActionContinuationOrResult() {
+      private SpawnContinuation spawnContinuation = first;
+
+      @Nullable
+      @Override
+      public ListenableFuture<?> getFuture() {
+        return spawnContinuation.getFuture();
+      }
+
+      @Override
+      public ActionContinuationOrResult execute()
+          throws ActionExecutionException, InterruptedException {
+        SpawnContinuation next;
+        try {
+          next = spawnContinuation.execute();
+          if (!next.isDone()) {
+            spawnContinuation = next;
+            return this;
+          }
+        } catch (ExecException e) {
+          throw e.toActionExecutionException(
+              "Error expanding template '" + Label.print(getOwner().getLabel()) + "'",
+              actionExecutionContext.getVerboseFailures(),
+              TemplateExpansionAction.this);
+        }
+        return ActionContinuationOrResult.of(ActionResult.create(next.get()));
+      }
+    }.execute();
+  }
+
   @Override
   public final ActionResult execute(ActionExecutionContext actionExecutionContext)
       throws ActionExecutionException, InterruptedException {
     TemplateExpansionContext expansionContext =
         actionExecutionContext.getContext(TemplateExpansionContext.class);
     try {
-      return ActionResult.create(expansionContext.expandTemplate(this, actionExecutionContext));
+      return ActionResult.create(
+          SpawnContinuation.completeBlocking(
+              expansionContext.expandTemplate(this, actionExecutionContext)));
     } catch (ExecException e) {
       throw e.toActionExecutionException(
           "Error expanding template '" + Label.print(getOwner().getLabel()) + "'",

src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionContext.java

@@ -17,11 +17,10 @@
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ExecException;
-import com.google.devtools.build.lib.actions.SpawnResult;
-import java.util.List;
+import com.google.devtools.build.lib.actions.SpawnContinuation;
 
 /** The action context for {@link TemplateExpansionAction} instances */
 public interface TemplateExpansionContext extends ActionContext {
-  List<SpawnResult> expandTemplate(TemplateExpansionAction action, ActionExecutionContext ctx)
+  SpawnContinuation expandTemplate(TemplateExpansionAction action, ActionExecutionContext ctx)
       throws ExecException, InterruptedException;
 }

src/main/java/com/google/devtools/build/lib/exec/FileWriteStrategy.java

@@ -14,15 +14,14 @@
 
 package com.google.devtools.build.lib.exec;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.actions.AbstractAction;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionStrategy;
 import com.google.devtools.build.lib.actions.RunningActionEvent;
-import com.google.devtools.build.lib.actions.SpawnResult;
+import com.google.devtools.build.lib.actions.SpawnContinuation;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction.DeterministicWriter;
 import com.google.devtools.build.lib.analysis.actions.FileWriteActionContext;
@@ -31,7 +30,6 @@
 import java.io.BufferedOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
-import java.util.List;
 import java.util.logging.Logger;
 
 /**
@@ -46,11 +44,12 @@ public FileWriteStrategy() {
   }
 
   @Override
-  public List<SpawnResult> writeOutputToFile(
+  public SpawnContinuation beginWriteOutputToFile(
       AbstractAction action,
       ActionExecutionContext actionExecutionContext,
       DeterministicWriter deterministicWriter,
-      boolean makeExecutable, boolean isRemotable)
+      boolean makeExecutable,
+      boolean isRemotable)
       throws ExecException {
     actionExecutionContext.getEventHandler().post(new RunningActionEvent(action, "local"));
     // TODO(ulfjack): Consider acquiring local resources here before trying to write the file.
@@ -72,6 +71,6 @@ public List<SpawnResult> writeOutputToFile(
         throw new EnvironmentalExecException("IOException during file write", e);
       }
     }
-    return ImmutableList.of();
+    return SpawnContinuation.immediate();
   }
 }