BLAKE3 digest function

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

@@ -0,0 +1,23 @@
+load("@rules_java//java:defs.bzl", "java_library")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    default_visibility = ["//src:__subpackages__"],
+)
+
+filegroup(
+    name = "srcs",
+    testonly = 0,
+    srcs = glob(["**"]),
+    visibility = ["//src:__subpackages__"],
+)
+
+java_library(
+    name = "hash",
+    srcs = glob(["*.java"]),
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib/jni",
+        "//third_party:guava",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+    ],
+)

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

@@ -0,0 +1,65 @@
+package com.google.devtools.build.lib.hash;
+
+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();
+  }
+}

src/main/java/com/google/devtools/build/lib/hash/Blake3Hasher.java

@@ -0,0 +1,184 @@
+package com.google.devtools.build.lib.hash;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.common.hash.Funnel;
+import com.google.common.hash.HashCode;
+import com.google.common.hash.Hasher;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.charset.Charset;
+
+final class Blake3Hasher implements Hasher {
+  // These constants match the native definitions in:
+  // https://github.com/BLAKE3-team/BLAKE3/blob/master/c/blake3.h
+  public static final int KEY_LEN = 32;
+  public static final int OUT_LEN = 32;
+
+  // To reduce the number of calls made via JNI, buffer up to this many bytes.
+  // If a call to "hash()" is made and less than this much data has been
+  // written, a single JNI call will be made that initializes, hashes, and
+  // cleans up the hasher, rather than making separate calls for each operation.
+  public static final int ONESHOT_THRESHOLD = 8 * 1024;
+
+  private ByteBuffer buffer;
+  private long hasher = -1;
+  private boolean isDone;
+
+  public Blake3Hasher() {
+    isDone = false;
+  }
+
+  private boolean isAllocated() {
+    return (hasher != -1);
+  }
+
+  private void initOnce() {
+    if (!isAllocated()) {
+      hasher = Blake3JNI.allocate_and_initialize_hasher();
+    }
+  }
+
+  private void resetBuffer(int minLength) {
+    int length = Math.max(ONESHOT_THRESHOLD, minLength);
+    if (buffer == null || buffer.capacity() < length) {
+      buffer = ByteBuffer.allocateDirect(length);
+      buffer.order(ByteOrder.nativeOrder());
+    }
+    buffer.clear();
+  }
+
+  private void update(byte[] data, int offset, int length) {
+    if (buffer == null) {
+      resetBuffer(length);
+    }
+
+    if (buffer.remaining() < length) {
+      initOnce();
+      Blake3JNI.blake3_hasher_update(hasher, buffer, 0, buffer.position());
+      resetBuffer(length);
+    }
+    buffer.put(data, offset, length);
+  }
+
+  private void update(byte[] data) {
+    update(data, 0, data.length);
+  }
+
+  private byte[] getOutput(int outputLength) throws IllegalArgumentException {
+    byte[] retByteArray = new byte[outputLength];
+
+    checkState(!isDone);
+    isDone = true;
+
+    if (!isAllocated() && buffer != null) {
+      Blake3JNI.blake3_hasher_oneshot(
+          hasher, buffer, buffer.position(), retByteArray, outputLength);
+    } else {
+      initOnce();
+      if (buffer != null && buffer.position() > 0) {
+        Blake3JNI.blake3_hasher_update(hasher, buffer, 0, buffer.position());
+      }
+      Blake3JNI.blake3_hasher_finalize_and_close(hasher, retByteArray, outputLength);
+      hasher = -1;
+    }
+
+    return retByteArray;
+  }
+
+  /* The following methods implement the {Hasher} interface. */
+
+  @CanIgnoreReturnValue
+  public Hasher putBytes(ByteBuffer b) {
+    buffer = b;
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putBytes(byte[] bytes, int off, int len) {
+    update(bytes, off, len);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putBytes(byte[] bytes) {
+    update(bytes, 0, bytes.length);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putByte(byte b) {
+    update(new byte[] {b});
+    return this;
+  }
+
+  public HashCode hash() {
+    return HashCode.fromBytes(getOutput(OUT_LEN));
+  }
+
+  @CanIgnoreReturnValue
+  public final Hasher putBoolean(boolean b) {
+    return putByte(b ? (byte) 1 : (byte) 0);
+  }
+
+  @CanIgnoreReturnValue
+  public final Hasher putDouble(double d) {
+    return putLong(Double.doubleToRawLongBits(d));
+  }
+
+  @CanIgnoreReturnValue
+  public final Hasher putFloat(float f) {
+    return putInt(Float.floatToRawIntBits(f));
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putUnencodedChars(CharSequence charSequence) {
+    for (int i = 0, len = charSequence.length(); i < len; i++) {
+      putChar(charSequence.charAt(i));
+    }
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putString(CharSequence charSequence, Charset charset) {
+    return putBytes(charSequence.toString().getBytes(charset));
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putShort(short s) {
+    putByte((byte) s);
+    putByte((byte) (s >>> 8));
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putInt(int i) {
+    putByte((byte) i);
+    putByte((byte) (i >>> 8));
+    putByte((byte) (i >>> 16));
+    putByte((byte) (i >>> 24));
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putLong(long l) {
+    for (int i = 0; i < 64; i += 8) {
+      putByte((byte) (l >>> i));
+    }
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public Hasher putChar(char c) {
+    putByte((byte) c);
+    putByte((byte) (c >>> 8));
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public <T extends Object> Hasher putObject(T instance, Funnel<? super T> funnel) {
+    funnel.funnel(instance, this);
+    return this;
+  }
+}

src/main/java/com/google/devtools/build/lib/hash/Blake3JNI.java

@@ -0,0 +1,37 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.hash;
+
+import com.google.devtools.build.lib.jni.JniLoader;
+import java.nio.ByteBuffer;
+
+final class Blake3JNI {
+  private Blake3JNI() {}
+
+  static {
+    JniLoader.loadJni();
+  }
+
+  public static final native long allocate_and_initialize_hasher();
+
+  public static final native void blake3_hasher_update(
+      long self, ByteBuffer inputBuf, int offset, int input_len);
+
+  public static final native void blake3_hasher_finalize_and_close(
+      long self, byte[] out, int out_len);
+
+  public static final native void blake3_hasher_oneshot(
+      long self, ByteBuffer inputBuf, int input_len, byte[] out, int out_len);
+}

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.");