[7.2.0] Implement XattrProvider which uses BatchStat to getFastDigest (bazelbuild#21946)

If `--experimental_remote_output_service` is set, a `XattrProvider`
which uses `BatchStat` RPC to get the digest of a file will be used for
Bazel to get digests for outputs.

Working towards bazelbuild#21630.

Closes bazelbuild#21812.

PiperOrigin-RevId: 619539349
Change-Id: I71364e78c45c37e834c71e120a109c8de602db77

Commit
bazelbuild@30aa5d0

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

Author: bazel-io

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

@@ -26,6 +26,8 @@
 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.BatchStatRequest;
+import com.google.devtools.build.lib.remote.BazelOutputServiceProto.BatchStatResponse;
 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;
@@ -53,9 +55,12 @@
 import com.google.devtools.build.lib.vfs.OutputService;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.vfs.XattrProvider;
+import com.google.devtools.build.lib.vfs.XattrProvider.DelegatingXattrProvider;
 import com.google.protobuf.Any;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
@@ -76,6 +81,7 @@ public class BazelOutputService implements OutputService {
   private final ReferenceCountedChannel channel;
 
   @Nullable private String buildId;
+  @Nullable private PathFragment outputPathTarget;
 
   public BazelOutputService(
       Path outputBase,
@@ -223,7 +229,8 @@ public ModifiedFileSet startBuild(
                   .build()));
     }
 
-    var outputPathTarget = constructOutputPathTarget(outputPathPrefix, response);
+    checkState(outputPathTarget == null, "outputPathTarget must be null");
+    outputPathTarget = constructOutputPathTarget(outputPathPrefix, response);
     prepareOutputPath(outputPath, outputPathTarget);
 
     if (finalizeActions) {
@@ -314,6 +321,7 @@ public void finalizeBuild(boolean buildSuccessful)
                   .build()));
     } finally {
       this.buildId = null;
+      this.outputPathTarget = null;
     }
   }
 
@@ -432,4 +440,72 @@ private CleanResponse clean(CleanRequest request) throws IOException, Interrupte
                   }
                 }));
   }
+
+  private BatchStatResponse batchStat(BatchStatRequest request)
+      throws IOException, InterruptedException {
+    return retrier.execute(
+        () ->
+            channel.withChannelBlocking(
+                channel -> {
+                  try {
+                    return BazelOutputServiceGrpc.newBlockingStub(channel).batchStat(request);
+                  } catch (StatusRuntimeException e) {
+                    throw new IOException(e);
+                  }
+                }));
+  }
+
+  @Override
+  public XattrProvider getXattrProvider(XattrProvider delegate) {
+    return new DelegatingXattrProvider(delegate) {
+      @Nullable
+      @Override
+      public byte[] getFastDigest(Path path) throws IOException {
+        var outputPath = outputPathSupplier.get();
+        var buildId = checkNotNull(BazelOutputService.this.buildId);
+        var outputPathTarget = checkNotNull(BazelOutputService.this.outputPathTarget);
+
+        String pathString = null;
+        if (path.startsWith(outputPath)) {
+          pathString = path.relativeTo(outputPath).toString();
+        } else if (path.startsWith(outputPathTarget)) {
+          pathString = path.asFragment().relativeTo(outputPathTarget).toString();
+        }
+        if (pathString == null) {
+          return super.getFastDigest(path);
+        }
+
+        var request =
+            BatchStatRequest.newBuilder().setBuildId(buildId).addPaths(pathString).build();
+        BatchStatResponse response;
+        try {
+          response = batchStat(request);
+        } catch (InterruptedException e) {
+          throw new IOException(e);
+        }
+
+        if (response.getResponsesCount() != 1) {
+          throw new IOException(
+              String.format(
+                  "BatchStat failed: expect 1 response, got %s", response.getResponsesCount()));
+        }
+
+        var statResponse = response.getResponses(0);
+        if (!statResponse.hasStat()) {
+          throw new FileNotFoundException(path.getPathString());
+        }
+
+        var stat = statResponse.getStat();
+        if (stat.hasFile()) {
+          var file = stat.getFile();
+          if (file.hasLocator()) {
+            var locator = file.getLocator().unpack(FileArtifactLocator.class);
+            return DigestUtil.toBinaryDigest(locator.getDigest());
+          }
+        }
+
+        return null;
+      }
+    };
+  }
 }

src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java

