Reduce Merkle tree footprint by using an ImmutableSortedMap

MODULE.bazel

@@ -105,6 +105,10 @@ single_version_override(
 # We need to package the lite runtime jar for Bazel's bootstrap build
 # 2. Partially revert of https://github.com/protocolbuffers/protobuf/commit/7ec5e19c0743b80c5097434a05cc0e972e7b4e7a
 # to support grpc
+# 3. Pull in the removal of the memoizedIsInitialized field for messages with no
+# required fields early. This reduces the memory usage of Bazel and avoids
+# spread in memory usage between Bazel and Blaze.
+# https://github.com/protocolbuffers/protobuf/blob/adf9e8b80fd10398b82e16c4ea72f2d2cffb0e9b/src/google/protobuf/compiler/java/full/message.cc#L884C1-L899C6
 single_version_override(
     module_name = "protobuf",
     patch_strip = 1,

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

@@ -15,7 +15,7 @@
 
 import build.bazel.remote.execution.v2.Digest;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
@@ -24,6 +24,7 @@
 import java.util.Collection;
 import java.util.Map;
 import java.util.Optional;
+import java.util.SortedMap;
 
 /**
  * A representation of the inputs to a remotely executed action represented as a Merkle tree.
@@ -79,11 +80,14 @@ record BlobsDiscarded(Digest digest, long inputFiles, long inputBytes) implement
    */
   final class Uploadable implements MerkleTree {
     private final RootOnly.BlobsUploaded root;
-    private final ImmutableMap<Digest, /* byte[] | Path | VirtualActionInput */ Object> blobs;
+    private final ImmutableSortedMap<Digest, /* byte[] | Path | VirtualActionInput */ Object> blobs;
 
-    Uploadable(RootOnly.BlobsUploaded root, ImmutableMap<Digest, Object> blobs) {
+    Uploadable(RootOnly.BlobsUploaded root, SortedMap<Digest, Object> blobs) {
       this.root = root;
-      this.blobs = blobs;
+      // A sorted map requires less memory than a regular hash map as it only stores two flat sorted
+      // arrays. Access performance is not critical since it's only used to find missing blobs,
+      // which always require network access.
+      this.blobs = ImmutableSortedMap.copyOfSorted(blobs);
     }
 
     @Override

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

@@ -22,6 +22,7 @@
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static com.google.common.util.concurrent.Futures.transform;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+import static com.google.devtools.build.lib.remote.util.DigestUtil.DIGEST_COMPARATOR;
 import static com.google.devtools.build.lib.util.StringEncoding.internalToUnicode;
 import static com.google.devtools.build.lib.vfs.PathFragment.HIERARCHICAL_COMPARATOR;
 import static java.util.Comparator.comparing;
@@ -88,6 +89,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -193,7 +195,7 @@ public MerkleTreeComputer(
     this.emptyDigest = digestUtil.compute(emptyBlob);
     this.emptyTree =
         new MerkleTree.Uploadable(
-            new MerkleTree.RootOnly.BlobsUploaded(emptyDigest, 0, 0), ImmutableMap.of());
+            new MerkleTree.RootOnly.BlobsUploaded(emptyDigest, 0, 0), ImmutableSortedMap.of());
   }
 
   /** Specifies which blobs should be retained in the Merkle tree. */
@@ -468,7 +470,8 @@ private MerkleTree buildWithPrecomputedSubTrees(
 
     long inputFiles = 0;
     long inputBytes = 0;
-    var blobs = ImmutableMap.<Digest, Object>builder();
+    var blobs =
+        new TreeMap<Digest, /* byte[] | Path | VirtualActionInput */ Object>(DIGEST_COMPARATOR);
     Deque<Directory.Builder> directoryStack = new ArrayDeque<>();
     directoryStack.push(Directory.newBuilder());
 
@@ -531,17 +534,16 @@ private MerkleTree buildWithPrecomputedSubTrees(
           inputBytes += directoryBlobDigest.getSizeBytes();
           var topDirectory = directoryStack.peek();
           if (topDirectory == null) {
-            var builtBlobs = blobs.buildKeepingLast();
             if (blobPolicy == BlobPolicy.DISCARD) {
               // Make sure that we didn't unnecessarily retain any blobs.
-              checkState(builtBlobs.isEmpty());
+              checkState(blobs.isEmpty());
               return new MerkleTree.RootOnly.BlobsDiscarded(
                   directoryBlobDigest, inputFiles, inputBytes);
             } else {
               return new MerkleTree.Uploadable(
                   new MerkleTree.RootOnly.BlobsUploaded(
                       directoryBlobDigest, inputFiles, inputBytes),
-                  builtBlobs);
+                  blobs);
             }
           }
           topDirectory

src/main/java/com/google/devtools/build/lib/remote/util/DigestUtil.java

@@ -15,6 +15,7 @@
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Comparator.comparing;
 
 import build.bazel.remote.execution.v2.Action;
 import build.bazel.remote.execution.v2.Digest;
@@ -31,13 +32,19 @@
 import com.google.devtools.build.lib.vfs.FileStatus;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.XattrProvider;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.util.Arrays;
+import java.util.Comparator;
 
 /** Utility methods to work with {@link Digest}. */
 public class DigestUtil {
+  public static final Comparator<Digest> DIGEST_COMPARATOR =
+      comparing(Digest::getHashBytes, ByteString.unsignedLexicographicalComparator())
+          .thenComparing(Digest::getSizeBytes);
+
   private final XattrProvider xattrProvider;
   private final DigestHashFunction hashFn;
   private final DigestFunction.Value digestFunction;

src/test/java/com/google/devtools/build/lib/remote/RemoteExecutionServiceTest.java

@@ -806,46 +806,41 @@ public void buildRemoteAction_goldenTest(@TestParameter({"1", "2", "3"}) int see
     // little memory as possible since they are kept in memory during the whole remote execution.
     // Memory usage isn't expected to differ by seed, so only check for one of them.
     if (seed == 1) {
-      var merkleTree =
-          service
-              .buildRemoteAction(spawn, context, MerkleTreeComputer.BlobPolicy.KEEP)
-              .getMerkleTree();
-      assertThat(merkleTree).isInstanceOf(MerkleTree.Uploadable.class);
-      // These are roots that are already retained while a remote action is being executed or are
-      // effectively global singletons.
-      var otherInput = ActionsTestUtil.createArtifact(artifactRoot, "some/other/input/file.txt");
-      fakeFileCache.createScratchInput(otherInput, "some other content");
-      var otherSpawn = new SpawnBuilder().withInput(otherInput).withOutput("foo").build();
-      var someOtherMerkleTree =
-          service.buildRemoteAction(otherSpawn, newSpawnExecutionContext(otherSpawn));
-
       // JOL tracks objects by their native address, so run GC to minimize noise from moved objects.
       System.gc();
-      var alreadyRetainedObjects =
-          GraphLayout.parseInstance(spawn, fakeFileCache, digestUtil, someOtherMerkleTree);
-      // This covers objects internal to JOL itself as well as metadata lazily attached to classes
-      // related to MerkleTree. Note that this has to be computed in a separate parseInstance call
-      // since the first call on someOtherMerkleTree may thus not have captured the metadata
-      // attached to the Class objects it references (e.g. cached field lists).
-      var jolInternalObjects =
-          GraphLayout.parseInstance(alreadyRetainedObjects, someOtherMerkleTree);
-      var merkleTreeTransitiveRetention = GraphLayout.parseInstance(merkleTree);
-
-      var merkleTreeOnlyRetention =
-          merkleTreeTransitiveRetention
-              .subtract(alreadyRetainedObjects)
-              .subtract(jolInternalObjects);
-      // Output the objects that are transitively retained by merkleTreeTransitiveRetention but
-      // neither by alreadyRetainedObjects nor jolInternalObjects for manual inspection.
+      // Keep building a Merkle tree and compute its retained size relative to all previous trees
+      // until the size stabilizes. The goal is to compute the size of the objects that are uniquely
+      // retained by the tree, as this is the effective overhead at runtime when building many trees
+      // in parallel. This is more delicate than it seems:
+      // 1. This has to be done in a loop since objects retain their respective Class and JOL's use
+      //    of reflection mutates the various caches in the Class objects in non-deterministic ways.
+      // 2. Subtracting previous trees is only correct under the assumption that the objects shared
+      //    with such trees would also be retained elsewhere (e.g. Artifact objects). If MerkleTree
+      //    ever uses techniques such as interning or weak caches, this strategy would have to be
+      //    revisited.
+      GraphLayout merkleTreeUniqueRetention;
+      var previousRoots = new ArrayList<>();
+      long stableRetainedSize = -1;
+      while (true) {
+        var merkleTree =
+            service
+                .buildRemoteAction(spawn, context, MerkleTreeComputer.BlobPolicy.KEEP)
+                .getMerkleTree();
+        assertThat(merkleTree).isInstanceOf(MerkleTree.Uploadable.class);
+        merkleTreeUniqueRetention =
+            GraphLayout.parseInstance(merkleTree)
+                .subtract(GraphLayout.parseInstance(previousRoots));
+        if (merkleTreeUniqueRetention.totalSize() == stableRetainedSize) {
+          break;
+        }
+        stableRetainedSize = merkleTreeUniqueRetention.totalSize();
+        previousRoots.add(merkleTree);
+      }
       var footprintOut =
           Paths.get(System.getenv("TEST_UNDECLARED_OUTPUTS_DIR"), "merkle_tree_footprint.txt");
-      Files.writeString(footprintOut, merkleTreeOnlyRetention.toFootprint());
+      Files.writeString(footprintOut, merkleTreeUniqueRetention.toFootprint());
       // TODO: Get this number down.
-      // TODO: Assert the size exactly after switching to an ImmutableSortedMap rather than an
-      //       ImmutableMap - proto messages don't have deterministic hash codes, which makes the
-      //       exact size vary slightly between runs.
-      assertThat(merkleTreeOnlyRetention.totalSize()).isAtLeast(8400);
-      assertThat(merkleTreeOnlyRetention.totalSize()).isAtMost(9500);
+      assertThat(stableRetainedSize).isEqualTo(7872);
     }
   }
 

third_party/protobuf.patch

@@ -45,3 +45,34 @@ index 426bf9124..fd17ac96c 100644
  # Amalgamation #################################################################
 
  upb_amalgamation(
+diff --git a/src/google/protobuf/compiler/java/full/message.cc b/src/google/protobuf/compiler/java/full/message.cc
+index dd290f98c..1f100c568 100644
+--- a/src/google/protobuf/compiler/java/full/message.cc
++++ b/src/google/protobuf/compiler/java/full/message.cc
+@@ -842,7 +842,25 @@ void ImmutableMessageGenerator::GenerateIsInitialized(io::Printer* printer) {
+   // Memoizes whether the protocol buffer is fully initialized (has all
+   // required fields). -1 means not yet computed. 0 means false and 1 means
+   // true.
+-  printer->Print("private byte memoizedIsInitialized = -1;\n");
++
++  // If the message transitively has no required fields or extensions,
++  // isInitialized() is always true.
++  if (!HasRequiredFields(descriptor_)) {
++    printer->Print(
++        "/**\n"
++        "  * @deprecated This always returns true for this type as it \n"
++        "  *   does not transitively contain any required fields.\n"
++        "  */\n"
++        "@java.lang.Deprecated\n"
++        "@java.lang.Override\n"
++        "public final boolean isInitialized() {\n"
++        "  return true;\n"
++        "}\n"
++        "\n");
++    return;
++  }
++
++  printer->Print("private transient byte memoizedIsInitialized = -1;\n");
+   printer->Print(
+       "@java.lang.Override\n"
+       "public final boolean isInitialized() {\n");