diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java index 18d2ae50950ecf..71494c18acf4df 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java @@ -52,8 +52,11 @@ import build.bazel.remote.execution.v2.RequestMetadata; import build.bazel.remote.execution.v2.SymlinkNode; import build.bazel.remote.execution.v2.Tree; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.base.Strings; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -69,6 +72,7 @@ import com.google.devtools.build.lib.actions.ExecException; import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue; import com.google.devtools.build.lib.actions.ForbiddenActionInputException; +import com.google.devtools.build.lib.actions.MetadataProvider; import com.google.devtools.build.lib.actions.Spawn; import com.google.devtools.build.lib.actions.Spawns; import com.google.devtools.build.lib.actions.UserExecException; @@ -116,6 +120,7 @@ import java.util.Objects; import java.util.SortedMap; import java.util.TreeSet; +import java.util.concurrent.ConcurrentLinkedQueue; import javax.annotation.Nullable; /** @@ -133,6 +138,7 @@ public class RemoteExecutionService { @Nullable private final RemoteExecutionClient remoteExecutor; private final ImmutableSet filesToDownload; @Nullable private final Path captureCorruptedOutputsDir; + private final Cache merkleTreeCache; public RemoteExecutionService( Path execRoot, @@ -154,6 +160,13 @@ public RemoteExecutionService( this.remoteCache = remoteCache; this.remoteExecutor = remoteExecutor; + CacheBuilder merkleTreeCacheBuilder = CacheBuilder.newBuilder() + .softValues(); + if (remoteOptions.remoteMerkleTreesCacheSize != -1) { + merkleTreeCacheBuilder.maximumSize(remoteOptions.remoteMerkleTreesCacheSize); + } + merkleTreeCache = merkleTreeCacheBuilder.build(); + ImmutableSet.Builder filesToDownloadBuilder = ImmutableSet.builder(); for (ActionInput actionInput : filesToDownload) { filesToDownloadBuilder.add(actionInput.getExecPath()); @@ -331,12 +344,53 @@ public boolean mayBeExecutedRemotely(Spawn spawn) { && Spawns.mayBeExecutedRemotely(spawn); } + private MerkleTree buildInputMerkleTree(SpawnExecutionContext context) + throws IOException, ForbiddenActionInputException { + if (remoteOptions.remoteMerkleTreesCacheSize != 0) { + MetadataProvider metadataProvider = context.getMetadataProvider(); + ConcurrentLinkedQueue subMerkleTrees = new ConcurrentLinkedQueue(); + remotePathResolver.walkInputs( + context, + (Object nodeKey, SpawnExecutionContext.InputWalker walker) -> { + subMerkleTrees.add(buildMerkleTreeVisitor(nodeKey, walker, metadataProvider)); + }); + return MerkleTree.merge(subMerkleTrees, digestUtil); + } else { + SortedMap inputMap = remotePathResolver.getInputMapping(context); + return MerkleTree.build(inputMap, context.getMetadataProvider(), execRoot, digestUtil); + } + } + + private MerkleTree buildMerkleTreeVisitor( + Object nodeKey, SpawnExecutionContext.InputWalker walker, MetadataProvider metadataProvider) + throws IOException, ForbiddenActionInputException { + MerkleTree result = merkleTreeCache.getIfPresent(nodeKey); + if (result == null) { + result = uncachedBuildMerkleTreeVisitor(walker, metadataProvider); + merkleTreeCache.put(nodeKey, result); + } + return result; + } + + @VisibleForTesting + public MerkleTree uncachedBuildMerkleTreeVisitor( + SpawnExecutionContext.InputWalker walker, MetadataProvider metadataProvider) + throws IOException, ForbiddenActionInputException { + ConcurrentLinkedQueue subMerkleTrees = new ConcurrentLinkedQueue(); + subMerkleTrees.add(MerkleTree.build( + walker.getLeavesInputMapping(), metadataProvider, execRoot, digestUtil)); + walker.visitNonLeaves( + (Object subNodeKey, SpawnExecutionContext.InputWalker subWalker) -> { + subMerkleTrees.add(buildMerkleTreeVisitor( + subNodeKey, subWalker, metadataProvider)); + }); + return MerkleTree.merge(subMerkleTrees, digestUtil); + } + /** Creates a new {@link RemoteAction} instance from spawn. */ public RemoteAction buildRemoteAction(Spawn spawn, SpawnExecutionContext context) throws IOException, UserExecException, ForbiddenActionInputException { - SortedMap inputMap = remotePathResolver.getInputMapping(context); - final MerkleTree merkleTree = - MerkleTree.build(inputMap, context.getMetadataProvider(), execRoot, digestUtil); + final MerkleTree merkleTree = buildInputMerkleTree(context); // Get the remote platform properties. Platform platform = PlatformUtils.getPlatformProto(spawn, remoteOptions); diff --git a/src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java b/src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java index cb2f1639807279..d08206be2ef9c7 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java +++ b/src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java @@ -504,6 +504,19 @@ public RemoteOutputsStrategyConverter() { + " discard the remotely cached values if they don't match the expected value.") public boolean remoteVerifyDownloads; + @Option( + name = "experimental_remote_merkle_tree_cache_size", + defaultValue = "0", + documentationCategory = OptionDocumentationCategory.REMOTE, + effectTags = {OptionEffectTag.UNKNOWN}, + help = + "The number of Merkle trees to memoize to improve the remote cache hit checking speed. " + + "Enabling this option will increase the memory foot print, although the cache is " + + "automatically pruned according to Java's handling of soft references. If set to " + + "0 (default), the memoization is disabled. If set to -1, the cache size is " + + "unlimited.") + public long remoteMerkleTreesCacheSize; + @Option( name = "remote_download_symlink_template", defaultValue = "", diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteExecutionServiceTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteExecutionServiceTest.java index 33f7fffdbef6d8..bf120fba8ac1ca 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/RemoteExecutionServiceTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/RemoteExecutionServiceTest.java @@ -46,6 +46,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.devtools.build.lib.actions.ActionInput; +import com.google.devtools.build.lib.actions.ActionInputHelper; import com.google.devtools.build.lib.actions.Artifact; import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact; import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact; @@ -58,6 +59,7 @@ import com.google.devtools.build.lib.actions.cache.MetadataInjector; import com.google.devtools.build.lib.actions.util.ActionsTestUtil; import com.google.devtools.build.lib.clock.JavaClock; +import com.google.devtools.build.lib.collect.nestedset.NestedSet; import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder; import com.google.devtools.build.lib.collect.nestedset.Order; import com.google.devtools.build.lib.exec.util.FakeOwner; @@ -1044,6 +1046,92 @@ public void downloadOutputs_missingInMemoryOutputWithMinimal_returnsNull() throw verify(injector, never()).injectFile(eq(a1), remoteFileMatchingDigest(d1)); } + @Test + public void buildMerkleTree_withMemoization_works() throws Exception { + // Test that Merkle tree building can be memoized. + + // TODO: Would like to check that NestedSet.getNonLeaves() is only called once per node, but + // cannot Mockito.spy on NestedSet as it is final. + + // arrange + /* + * First: + * /bar/file + * /foo1/file + * Second: + * /bar/file + * /foo2/file + */ + + // arrange + // Single node NestedSets are folded, so always add a dummy file everywhere. + ActionInput dummyFile = ActionInputHelper.fromPath("dummy"); + fakeFileCache.createScratchInput(dummyFile, "dummy"); + + ActionInput barFile = ActionInputHelper.fromPath("bar/file"); + NestedSet nodeBar = NestedSetBuilder.create( + Order.STABLE_ORDER, dummyFile, barFile); + fakeFileCache.createScratchInput(barFile, "bar"); + + ActionInput foo1File = ActionInputHelper.fromPath("foo1/file"); + NestedSet nodeFoo1 = NestedSetBuilder.create( + Order.STABLE_ORDER, dummyFile, foo1File); + fakeFileCache.createScratchInput(foo1File, "foo1"); + + ActionInput foo2File = ActionInputHelper.fromPath("foo2/file"); + NestedSet nodeFoo2 = NestedSetBuilder.create( + Order.STABLE_ORDER, dummyFile, foo2File); + fakeFileCache.createScratchInput(foo2File, "foo2"); + + NestedSet nodeRoot1 = new NestedSetBuilder(Order.STABLE_ORDER) + .add(dummyFile) + .addTransitive(nodeBar) + .addTransitive(nodeFoo1) + .build(); + NestedSet nodeRoot2 = new NestedSetBuilder(Order.STABLE_ORDER) + .add(dummyFile) + .addTransitive(nodeBar) + .addTransitive(nodeFoo2) + .build(); + + Spawn spawn1 = new SimpleSpawn( + new FakeOwner("foo", "bar", "//dummy:label"), + /*arguments=*/ ImmutableList.of(), + /*environment=*/ ImmutableMap.of(), + /*executionInfo=*/ ImmutableMap.of(), + /*inputs=*/ nodeRoot1, + /*outputs=*/ ImmutableSet.of(), + ResourceSet.ZERO); + Spawn spawn2 = new SimpleSpawn( + new FakeOwner("foo", "bar", "//dummy:label"), + /*arguments=*/ ImmutableList.of(), + /*environment=*/ ImmutableMap.of(), + /*executionInfo=*/ ImmutableMap.of(), + /*inputs=*/ nodeRoot2, + /*outputs=*/ ImmutableSet.of(), + ResourceSet.ZERO); + + FakeSpawnExecutionContext context1 = newSpawnExecutionContext(spawn1); + FakeSpawnExecutionContext context2 = newSpawnExecutionContext(spawn2); + RemoteOptions remoteOptions = Options.getDefaults(RemoteOptions.class); + remoteOptions.remoteMerkleTreesCacheSize = -1; + RemoteExecutionService service = spy(newRemoteExecutionService(remoteOptions)); + + // act first time + service.buildRemoteAction(spawn1, context1); + + // assert first time + // Called for: manifests, runfiles, nodeRoot1, nodeFoo1 and nodeBar. + verify(service, times(5)).uncachedBuildMerkleTreeVisitor(any(), any()); + + // act second time + service.buildRemoteAction(spawn2, context2); + + // assert second time + // Called again for: manifests, runfiles, nodeRoot2 and nodeFoo2 but not nodeBar (cached). + verify(service, times(5 + 4)).uncachedBuildMerkleTreeVisitor(any(), any()); + } + private Spawn newSpawnFromResult(RemoteActionResult result) { return newSpawnFromResult(ImmutableMap.of(), result); }