Store MerkleTree nodes unwrapped

src/main/java/com/google/devtools/build/lib/actions/cache/VirtualActionInput.java

@@ -118,7 +118,8 @@ protected byte[] writeTo(Path target) throws IOException {
    * Gets a {@link ByteString} representation of the fake file. Used to avoid copying if the fake
    * file is internally represented as a {@link ByteString}.
    *
-   * <p>Prefer {@link #writeTo} to this method to avoid materializing the entire file in memory.
+   * <p>Prefer {@link #writeTo} to this method to avoid materializing the entire file in memory. The
+   * return value should not be retained.
    */
   public ByteString getBytes() {
     ByteString.Output out = ByteString.newOutput();

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

@@ -187,10 +187,11 @@ private ListenableFuture<Void> uploadBlob(
     ContentSource file = merkleTree.getFileByDigest(digest);
     if (file != null) {
       return switch (file) {
-        case ContentSource.VirtualActionInput(VirtualActionInput virtualActionInput) ->
-            // TODO: Avoid materializing the entire file in memory.
+        case ContentSource.VirtualActionInputSource(VirtualActionInput virtualActionInput) ->
+            // TODO: Avoid materializing the entire file in memory. This requires changing the
+            // upload to be driven by an OutputStream rather than consuming an InputStream.
             remoteCacheClient.uploadBlob(context, digest, virtualActionInput.getBytes());
-        case ContentSource.Path(Path path) -> {
+        case ContentSource.PathSource(Path path) -> {
           try {
             if (remotePathChecker.isRemote(context, path)) {
               // If we get here, the remote input was determined to exist in the remote or disk

src/main/java/com/google/devtools/build/lib/remote/merkletree/MerkleTree.java

@@ -32,6 +32,7 @@
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.InputMetadataProvider;
+import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
 import com.google.devtools.build.lib.remote.Scrubber.SpawnScrubber;
@@ -55,10 +56,9 @@ public class MerkleTree {
   private static final String BAZEL_TOOL_INPUT_MARKER = "bazel_tool_input";
 
   public sealed interface ContentSource {
-    record Path(com.google.devtools.build.lib.vfs.Path path) implements ContentSource {}
+    record PathSource(Path path) implements ContentSource {}
 
-    record VirtualActionInput(
-        com.google.devtools.build.lib.actions.cache.VirtualActionInput virtualActionInput)
+    record VirtualActionInputSource(VirtualActionInput virtualActionInput)
         implements ContentSource {}
   }
 
@@ -75,7 +75,8 @@ private interface MerkleTreeDirectoryVisitor {
   }
 
   private Map<Digest, Directory> digestDirectoryMap;
-  private Map<Digest, ContentSource> digestFileMap;
+  // Object is an unwrapped ContentSource to reduce retained memory when caching MerkleTrees.
+  private Map<Digest, Object> digestFileMap;
   @Nullable private final Directory rootProto;
   private final Digest rootDigest;
   private final SortedSet<DirectoryTree.FileNode> files;
@@ -157,19 +158,17 @@ private Map<Digest, Directory> getDigestDirectoryMap() {
     return this.digestDirectoryMap;
   }
 
-  private Map<Digest, ContentSource> getDigestFileMap() {
+  private Map<Digest, Object> getDigestFileMap() {
     if (this.digestFileMap == null) {
-      Map<Digest, ContentSource> newDigestMap = Maps.newHashMap();
+      Map<Digest, Object> newDigestMap = Maps.newHashMap();
       visitTree(
           (dir) -> {
             for (DirectoryTree.FileNode file : dir.getFiles()) {
-              ContentSource contentSource;
               if (file.getPath() != null) {
-                contentSource = new ContentSource.Path(file.getPath());
+                newDigestMap.put(file.getDigest(), file.getPath());
               } else {
-                contentSource = new ContentSource.VirtualActionInput(file.getVirtualActionInput());
+                newDigestMap.put(file.getDigest(), file.getVirtualActionInput());
               }
-              newDigestMap.put(file.getDigest(), contentSource);
             }
           });
       this.digestFileMap = newDigestMap;
@@ -184,7 +183,13 @@ public Directory getDirectoryByDigest(Digest digest) {
 
   @Nullable
   public ContentSource getFileByDigest(Digest digest) {
-    return getDigestFileMap().get(digest);
+    Object pathOrVirtualActionInput = getDigestFileMap().get(digest);
+    if (pathOrVirtualActionInput instanceof Path path) {
+      return new ContentSource.PathSource(path);
+    } else {
+      return new ContentSource.VirtualActionInputSource(
+          (VirtualActionInput) pathOrVirtualActionInput);
+    }
   }
 
   /**

src/test/java/com/google/devtools/build/lib/remote/merkletree/MerkleTreeTest.java

@@ -140,13 +140,13 @@ public void buildMerkleTree() throws IOException {
           digestUtil.computeAsUtf8("fizzbuzz")
         };
     assertThat(tree.getFileByDigest(inputDigests[0]))
-        .isEqualTo(new ContentSource.Path(foo.getPath()));
+        .isEqualTo(new ContentSource.PathSource(foo.getPath()));
     assertThat(tree.getFileByDigest(inputDigests[1]))
-        .isEqualTo(new ContentSource.Path(bar.getPath()));
+        .isEqualTo(new ContentSource.PathSource(bar.getPath()));
     assertThat(tree.getFileByDigest(inputDigests[2]))
-        .isEqualTo(new ContentSource.Path(buzz.getPath()));
+        .isEqualTo(new ContentSource.PathSource(buzz.getPath()));
     assertThat(tree.getFileByDigest(inputDigests[3]))
-        .isEqualTo(new ContentSource.Path(fizzbuzz.getPath()));
+        .isEqualTo(new ContentSource.PathSource(fizzbuzz.getPath()));
 
     Digest[] allDigests = Iterables.toArray(tree.getAllDigests(), Digest.class);
     assertThat(allDigests.length).isEqualTo(dirDigests.length + inputDigests.length);