@@ -217,7 +217,7 @@ public void injectTree(SpecialArtifact output, TreeArtifactValue tree) {
   @Nullable private ActionCompletedReceiver completionReceiver;
 
   private final AtomicReference<ActionExecutionStatusReporter> statusReporterRef;
-  private OutputService outputService;
+  @Nullable private OutputService outputService;
   private boolean finalizeActions;
   private boolean rewindingEnabled;
   private final Supplier<ImmutableList<Root>> sourceRootSupplier;
@@ -283,7 +283,7 @@ void prepareForExecution(
       OptionsProvider options,
       ActionCacheChecker actionCacheChecker,
       ActionOutputDirectoryHelper outputDirectoryHelper,
-      OutputService outputService,
+      @Nullable OutputService outputService,
       boolean trackIncrementalState) {
     this.reporter = checkNotNull(reporter);
     this.executorEngine = checkNotNull(executor);
@@ -361,6 +361,10 @@ OutputPermissions getOutputPermissions() {
   }
 
   XattrProvider getXattrProvider() {
+    if (outputService != null) {
+      return checkNotNull(outputService.getXattrProvider(syscallCache));
+    }
+
     return syscallCache;
   }
 
@@ -369,14 +373,15 @@ FileSystem createActionFileSystem(
       String relativeOutputPath,
       ActionInputMap inputArtifactData,
       Iterable<Artifact> outputArtifacts) {
-    return outputService.createActionFileSystem(
-        executorEngine.getFileSystem(),
-        executorEngine.getExecRoot().asFragment(),
-        relativeOutputPath,
-        sourceRootSupplier.get(),
-        inputArtifactData,
-        outputArtifacts,
-        rewindingEnabled);
+    return checkNotNull(outputService)
+        .createActionFileSystem(
+            executorEngine.getFileSystem(),
+            executorEngine.getExecRoot().asFragment(),
+            relativeOutputPath,
+            sourceRootSupplier.get(),
+            inputArtifactData,
+            outputArtifacts,
+            rewindingEnabled);
   }
 
   private void updateActionFileSystemContext(
@@ -385,8 +390,8 @@ private void updateActionFileSystemContext(
       Environment env,
       MetadataInjector metadataInjector,
       ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
-    outputService.updateActionFileSystemContext(
-        action, actionFileSystem, env, metadataInjector, filesets);
+    checkNotNull(outputService)
+        .updateActionFileSystemContext(action, actionFileSystem, env, metadataInjector, filesets);
   }
 
   void executionOver() {

src/main/java/com/google/devtools/build/lib/vfs/OutputService.java

@@ -241,4 +241,8 @@ default ArtifactPathResolver createPathResolverForArtifactValues(
   default BulkDeleter bulkDeleter() {
     return null;
   }
+
+  default XattrProvider getXattrProvider(XattrProvider delegate) {
+    return delegate;
+  }
 }

src/main/java/com/google/devtools/build/lib/vfs/XattrProvider.java

@@ -38,4 +38,29 @@ default byte[] getxattr(Path path, String xattrName) throws IOException {
   default byte[] getxattr(Path path, String xattrName, Symlinks followSymlinks) throws IOException {
     return path.getxattr(xattrName, followSymlinks);
   }
+
+  /** Delegates to another {@link XattrProvider}. */
+  class DelegatingXattrProvider implements XattrProvider {
+    private final XattrProvider delegate;
+
+    public DelegatingXattrProvider(XattrProvider delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public byte[] getFastDigest(Path path) throws IOException {
+      return delegate.getFastDigest(path);
+    }
+
+    @Override
+    public byte[] getxattr(Path path, String xattrName) throws IOException {
+      return delegate.getxattr(path, xattrName);
+    }
+
+    @Override
+    public byte[] getxattr(Path path, String xattrName, Symlinks followSymlinks)
+        throws IOException {
+      return delegate.getxattr(path, xattrName, followSymlinks);
+    }
+  }
 }

src/test/java/com/google/devtools/build/lib/skyframe/OutputsInvalidationIntegrationTest.java

@@ -57,6 +57,7 @@ public void prepareOutputServiceMock()
     when(outputService.getFilesSystemName()).thenReturn("fileSystemName");
     when(outputService.startBuild(any(), any(), anyBoolean()))
         .thenReturn(ModifiedFileSet.EVERYTHING_MODIFIED);
+    when(outputService.getXattrProvider(any())).thenAnswer(i -> i.getArgument(0));
   }
 
   @Override