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

10 files changed, 151 insertions, 118 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/remote/AbstractRemoteActionCache.java b/src/main/java/com/google/devtools/build/lib/remote/AbstractRemoteActionCache.java
index 66f29c5120..ccc6de6dee 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/AbstractRemoteActionCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/AbstractRemoteActionCache.java
@@ -518,7 +518,8 @@ public abstract class AbstractRemoteActionCache implements AutoCloseable {
 
       byte[] blob = tree.build().toByteArray();
       Digest digest = digestUtil.compute(blob);
-      Chunker chunker = new Chunker(blob, blob.length, digestUtil);
+      Chunker chunker =
+          Chunker.builder(digestUtil).setInput(digest, blob).setChunkSize(blob.length).build();
 
       if (result != null) {
         result
diff --git a/src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java b/src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java
index 3f3308c5e9..d44bfcf12e 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/ByteStreamBuildEventArtifactUploader.java
@@ -21,6 +21,7 @@ import com.google.common.util.concurrent.MoreExecutors;
 import com.google.devtools.build.lib.buildeventstream.BuildEvent.LocalFile;
 import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader;
 import com.google.devtools.build.lib.buildeventstream.PathConverter;
+import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.remoteexecution.v1test.Digest;
 import io.grpc.Context;
@@ -65,8 +66,9 @@ class ByteStreamBuildEventArtifactUploader implements BuildEventArtifactUploader
     Context prevCtx = ctx.attach();
     try {
       for (Path file : files.keySet()) {
-        Chunker chunker = new Chunker(file);
-        Digest digest = chunker.digest();
+        DigestUtil digestUtil = new DigestUtil(file.getFileSystem().getDigestFunction());
+        Digest digest = digestUtil.compute(file);
+        Chunker chunker = Chunker.builder(digestUtil).setInput(digest, file).build();
         ListenableFuture<PathDigestPair> upload =
             Futures.transform(
                 uploader.uploadBlobAsync(chunker, /*forceUpload=*/false),
diff --git a/src/main/java/com/google/devtools/build/lib/remote/Chunker.java b/src/main/java/com/google/devtools/build/lib/remote/Chunker.java
index 7e960ac43d..6d313dcc07 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/Chunker.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/Chunker.java
@@ -18,10 +18,10 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
 import com.google.common.io.ByteStreams;
 import com.google.devtools.build.lib.actions.ActionInput;
-import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.vfs.Path;
@@ -115,65 +115,7 @@ public final class Chunker {
   // lazily on the first call to next(), as opposed to opening it in the constructor or on reset().
   private boolean initialized;
 
-  public Chunker(byte[] data, DigestUtil digestUtil) throws IOException {
-    this(data, getDefaultChunkSize(), digestUtil);
-  }
-
-  public Chunker(byte[] data, int chunkSize, DigestUtil digestUtil) throws IOException {
-    this(() -> new ByteArrayInputStream(data), digestUtil.compute(data), chunkSize, digestUtil);
-  }
-
-  public Chunker(Path file) throws IOException {
-    this(file, getDefaultChunkSize());
-  }
-
-  public Chunker(Path file, int chunkSize) throws IOException {
-    this(
-        () -> {
-          try {
-            return file.getInputStream();
-          } catch (IOException e) {
-            throw new RuntimeException(e);
-          }
-        },
-        new DigestUtil(file.getFileSystem().getDigestFunction()).compute(file),
-        chunkSize,
-        new DigestUtil(file.getFileSystem().getDigestFunction()));
-  }
-
-  public Chunker(
-      ActionInput actionInput, MetadataProvider inputCache, Path execRoot, DigestUtil digestUtil)
-      throws IOException {
-    this(actionInput, inputCache, execRoot, getDefaultChunkSize(), digestUtil);
-  }
-
-  public Chunker(VirtualActionInput actionInput, DigestUtil digestUtil) throws IOException {
-    this(actionInput.getBytes().toByteArray(), digestUtil);
-  }
-
-  public Chunker(
-      ActionInput actionInput,
-      MetadataProvider inputCache,
-      Path execRoot,
-      int chunkSize,
-      DigestUtil digestUtil)
-      throws IOException {
-    this(
-        () -> {
-          try {
-            return execRoot.getRelative(actionInput.getExecPathString()).getInputStream();
-          } catch (IOException e) {
-            throw new RuntimeException(e);
-          }
-        },
-        DigestUtil.getFromInputCache(actionInput, inputCache),
-        chunkSize,
-        digestUtil);
-  }
-
-  @VisibleForTesting
-  Chunker(Supplier<InputStream> dataSupplier, Digest digest, int chunkSize, DigestUtil digestUtil)
-      throws IOException {
+  Chunker(Supplier<InputStream> dataSupplier, Digest digest, int chunkSize, DigestUtil digestUtil) {
     this.dataSupplier = checkNotNull(dataSupplier);
     this.digest = checkNotNull(digest);
     this.chunkSize = chunkSize;
@@ -282,4 +224,83 @@ public final class Chunker {
     }
     initialized = true;
   }
+
+  public static Builder builder(DigestUtil digestUtil) {
+    return new Builder(digestUtil);
+  }
+
+  /** Builder class for the Chunker */
+  public static class Builder {
+    private final DigestUtil digestUtil;
+    private int chunkSize = getDefaultChunkSize();
+    private Digest digest;
+    private Supplier<InputStream> inputStream;
+
+    Builder(DigestUtil digestUtil) {
+      this.digestUtil = digestUtil;
+    }
+
+    public Builder setInput(byte[] data) {
+      Preconditions.checkState(inputStream == null);
+      digest = digestUtil.compute(data);
+      inputStream = () -> new ByteArrayInputStream(data);
+      return this;
+    }
+
+    public Builder setInput(Digest digest, byte[] data) {
+      Preconditions.checkState(inputStream == null);
+      this.digest = digest;
+      inputStream = () -> new ByteArrayInputStream(data);
+      return this;
+    }
+
+    public Builder setInput(Digest digest, Path file) {
+      Preconditions.checkState(inputStream == null);
+      this.digest = digest;
+      inputStream =
+          () -> {
+            try {
+              return file.getInputStream();
+            } catch (IOException e) {
+              throw new RuntimeException(e);
+            }
+          };
+      return this;
+    }
+
+    public Builder setInput(Digest digest, ActionInput actionInput, Path execRoot) {
+      Preconditions.checkState(inputStream == null);
+      this.digest = digest;
+      if (actionInput instanceof VirtualActionInput) {
+        this.inputStream =
+            () -> {
+              try {
+                return ((VirtualActionInput) actionInput).getBytes().newInput();
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            };
+      } else {
+        inputStream =
+            () -> {
+              try {
+                return execRoot.getRelative(actionInput.getExecPathString()).getInputStream();
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            };
+      }
+      return this;
+    }
+
+    public Builder setChunkSize(int chunkSize) {
+      this.chunkSize = chunkSize;
+      return this;
+    }
+
+    public Chunker build() {
+      Preconditions.checkNotNull(inputStream, digest);
+      return new Chunker(inputStream, digest, chunkSize, digestUtil);
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteCache.java b/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteCache.java
index c1755599a0..39f443ff54 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteCache.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteCache.java
@@ -27,8 +27,6 @@ import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ExecException;
-import com.google.devtools.build.lib.actions.MetadataProvider;
-import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.remote.Retrier.RetryException;
 import com.google.devtools.build.lib.remote.TreeNodeRepository.TreeNode;
@@ -57,6 +55,7 @@ import java.io.IOException;
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -153,28 +152,28 @@ public class GrpcRemoteCache extends AbstractRemoteActionCache {
 
     List<Chunker> toUpload = new ArrayList<>();
     // Only upload data that was missing from the cache.
-    ArrayList<ActionInput> missingActionInputs = new ArrayList<>();
-    ArrayList<Directory> missingTreeNodes = new ArrayList<>();
+    Map<Digest, ActionInput> missingActionInputs = new HashMap<>();
+    Map<Digest, Directory> missingTreeNodes = new HashMap<>();
     HashSet<Digest> missingTreeDigests = new HashSet<>(missingDigests);
     missingTreeDigests.remove(commandDigest);
     repository.getDataFromDigests(missingTreeDigests, missingActionInputs, missingTreeNodes);
 
     if (missingDigests.contains(commandDigest)) {
-      toUpload.add(new Chunker(command.toByteArray(), digestUtil));
+      toUpload.add(
+          Chunker.builder(digestUtil).setInput(commandDigest, command.toByteArray()).build());
     }
     if (!missingTreeNodes.isEmpty()) {
-      for (Directory d : missingTreeNodes) {
-        toUpload.add(new Chunker(d.toByteArray(), digestUtil));
+      for (Map.Entry<Digest, Directory> entry : missingTreeNodes.entrySet()) {
+        Digest digest = entry.getKey();
+        Directory d = entry.getValue();
+        toUpload.add(Chunker.builder(digestUtil).setInput(digest, d.toByteArray()).build());
       }
     }
     if (!missingActionInputs.isEmpty()) {
-      MetadataProvider inputFileCache = repository.getInputFileCache();
-      for (ActionInput actionInput : missingActionInputs) {
-        if (actionInput instanceof VirtualActionInput) {
-          toUpload.add(new Chunker((VirtualActionInput) actionInput, digestUtil));
-        } else {
-          toUpload.add(new Chunker(actionInput, inputFileCache, execRoot, digestUtil));
-        }
+      for (Map.Entry<Digest, ActionInput> entry : missingActionInputs.entrySet()) {
+        Digest digest = entry.getKey();
+        ActionInput actionInput = entry.getValue();
+        toUpload.add(Chunker.builder(digestUtil).setInput(digest, actionInput, execRoot).build());
       }
     }
     uploader.uploadBlobs(toUpload, true);
@@ -290,7 +289,7 @@ public class GrpcRemoteCache extends AbstractRemoteActionCache {
       Chunker chunker;
       Path file = digestToFile.get(digest);
       if (file != null) {
-        chunker = new Chunker(file);
+        chunker = Chunker.builder(digestUtil).setInput(digest, file).build();
       } else {
         chunker = digestToChunkers.get(digest);
         if (chunker == null) {
@@ -326,7 +325,7 @@ public class GrpcRemoteCache extends AbstractRemoteActionCache {
     Digest digest = digestUtil.compute(file);
     ImmutableSet<Digest> missing = getMissingDigests(ImmutableList.of(digest));
     if (!missing.isEmpty()) {
-      uploader.uploadBlob(new Chunker(file), true);
+      uploader.uploadBlob(Chunker.builder(digestUtil).setInput(digest, file).build(), true);
     }
     return digest;
   }
@@ -335,7 +334,7 @@ public class GrpcRemoteCache extends AbstractRemoteActionCache {
     Digest digest = digestUtil.compute(blob);
     ImmutableSet<Digest> missing = getMissingDigests(ImmutableList.of(digest));
     if (!missing.isEmpty()) {
-      uploader.uploadBlob(new Chunker(blob, digestUtil), true);
+      uploader.uploadBlob(Chunker.builder(digestUtil).setInput(digest, blob).build(), true);
     }
     return digest;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/remote/TreeNodeRepository.java b/src/main/java/com/google/devtools/build/lib/remote/TreeNodeRepository.java
index 114ce46ddc..6f4366ea3d 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/TreeNodeRepository.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/TreeNodeRepository.java
@@ -475,11 +475,13 @@ public final class TreeNodeRepository {
    * or Directory messages by cached digests and adds them to the lists.
    */
   public void getDataFromDigests(
-      Iterable<Digest> digests, List<ActionInput> actionInputs, List<Directory> nodes) {
+      Iterable<Digest> digests,
+      Map<Digest, ActionInput> actionInputs,
+      Map<Digest, Directory> nodes) {
     for (Digest digest : digests) {
       TreeNode treeNode = digestTreeNodeCache.get(digest);
       if (treeNode != null) {
-        nodes.add(Preconditions.checkNotNull(directoryCache.get(treeNode)));
+        nodes.put(digest, Preconditions.checkNotNull(directoryCache.get(treeNode)));
       } else { // If not there, it must be an ActionInput.
         ByteString hexDigest = ByteString.copyFromUtf8(digest.getHash());
         ActionInput input = reverseInputMap.get(hexDigest);
@@ -487,7 +489,7 @@ public final class TreeNodeRepository {
           // ... or a VirtualActionInput.
           input = digestVirtualInputCache.get(digest);
         }
-        actionInputs.add(Preconditions.checkNotNull(input));
+        actionInputs.put(digest, Preconditions.checkNotNull(input));
       }
     }
   }
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);
   }
diff --git a/src/tools/remote/src/main/java/com/google/devtools/build/remote/worker/ByteStreamServer.java b/src/tools/remote/src/main/java/com/google/devtools/build/remote/worker/ByteStreamServer.java
index 066d93cf7c..08c4163621 100644
--- a/src/tools/remote/src/main/java/com/google/devtools/build/remote/worker/ByteStreamServer.java
+++ b/src/tools/remote/src/main/java/com/google/devtools/build/remote/worker/ByteStreamServer.java
@@ -81,7 +81,10 @@ final class ByteStreamServer extends ByteStreamImplBase {
     try {
       // This still relies on the blob size to be small enough to fit in memory.
       // TODO(olaola): refactor to fix this if the need arises.
-      Chunker c = new Chunker(getFromFuture(cache.downloadBlob(digest)), digestUtil);
+      Chunker c =
+          Chunker.builder(digestUtil)
+              .setInput(digest, getFromFuture(cache.downloadBlob(digest)))
+              .build();
       while (c.hasNext()) {
         responseObserver.onNext(
             ReadResponse.newBuilder().setData(c.next().getData()).build());