Pass digest to Chunker construction when available.

* Refactor Chunker constructor to a builder to reduce constructor overload.
* Pass digest into this where we have it
* Redo ensureInputsPresent to not lose the missing digests during processing so we can pass them to the Chunker constructor.

RELNOTES: None
PiperOrigin-RevId: 207297915

4 files changed, 39 insertions, 34 deletions

diff --git a/src/test/java/com/google/devtools/build/lib/remote/ByteStreamUploaderTest.java b/src/test/java/com/google/devtools/build/lib/remote/ByteStreamUploaderTest.java
index 06ad3914ed..0e142de538 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/ByteStreamUploaderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/ByteStreamUploaderTest.java
@@ -143,7 +143,7 @@ public class ByteStreamUploaderTest {
     byte[] blob = new byte[CHUNK_SIZE * 2 + 1];
     new Random().nextBytes(blob);
 
-    Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker = Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
 
     serviceRegistry.addService(new ByteStreamImplBase() {
           @Override
@@ -221,7 +221,8 @@ public class ByteStreamUploaderTest {
       int blobSize = rand.nextInt(CHUNK_SIZE * 10) + CHUNK_SIZE;
       byte[] blob = new byte[blobSize];
       rand.nextBytes(blob);
-      Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+      Chunker chunker =
+          Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
       builders.add(chunker);
       blobsByHash.put(chunker.digest().getHash(), blob);
     }
@@ -250,7 +251,8 @@ public class ByteStreamUploaderTest {
     List<Chunker> builders = new ArrayList<>(toUpload.size());
     Map<String, Integer> uploadsFailed = new HashMap<>();
     for (String s : toUpload) {
-      Chunker chunker = new Chunker(s.getBytes(UTF_8), /* chunkSize=*/ 3, DIGEST_UTIL);
+      Chunker chunker =
+          Chunker.builder(DIGEST_UTIL).setInput(s.getBytes(UTF_8)).setChunkSize(3).build();
       builders.add(chunker);
       uploadsFailed.put(chunker.digest().getHash(), 0);
     }
@@ -339,7 +341,7 @@ public class ByteStreamUploaderTest {
         new ReferenceCountedChannel(channel), null, 3, retrier);
 
     byte[] blob = new byte[CHUNK_SIZE * 10];
-    Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker = Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
 
     AtomicInteger numWriteCalls = new AtomicInteger();
     CountDownLatch blocker = new CountDownLatch(1);
@@ -402,7 +404,7 @@ public class ByteStreamUploaderTest {
         new ReferenceCountedChannel(channel), null, 3, retrier);
 
     byte[] blob = new byte[CHUNK_SIZE];
-    Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker = Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
 
     serviceRegistry.addService(new ByteStreamImplBase() {
       @Override
@@ -460,10 +462,12 @@ public class ByteStreamUploaderTest {
     serviceRegistry.addService(service);
 
     byte[] blob1 = new byte[CHUNK_SIZE];
-    Chunker chunker1 = new Chunker(blob1, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker1 =
+        Chunker.builder(DIGEST_UTIL).setInput(blob1).setChunkSize(CHUNK_SIZE).build();
 
     byte[] blob2 = new byte[CHUNK_SIZE + 1];
-    Chunker chunker2 = new Chunker(blob2, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker2 =
+        Chunker.builder(DIGEST_UTIL).setInput(blob2).setChunkSize(CHUNK_SIZE).build();
 
     ListenableFuture<Void> f1 = uploader.uploadBlobAsync(chunker1, true);
     ListenableFuture<Void> f2 = uploader.uploadBlobAsync(chunker2, true);
@@ -508,7 +512,7 @@ public class ByteStreamUploaderTest {
     assertThat(retryService.isShutdown()).isTrue();
 
     byte[] blob = new byte[1];
-    Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker = Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
     try {
       uploader.uploadBlob(chunker, true);
       fail("Should have thrown an exception.");
@@ -553,7 +557,7 @@ public class ByteStreamUploaderTest {
     });
 
     byte[] blob = new byte[1];
-    Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker = Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
 
     uploader.uploadBlob(chunker, true);
 
@@ -585,7 +589,7 @@ public class ByteStreamUploaderTest {
     });
 
     byte[] blob = new byte[1];
-    Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker = Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
 
     try {
       uploader.uploadBlob(chunker, true);
@@ -608,7 +612,7 @@ public class ByteStreamUploaderTest {
     byte[] blob = new byte[CHUNK_SIZE * 2 + 1];
     new Random().nextBytes(blob);
 
-    Chunker chunker = new Chunker(blob, CHUNK_SIZE, DIGEST_UTIL);
+    Chunker chunker = Chunker.builder(DIGEST_UTIL).setInput(blob).setChunkSize(CHUNK_SIZE).build();
 
     AtomicInteger numUploads = new AtomicInteger();
     serviceRegistry.addService(new ByteStreamImplBase() {
diff --git a/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java b/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java
index 8857abbd85..a157260b68 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/ChunkerTest.java
@@ -47,7 +47,7 @@ public class ChunkerTest {
     rand.nextBytes(expectedData);
     Digest expectedDigest = digestUtil.compute(expectedData);
 
-    Chunker chunker = new Chunker(expectedData, 10, digestUtil);
+    Chunker chunker = Chunker.builder(digestUtil).setInput(expectedData).setChunkSize(10).build();
 
     ByteArrayOutputStream actualData = new ByteArrayOutputStream();
 
@@ -74,13 +74,13 @@ public class ChunkerTest {
 
     assertThat(chunker.hasNext()).isFalse();
 
-    assertThat(expectedData).isEqualTo(actualData.toByteArray());
+    assertThat(actualData.toByteArray()).isEqualTo(expectedData);
   }
 
   @Test
   public void nextShouldThrowIfNoMoreData() throws IOException {
     byte[] data = new byte[10];
-    Chunker chunker = new Chunker(data, 10, digestUtil);
+    Chunker chunker = Chunker.builder(digestUtil).setInput(data).setChunkSize(10).build();
 
     assertThat(chunker.hasNext()).isTrue();
     assertThat(chunker.next()).isNotNull();
@@ -98,7 +98,7 @@ public class ChunkerTest {
   @Test
   public void emptyData() throws Exception {
     byte[] data = new byte[0];
-    Chunker chunker = new Chunker(data, digestUtil);
+    Chunker chunker = Chunker.builder(digestUtil).setInput(data).build();
 
     assertThat(chunker.hasNext()).isTrue();
 
@@ -121,7 +121,7 @@ public class ChunkerTest {
   @Test
   public void reset() throws Exception {
     byte[] data = new byte[]{1, 2, 3};
-    Chunker chunker = new Chunker(data, 1, digestUtil);
+    Chunker chunker = Chunker.builder(digestUtil).setInput(data).setChunkSize(1).build();
 
     assertNextEquals(chunker, (byte) 1);
     assertNextEquals(chunker, (byte) 2);
diff --git a/src/test/java/com/google/devtools/build/lib/remote/GrpcRemoteCacheTest.java b/src/test/java/com/google/devtools/build/lib/remote/GrpcRemoteCacheTest.java
index de453b3e68..a2391dbfd0 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/GrpcRemoteCacheTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/GrpcRemoteCacheTest.java
@@ -533,11 +533,11 @@ public class GrpcRemoteCacheTest {
         StreamObserver<WriteResponse> responseObserver, String contents, int chunkSizeBytes) {
       this.responseObserver = responseObserver;
       this.contents = contents;
-      try {
-        chunker = new Chunker(contents.getBytes(UTF_8), chunkSizeBytes, DIGEST_UTIL);
-      } catch (IOException e) {
-        fail("An error occurred:" + e);
-      }
+      chunker =
+          Chunker.builder(DIGEST_UTIL)
+              .setInput(contents.getBytes(UTF_8))
+              .setChunkSize(chunkSizeBytes)
+              .build();
     }
 
     @Override
diff --git a/src/test/java/com/google/devtools/build/lib/remote/TreeNodeRepositoryTest.java b/src/test/java/com/google/devtools/build/lib/remote/TreeNodeRepositoryTest.java
index 2f19f4669a..e31a337f30 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/TreeNodeRepositoryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/TreeNodeRepositoryTest.java
@@ -34,7 +34,8 @@ import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
 import com.google.devtools.remoteexecution.v1test.Digest;
 import com.google.devtools.remoteexecution.v1test.Directory;
 import java.io.IOException;
-import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.SortedMap;
 import java.util.TreeMap;
 import org.junit.Before;
@@ -107,15 +108,15 @@ public class TreeNodeRepositoryTest {
     Digest barDigest = repo.getMerkleDigest(barNode);
     assertThat(digests).containsExactly(rootDigest, aDigest, barDigest, fooDigest);
 
-    ArrayList<Directory> directories = new ArrayList<>();
-    ArrayList<ActionInput> actionInputs = new ArrayList<>();
+    Map<Digest, Directory> directories = new HashMap<>();
+    Map<Digest, ActionInput> actionInputs = new HashMap<>();
     repo.getDataFromDigests(digests, actionInputs, directories);
-    assertThat(actionInputs).containsExactly(bar, foo);
+    assertThat(actionInputs.values()).containsExactly(bar, foo);
     assertThat(directories).hasSize(2);
-    Directory rootDirectory = directories.get(0);
+    Directory rootDirectory = directories.get(rootDigest);
     assertThat(rootDirectory.getDirectories(0).getName()).isEqualTo("a");
     assertThat(rootDirectory.getDirectories(0).getDigest()).isEqualTo(aDigest);
-    Directory aDirectory = directories.get(1);
+    Directory aDirectory = directories.get(aDigest);
     assertThat(aDirectory.getFiles(0).getName()).isEqualTo("bar");
     assertThat(aDirectory.getFiles(0).getDigest()).isEqualTo(barDigest);
     assertThat(aDirectory.getFiles(1).getName()).isEqualTo("foo");
@@ -191,27 +192,27 @@ public class TreeNodeRepositoryTest {
             aClientDigest,
             bazDigest);
 
-    ArrayList<Directory> directories = new ArrayList<>();
-    ArrayList<ActionInput> actionInputs = new ArrayList<>();
+    Map<Digest, Directory> directories = new HashMap<>();
+    Map<Digest, ActionInput> actionInputs = new HashMap<>();
     repo.getDataFromDigests(digests, actionInputs, directories);
-    assertThat(actionInputs).containsExactly(bar, fooH, fooCc, baz);
+    assertThat(actionInputs.values()).containsExactly(bar, fooH, fooCc, baz);
     assertThat(directories).hasSize(4); // root, root/a, root/a/foo, and root/a-client
-    Directory rootDirectory = directories.get(0);
+    Directory rootDirectory = directories.get(rootDigest);
     assertThat(rootDirectory.getDirectories(0).getName()).isEqualTo("a");
     assertThat(rootDirectory.getDirectories(0).getDigest()).isEqualTo(aDigest);
     assertThat(rootDirectory.getDirectories(1).getName()).isEqualTo("a-client");
     assertThat(rootDirectory.getDirectories(1).getDigest()).isEqualTo(aClientDigest);
-    Directory aDirectory = directories.get(1);
+    Directory aDirectory = directories.get(aDigest);
     assertThat(aDirectory.getFiles(0).getName()).isEqualTo("bar.txt");
     assertThat(aDirectory.getFiles(0).getDigest()).isEqualTo(barDigest);
     assertThat(aDirectory.getDirectories(0).getName()).isEqualTo("foo");
     assertThat(aDirectory.getDirectories(0).getDigest()).isEqualTo(fooDigest);
-    Directory fooDirectory = directories.get(2);
+    Directory fooDirectory = directories.get(fooDigest);
     assertThat(fooDirectory.getFiles(0).getName()).isEqualTo("foo.cc");
     assertThat(fooDirectory.getFiles(0).getDigest()).isEqualTo(fooCcDigest);
     assertThat(fooDirectory.getFiles(1).getName()).isEqualTo("foo.h");
     assertThat(fooDirectory.getFiles(1).getDigest()).isEqualTo(fooHDigest);
-    Directory aClientDirectory = directories.get(3);
+    Directory aClientDirectory = directories.get(aClientDigest);
     assertThat(aClientDirectory.getFiles(0).getName()).isEqualTo("baz.txt");
     assertThat(aClientDirectory.getFiles(0).getDigest()).isEqualTo(bazDigest);
   }