Set the digest_function field as part of all relevant gRPC requests

EdSchouten · copybara-github · commit 0a8380bec9dc · 2023-03-28T10:54:07.000-07:00
In the following PR we extended most of the REv2 *Request messages to take an explicit digest function: bazelbuild/remote-apis#235 This eliminates the ambiguity between digest functions that have the same hash length (e.g., MD5 and MURMUR3, SHA256 and SHA256TREE). This change extends the REv2 client in Bazel to set the digest_function field explicitly, so that the server does not need to use any heuristics to pick a digest function when processing requests. As we are now going to call DigestUtil.getDigestFunction() far more frequently, alter DigestUtil to precompute the value upon construction. This change was tested by letting Bazel forward all traffic through an instance of bb_clientd that was patched up to require the use of `digest_function`. Closes bazelbuild#17772. PiperOrigin-RevId: 520074622 Change-Id: Ie92dc368c502b9bc3fddef56a5eb0a238760f673
diff --git a/src/main/java/com/google/devtools/build/lib/remote/GrpcCacheClient.java b/src/main/java/com/google/devtools/build/lib/remote/GrpcCacheClient.java
@@ -117,6 +117,7 @@ private int computeMaxMissingBlobsDigestsPerMessage() {
     final int overhead =
         FindMissingBlobsRequest.newBuilder()
             .setInstanceName(options.remoteInstanceName)
+            .setDigestFunction(digestUtil.getDigestFunction())
             .build()
             .getSerializedSize();
     final int tagSize =
@@ -185,7 +186,9 @@ public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(
     }
     // Need to potentially split the digests into multiple requests.
     FindMissingBlobsRequest.Builder requestBuilder =
-        FindMissingBlobsRequest.newBuilder().setInstanceName(options.remoteInstanceName);
+        FindMissingBlobsRequest.newBuilder()
+            .setInstanceName(options.remoteInstanceName)
+            .setDigestFunction(digestUtil.getDigestFunction());
     List<ListenableFuture<FindMissingBlobsResponse>> getMissingDigestCalls = new ArrayList<>();
     for (Digest digest : digests) {
       requestBuilder.addBlobDigests(digest);
@@ -260,6 +263,7 @@ public ListenableFuture<CachedActionResult> downloadActionResult(
     GetActionResultRequest request =
         GetActionResultRequest.newBuilder()
             .setInstanceName(options.remoteInstanceName)
+            .setDigestFunction(digestUtil.getDigestFunction())
             .setActionDigest(actionKey.getDigest())
             .setInlineStderr(inlineOutErr)
             .setInlineStdout(inlineOutErr)
@@ -289,6 +293,7 @@ public ListenableFuture<Void> uploadActionResult(
                                         .updateActionResult(
                                             UpdateActionResultRequest.newBuilder()
                                                 .setInstanceName(options.remoteInstanceName)
+                                                .setDigestFunction(digestUtil.getDigestFunction())
                                                 .setActionDigest(actionKey.getDigest())
                                                 .setActionResult(actionResult)
                                                 .build())),
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java
@@ -1486,6 +1486,7 @@ public RemoteActionResult executeRemotely(
     ExecuteRequest.Builder requestBuilder =
         ExecuteRequest.newBuilder()
             .setInstanceName(remoteOptions.remoteInstanceName)
+            .setDigestFunction(digestUtil.getDigestFunction())
             .setActionDigest(action.getActionKey().getDigest())
             .setSkipCacheLookup(!acceptCachedResult);
     if (remoteOptions.remoteResultCachePriority != 0) {
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteRepositoryRemoteExecutor.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteRepositoryRemoteExecutor.java
@@ -167,6 +167,7 @@ public ExecutionResult execute(
             ExecuteRequest.newBuilder()
                 .setActionDigest(actionDigest)
                 .setInstanceName(remoteInstanceName)
+                .setDigestFunction(digestUtil.getDigestFunction())
                 .setSkipCacheLookup(!acceptCached)
                 .build();
 
diff --git a/src/main/java/com/google/devtools/build/lib/remote/util/DigestUtil.java b/src/main/java/com/google/devtools/build/lib/remote/util/DigestUtil.java
@@ -35,14 +35,15 @@
 public class DigestUtil {
   private final XattrProvider xattrProvider;
   private final DigestHashFunction hashFn;
+  private final DigestFunction.Value digestFunction;
 
   public DigestUtil(XattrProvider xattrProvider, DigestHashFunction hashFn) {
     this.xattrProvider = xattrProvider;
     this.hashFn = hashFn;
+    this.digestFunction = getDigestFunctionFromHashFunction(hashFn);
   }
 
-  /** Returns the currently used digest function. */
-  public DigestFunction.Value getDigestFunction() {
+  private static DigestFunction.Value getDigestFunctionFromHashFunction(DigestHashFunction hashFn) {
     for (String name : hashFn.getNames()) {
       try {
         return DigestFunction.Value.valueOf(name);
@@ -53,6 +54,11 @@ public DigestFunction.Value getDigestFunction() {
     return DigestFunction.Value.UNKNOWN;
   }
 
+  /** Returns the currently used digest function. */
+  public DigestFunction.Value getDigestFunction() {
+    return digestFunction;
+  }
+
   public Digest compute(byte[] blob) {
     return buildDigest(hashFn.getHashFunction().hashBytes(blob).toString(), blob.length);
   }
diff --git a/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java b/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java
@@ -30,6 +30,7 @@
 import build.bazel.remote.execution.v2.Command;
 import build.bazel.remote.execution.v2.ContentAddressableStorageGrpc.ContentAddressableStorageImplBase;
 import build.bazel.remote.execution.v2.Digest;
+import build.bazel.remote.execution.v2.DigestFunction;
 import build.bazel.remote.execution.v2.Directory;
 import build.bazel.remote.execution.v2.DirectoryNode;
 import build.bazel.remote.execution.v2.FileNode;
@@ -798,6 +799,7 @@ public void updateActionResult(
             assertThat(request)
                 .isEqualTo(
                     UpdateActionResultRequest.newBuilder()
+                        .setDigestFunction(DigestFunction.Value.SHA256)
                         .setActionDigest(fooDigest)
                         .setActionResult(result)
                         .build());
