Add BLAKE3 hasher to vfs

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

@@ -21,6 +21,7 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 import build.bazel.remote.execution.v2.Digest;
+import build.bazel.remote.execution.v2.DigestFunction;
 import com.google.bytestream.ByteStreamGrpc;
 import com.google.bytestream.ByteStreamGrpc.ByteStreamFutureStub;
 import com.google.bytestream.ByteStreamGrpc.ByteStreamStub;
@@ -69,6 +70,7 @@ final class ByteStreamUploader {
   private final CallCredentialsProvider callCredentialsProvider;
   private final long callTimeoutSecs;
   private final RemoteRetrier retrier;
+  private final DigestFunction.Value digestFunction;
 
   @Nullable private final Semaphore openedFilePermits;
 
@@ -89,14 +91,16 @@ final class ByteStreamUploader {
       CallCredentialsProvider callCredentialsProvider,
       long callTimeoutSecs,
       RemoteRetrier retrier,
-      int maximumOpenFiles) {
+      int maximumOpenFiles,
+      DigestFunction.Value digestFunction) {
     checkArgument(callTimeoutSecs > 0, "callTimeoutSecs must be gt 0.");
     this.instanceName = instanceName;
     this.channel = channel;
     this.callCredentialsProvider = callCredentialsProvider;
     this.callTimeoutSecs = callTimeoutSecs;
     this.retrier = retrier;
     this.openedFilePermits = maximumOpenFiles != -1 ? new Semaphore(maximumOpenFiles) : null;
+    this.digestFunction = digestFunction;
   }
 
   @VisibleForTesting
@@ -175,11 +179,34 @@ public ListenableFuture<Void> uploadBlobAsync(
         MoreExecutors.directExecutor());
   }
 
