[7.2.0] Implement FinalizeArtifacts (bazelbuild#21942)

If `--experimental_remote_output_service` is set, Bazel issues a
FinalizeArtifacts RPC to the output service after an action has been
executed.

Working towards bazelbuild#21630.

Closes bazelbuild#21759.

PiperOrigin-RevId: 618190373
Change-Id: I8ed981fcea7ebe146de0b1c48b1933392c86253f

Commit
bazelbuild@586bd13

Co-authored-by: Chi Wang <[email protected]>

Author: bazel-io

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

@@ -20,16 +20,19 @@
 
 import build.bazel.remote.execution.v2.DigestFunction;
 import com.google.devtools.build.lib.actions.Action;
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.cache.OutputMetadataStore;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.CleanRequest;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.CleanResponse;
+import com.google.devtools.build.lib.remote.BazelOutputServiceProto.FinalizeArtifactsRequest;
+import com.google.devtools.build.lib.remote.BazelOutputServiceProto.FinalizeArtifactsResponse;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.FinalizeBuildRequest;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.FinalizeBuildResponse;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.StageArtifactsRequest;
-import com.google.devtools.build.lib.remote.BazelOutputServiceProto.StageArtifactsRequest.Artifact;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.StageArtifactsResponse;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.StartBuildRequest;
 import com.google.devtools.build.lib.remote.BazelOutputServiceProto.StartBuildResponse;
@@ -64,6 +67,7 @@
 public class BazelOutputService implements OutputService {
 
   private final String outputBaseId;
+  private final Supplier<Path> execRootSupplier;
   private final Supplier<Path> outputPathSupplier;
   private final DigestFunction.Value digestFunction;
   private final RemoteOptions remoteOptions;
@@ -75,13 +79,15 @@ public class BazelOutputService implements OutputService {
 
   public BazelOutputService(
       Path outputBase,
+      Supplier<Path> execRootSupplier,
       Supplier<Path> outputPathSupplier,
       DigestFunction.Value digestFunction,
       RemoteOptions remoteOptions,
       boolean verboseFailures,
       RemoteRetrier retrier,
       ReferenceCountedChannel channel) {
     this.outputBaseId = DigestUtil.hashCodeToString(md5().hashString(outputBase.toString(), UTF_8));
+    this.execRootSupplier = execRootSupplier;
     this.outputPathSupplier = outputPathSupplier;
     this.digestFunction = digestFunction;
     this.remoteOptions = remoteOptions;
@@ -247,7 +253,7 @@ protected void stageArtifacts(List<FileMetadata> files) throws IOException, Inte
     request.setBuildId(buildId);
     for (var file : files) {
       request.addArtifacts(
-          Artifact.newBuilder()
+          StageArtifactsRequest.Artifact.newBuilder()
               .setPath(file.path().relativeTo(outputPath).toString())
               .setLocator(
                   Any.pack(FileArtifactLocator.newBuilder().setDigest(file.digest()).build()))
@@ -328,7 +334,63 @@ private FinalizeBuildResponse finalizeBuild(FinalizeBuildRequest request)
   @Override
   public void finalizeAction(Action action, OutputMetadataStore outputMetadataStore)
       throws IOException, EnvironmentalExecException, InterruptedException {
-    // TODO(chiwang): implement this
+    var execRoot = execRootSupplier.get();
+    var outputPath = outputPathSupplier.get();
+
+    var request = FinalizeArtifactsRequest.newBuilder();
+    request.setBuildId(buildId);
+    for (var output : action.getOutputs()) {
+      if (outputMetadataStore.artifactOmitted(output)) {
+        continue;
+      }
+
+      if (output.isTreeArtifact()) {
+        // TODO(chiwang): Use TreeArtifactLocator
+        var children = outputMetadataStore.getTreeArtifactChildren((SpecialArtifact) output);
+        for (var child : children) {
+          addArtifact(outputMetadataStore, execRoot, outputPath, request, child);
+        }
+      } else {
+        addArtifact(outputMetadataStore, execRoot, outputPath, request, output);
+      }
+    }
+
+    var unused = finalizeArtifacts(request.build());
+  }
+
+  private FinalizeArtifactsResponse finalizeArtifacts(FinalizeArtifactsRequest request)
+      throws IOException, InterruptedException {
+    return retrier.execute(
+        () ->
+            channel.withChannelBlocking(
+                channel -> {
+                  try {
+                    return BazelOutputServiceGrpc.newBlockingStub(channel)
+                        .finalizeArtifacts(request);
+                  } catch (StatusRuntimeException e) {
+                    throw new IOException(e);
+                  }
+                }));
+  }
+
+  private void addArtifact(
+      OutputMetadataStore outputMetadataStore,
+      Path execRoot,
+      Path outputPath,
+      FinalizeArtifactsRequest.Builder builder,
+      Artifact output)
+      throws IOException, InterruptedException {
+    checkState(!output.isTreeArtifact());
+    var metadata = outputMetadataStore.getOutputMetadata(output);
+    if (metadata.getType().isFile()) {
+      var digest = DigestUtil.buildDigest(metadata.getDigest(), metadata.getSize());
+      var path = execRoot.getRelative(output.getExecPath()).relativeTo(outputPath).toString();
+      builder.addArtifacts(
+          FinalizeArtifactsRequest.Artifact.newBuilder()
+              .setPath(path)
+              .setLocator(Any.pack(FileArtifactLocator.newBuilder().setDigest(digest).build()))
+              .build());
+    }
   }
 
   @Override

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

@@ -446,6 +446,7 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
       outputService =
           new BazelOutputService(
               env.getOutputBase(),
+              env::getExecRoot,
               () -> env.getDirectories().getOutputPath(env.getWorkspaceName()),
               digestUtil.getDigestFunction(),
               remoteOptions,