-  private static String buildUploadResourceName(
+  private boolean isOldStyleDigestFunction() {
+    // Old-style digest functions (SHA256, etc) are distinguishable by the length
+    // of their hash alone and do not require extra specification, but newer
+    // digest functions (which may have the same length hashes as the older
+    // functions!) must be explicitly specified in the upload resource name.
+    return digestFunction.getNumber() <= 7;
+  }
+
+  private String buildUploadResourceName(
       String instanceName, UUID uuid, Digest digest, boolean compressed) {
-    String template =
-        compressed ? "uploads/%s/compressed-blobs/zstd/%s/%d" : "uploads/%s/blobs/%s/%d";
-    String resourceName = format(template, uuid, digest.getHash(), digest.getSizeBytes());
+
+    String resourceName;
+
+    if (isOldStyleDigestFunction()) {
+      String template =
+          compressed ? "uploads/%s/compressed-blobs/zstd/%s/%d" : "uploads/%s/blobs/%s/%d";
+      resourceName = format(template, uuid, digest.getHash(), digest.getSizeBytes());
+    } else {
+      String template =
+          compressed ? "uploads/%s/compressed-blobs/zstd/%s/%s/%d" : "uploads/%s/blobs/%s/%s/%d";
+      resourceName =
+          format(
+              template,
+              uuid,
+              digestFunction.getValueDescriptor().getName().toLowerCase(),
+              digest.getHash(),
+              digest.getSizeBytes());
+    }
     if (!Strings.isNullOrEmpty(instanceName)) {
       resourceName = instanceName + "/" + resourceName;
     }

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

@@ -107,7 +107,8 @@ public GrpcCacheClient(
             callCredentialsProvider,
             options.remoteTimeout.getSeconds(),
             retrier,
-            options.maximumOpenFiles);
+            options.maximumOpenFiles,
+            digestUtil.getDigestFunction());
     maxMissingBlobsDigestsPerMessage = computeMaxMissingBlobsDigestsPerMessage();
     Preconditions.checkState(
         maxMissingBlobsDigestsPerMessage > 0, "Error: gRPC message size too small.");

src/main/java/com/google/devtools/build/lib/util/Fingerprint.java

@@ -14,26 +14,25 @@
 
 package com.google.devtools.build.lib.util;
 
-import static java.nio.charset.StandardCharsets.UTF_8;
 
-import com.google.common.io.ByteStreams;
+import com.google.common.hash.Funnels;
+import com.google.common.hash.HashFunction;
+import com.google.common.hash.Hasher;
 import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
-import java.security.DigestException;
-import java.security.DigestOutputStream;
-import java.security.MessageDigest;
+import java.nio.charset.StandardCharsets;
 import java.util.Collection;
 import java.util.Map;
 import java.util.UUID;
 import javax.annotation.Nullable;
 
 /**
- * Simplified wrapper for using {@link MessageDigest} to generate fingerprints.
+ * Simplified wrapper for using {@link HashFunction} to generate fingerprints.
  *
  * <p>A fingerprint is a cryptographic hash of a message that encodes the representation of an
  * object. Two objects of the same type have the same fingerprint if and only if they are equal.
@@ -64,18 +63,14 @@ public final class Fingerprint {
 
   // Make novel use of a CodedOutputStream, which is good at efficiently serializing data. By
   // flushing at the end of each digest we can continue to use the stream.
-  private final CodedOutputStream codedOut;
-  private final MessageDigest messageDigest;
+  private final HashFunction hashFunction;
+  private Hasher hasher;
+  private CodedOutputStream codedOut;
 
   /** Creates and initializes a new instance. */
   public Fingerprint(DigestHashFunction digestFunction) {
-    messageDigest = digestFunction.cloneOrCreateMessageDigest();
-    // This is a lot of indirection, but CodedOutputStream does a reasonable job of converting
-    // strings to bytes without creating a whole bunch of garbage, which pays off.
-    codedOut =
-        CodedOutputStream.newInstance(
-            new DigestOutputStream(ByteStreams.nullOutputStream(), messageDigest),
-            /*bufferSize=*/ 1024);
+    hashFunction = digestFunction.getHashFunction();
+    reset();
   }
 
   public Fingerprint() {
@@ -84,6 +79,14 @@ public Fingerprint() {
     this(DigestHashFunction.SHA256);
   }
 
+  private void reset() {
+    hasher = hashFunction.newHasher();
+    // This is a lot of indirection, but CodedOutputStream does a reasonable job of converting
+    // strings to bytes without creating a whole bunch of garbage, which pays off.
+    codedOut =
+        CodedOutputStream.newInstance(Funnels.asOutputStream(hasher), /* bufferSize= */ 1024);
+  }
+
   /**
    * Completes the hash computation by doing final operations and resets the underlying state,
    * allowing this instance to be used again.
@@ -97,7 +100,9 @@ public byte[] digestAndReset() {
     } catch (IOException e) {
       throw new IllegalStateException("failed to flush", e);
     }
-    return messageDigest.digest();
+    byte[] digest = hasher.hash().asBytes();
+    reset();
+    return digest;
   }
 
   /**
@@ -112,12 +117,12 @@ public byte[] digestAndReset() {
   public void digestAndReset(byte[] buf, int offset, int len) {
     try {
       codedOut.flush();
-      messageDigest.digest(buf, offset, len);
     } catch (IOException e) {
       throw new IllegalStateException("failed to flush", e);
-    } catch (DigestException e) {
-      throw new IllegalStateException("failed to digest", e);
     }
+
+    hasher.hash().writeBytesTo(buf, offset, len);
+    reset();
   }
 
   /** Same as {@link #digestAndReset()}, except returns the digest in hex string form. */
@@ -343,7 +348,9 @@ public static String getHexDigest(String input) {
     // use the value from DigestHashFunction.getDefault(). However, this gets called during class
     // loading in a few places, before setDefault() has been called, so these call-sites should be
     // removed before this can be done safely.
-    return hexDigest(
-        DigestHashFunction.SHA256.cloneOrCreateMessageDigest().digest(input.getBytes(UTF_8)));
+    return DigestHashFunction.SHA256
+        .getHashFunction()
+        .hashString(input, StandardCharsets.UTF_8)
+        .toString();
   }
 }

src/main/java/com/google/devtools/build/lib/vfs/BUILD

@@ -66,6 +66,7 @@ java_library(
         ":pathfragment",
         "//src/main/java/com/google/devtools/build/lib/clock",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
+        "//src/main/java/com/google/devtools/build/lib/jni",
         "//src/main/java/com/google/devtools/build/lib/io:file_symlink_exception",
         "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
@@ -77,6 +78,7 @@ java_library(
         "//third_party:guava",
         "//third_party:jsr305",
         "//third_party/protobuf:protobuf_java",
+        "@maven//:com_google_errorprone_error_prone_annotations",
     ],
 )
 

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

@@ -0,0 +1,65 @@
+package com.google.devtools.build.lib.vfs;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkPositionIndexes;
+
+import com.google.common.hash.Funnel;
+import com.google.common.hash.HashCode;
+import com.google.common.hash.HashFunction;
+import com.google.common.hash.Hasher;
+import com.google.errorprone.annotations.Immutable;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+
+@Immutable
+public final class Blake3HashFunction implements HashFunction {
+  public int bits() {
+    return 256;
+  }
+
+  public Hasher newHasher() {
+    return new Blake3Hasher();
+  }
+
+  /* The following methods implement the {Hasher} interface. */
+
+  public <T extends Object> HashCode hashObject(T instance, Funnel<? super T> funnel) {
+    return newHasher().putObject(instance, funnel).hash();
+  }
+
+  public HashCode hashUnencodedChars(CharSequence input) {
+    int len = input.length();
+    return newHasher(len * 2).putUnencodedChars(input).hash();
+  }
+
+  public HashCode hashString(CharSequence input, Charset charset) {
+    return newHasher().putString(input, charset).hash();
+  }
+
+  public HashCode hashInt(int input) {
+    return newHasher(4).putInt(input).hash();
+  }
+
+  public HashCode hashLong(long input) {
+    return newHasher(8).putLong(input).hash();
+  }
+
+  public HashCode hashBytes(byte[] input) {
+    return hashBytes(input, 0, input.length);
+  }
+
+  public HashCode hashBytes(byte[] input, int off, int len) {
+    checkPositionIndexes(off, off + len, input.length);
+    return newHasher(len).putBytes(input, off, len).hash();
+  }
+
+  public HashCode hashBytes(ByteBuffer input) {
+    return newHasher(input.remaining()).putBytes(input).hash();
+  }
+
+  public Hasher newHasher(int expectedInputSize) {
+    checkArgument(
+        expectedInputSize >= 0, "expectedInputSize must be >= 0 but was %s", expectedInputSize);
+    return newHasher();
+  }
+}