From b024eea67edd76bcffec3550bc88c991859e6500 Mon Sep 17 00:00:00 2001 From: Googler Date: Thu, 5 Dec 2024 13:05:13 -0800 Subject: [PATCH] Widen interface for action input prefetching. RELNOTES: None PiperOrigin-RevId: 703221309 Change-Id: Ie7a4607b530a01e49f671bc6804a307740c81175 --- .../lib/actions/ActionInputPrefetcher.java | 19 ++----- .../build/lib/exec/AbstractSpawnStrategy.java | 2 +- .../remote/AbstractActionInputPrefetcher.java | 38 ++++++++++++- .../lib/remote/RemoteActionFileSystem.java | 5 +- .../lib/skyframe/CompletionFunction.java | 5 +- .../remote/ActionInputPrefetcherTestBase.java | 56 +++++++++++-------- .../remote/RemoteActionInputFetcherTest.java | 6 +- 7 files changed, 82 insertions(+), 49 deletions(-) diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java b/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java index 9ac0461ba03cce..c237d171f83a93 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java @@ -20,21 +20,10 @@ /** Prefetches files to local disk. */ public interface ActionInputPrefetcher { - /** - * Returns the metadata for an {@link ActionInput}. - * - *

This will generally call through to a {@link InputMetadataProvider} and ask for the metadata - * of either an input or an output artifact. - */ - public interface MetadataSupplier { - FileArtifactValue getMetadata(ActionInput actionInput) throws IOException, InterruptedException; - } - public static final ActionInputPrefetcher NONE = - (action, inputs, metadataSupplier, priority) -> { - // Do nothing. - return immediateVoidFuture(); - }; + (action, inputs, metadataProvider, priority) -> + // Do nothing. + immediateVoidFuture(); /** Priority for the staging task. */ public enum Priority { @@ -70,6 +59,6 @@ public enum Priority { ListenableFuture prefetchFiles( ActionExecutionMetadata action, Iterable inputs, - MetadataSupplier metadataSupplier, + InputMetadataProvider metadataProvider, Priority priority); } diff --git a/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java index c689df66c18a21..9618b5f4b3aa74 100644 --- a/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java +++ b/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java @@ -289,7 +289,7 @@ public ListenableFuture prefetchInputs() throws ForbiddenActionInputExcept spawn.getResourceOwner(), getInputMapping(PathFragment.EMPTY_FRAGMENT, /* willAccessRepeatedly= */ true) .values(), - getInputMetadataProvider()::getInputMetadata, + getInputMetadataProvider(), Priority.MEDIUM), BulkTransferException.class, (BulkTransferException e) -> { diff --git a/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java b/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java index 692ba9bfbcd768..b965aeb5ddd10f 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java +++ b/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java @@ -40,6 +40,7 @@ import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact; import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact; import com.google.devtools.build.lib.actions.FileArtifactValue; +import com.google.devtools.build.lib.actions.InputMetadataProvider; import com.google.devtools.build.lib.actions.cache.OutputMetadataStore; import com.google.devtools.build.lib.actions.cache.VirtualActionInput; import com.google.devtools.build.lib.events.Reporter; @@ -89,6 +90,17 @@ enum DirectoryState { OUTPUT_PERMISSIONS } + /** + * Returns the metadata for an {@link ActionInput}. + * + *

This will generally call through to a {@link InputMetadataProvider} and ask for the metadata + * of either an input or an output artifact. + */ + @VisibleForTesting + public interface MetadataSupplier { + FileArtifactValue getMetadata(ActionInput actionInput) throws IOException, InterruptedException; + } + /** * Tracks directory permissions to minimize filesystem operations. * @@ -260,6 +272,30 @@ protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOExc */ @Override public ListenableFuture prefetchFiles( + ActionExecutionMetadata action, + Iterable inputs, + InputMetadataProvider metadataProvider, + Priority priority) { + return prefetchFilesInterruptibly(action, inputs, metadataProvider::getInputMetadata, priority); + } + + /** + * Fetches remotely stored action outputs, that are inputs to this spawn, and stores them under + * their path in the output base. + * + *

The {@code inputs} may not contain any unexpanded directories. + * + *

This method is safe to be called concurrently from spawn runners before running any local + * spawn. + * + *

This method is similar to #prefetchFiles() above, but note that {@code metadataSupplier} may + * throw {@link InterruptedException}. If it does, this method will propagate this exception in + * the returned future. + * + * @return a future that is completed once all downloads have finished. + */ + @VisibleForTesting + public ListenableFuture prefetchFilesInterruptibly( ActionExecutionMetadata action, Iterable inputs, MetadataSupplier metadataSupplier, @@ -686,7 +722,7 @@ public void finalizeAction(Action action, OutputMetadataStore outputMetadataStor if (!outputsToDownload.isEmpty()) { try (var s = Profiler.instance().profile(ProfilerTask.REMOTE_DOWNLOAD, "Download outputs")) { getFromFuture( - prefetchFiles( + prefetchFilesInterruptibly( action, outputsToDownload, outputMetadataStore::getOutputMetadata, Priority.HIGH)); } } diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java index 55cad4046c7df1..6e3dbfd43361ae 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java @@ -735,8 +735,7 @@ ActionInput getInput(String execPath) { @Nullable @VisibleForTesting FileArtifactValue getInputMetadata(ActionInput input) { - PathFragment execPath = input.getExecPath(); - return inputArtifactData.getMetadata(execPath); + return inputArtifactData.getInputMetadata(input); } private void downloadFileIfRemote(PathFragment path) throws IOException { @@ -753,7 +752,7 @@ private void downloadFileIfRemote(PathFragment path) throws IOException { } getFromFuture( inputFetcher.prefetchFiles( - action, ImmutableList.of(input), this::getInputMetadata, Priority.CRITICAL)); + action, ImmutableList.of(input), inputArtifactData, Priority.CRITICAL)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IOException(String.format("Received interrupt while fetching file '%s'", path), e); diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java index 79db0fbbf01afe..6579afe958264a 100644 --- a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java +++ b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java @@ -468,8 +468,7 @@ private void downloadArtifact( var action = ActionUtils.getActionForLookupData(env, derivedArtifact.getGeneratingActionKey()); var future = - actionInputPrefetcher.prefetchFiles( - action, filesToDownload, inputMap::getInputMetadata, Priority.LOW); + actionInputPrefetcher.prefetchFiles(action, filesToDownload, inputMap, Priority.LOW); futures.add(future); } } else { @@ -485,7 +484,7 @@ private void downloadArtifact( ActionUtils.getActionForLookupData(env, derivedArtifact.getGeneratingActionKey()); var future = actionInputPrefetcher.prefetchFiles( - action, ImmutableList.of(artifact), inputMap::getInputMetadata, Priority.LOW); + action, ImmutableList.of(artifact), inputMap, Priority.LOW); futures.add(future); } } diff --git a/src/test/java/com/google/devtools/build/lib/remote/ActionInputPrefetcherTestBase.java b/src/test/java/com/google/devtools/build/lib/remote/ActionInputPrefetcherTestBase.java index eef215f0e4de98..88af250a24506b 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/ActionInputPrefetcherTestBase.java +++ b/src/test/java/com/google/devtools/build/lib/remote/ActionInputPrefetcherTestBase.java @@ -45,7 +45,6 @@ import com.google.common.util.concurrent.SettableFuture; import com.google.devtools.build.lib.actions.ActionExecutionMetadata; import com.google.devtools.build.lib.actions.ActionInput; -import com.google.devtools.build.lib.actions.ActionInputPrefetcher.MetadataSupplier; import com.google.devtools.build.lib.actions.ActionInputPrefetcher.Priority; import com.google.devtools.build.lib.actions.Artifact; import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact; @@ -56,6 +55,7 @@ import com.google.devtools.build.lib.actions.FileArtifactValue; import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue; import com.google.devtools.build.lib.actions.util.ActionsTestUtil; +import com.google.devtools.build.lib.remote.AbstractActionInputPrefetcher.MetadataSupplier; import com.google.devtools.build.lib.remote.common.LostInputsEvent; import com.google.devtools.build.lib.skyframe.TreeArtifactValue; import com.google.devtools.build.lib.testing.vfs.SpiedFileSystem; @@ -296,7 +296,9 @@ public void prefetchFiles_fileExists_doNotDownload() FileSystemUtils.writeContent(a.getPath(), "hello world".getBytes(UTF_8)); AbstractActionInputPrefetcher prefetcher = spy(createPrefetcher(cas)); - wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM)); + wait( + prefetcher.prefetchFilesInterruptibly( + action, metadata.keySet(), metadata::get, Priority.MEDIUM)); verify(prefetcher, never()).doDownloadFile(eq(action), any(), any(), any(), any(), any()); assertThat(prefetcher.downloadedFiles()).containsExactly(a.getPath()); @@ -312,7 +314,9 @@ public void prefetchFiles_fileExistsButContentMismatches_download() FileSystemUtils.writeContent(a.getPath(), "hello world local".getBytes(UTF_8)); AbstractActionInputPrefetcher prefetcher = spy(createPrefetcher(cas)); - wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM)); + wait( + prefetcher.prefetchFilesInterruptibly( + action, metadata.keySet(), metadata::get, Priority.MEDIUM)); verify(prefetcher).doDownloadFile(eq(action), any(), any(), eq(a.getExecPath()), any(), any()); assertThat(prefetcher.downloadedFiles()).containsExactly(a.getPath()); @@ -328,7 +332,9 @@ public void prefetchFiles_downloadRemoteFiles() throws Exception { Artifact a2 = createRemoteArtifact("file2", "fizz buzz", metadata, cas); AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas); - wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM)); + wait( + prefetcher.prefetchFilesInterruptibly( + action, metadata.keySet(), metadata::get, Priority.MEDIUM)); assertThat(FileSystemUtils.readContent(a1.getPath(), UTF_8)).isEqualTo("hello world"); assertReadableNonWritableAndExecutable(a1.getPath()); @@ -346,7 +352,9 @@ public void prefetchFiles_downloadRemoteFiles_withMaterializationExecPath() thro Artifact a = createRemoteArtifact("file", "hello world", targetExecPath, metadata, cas); AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas); - wait(prefetcher.prefetchFiles(action, metadata.keySet(), metadata::get, Priority.MEDIUM)); + wait( + prefetcher.prefetchFilesInterruptibly( + action, metadata.keySet(), metadata::get, Priority.MEDIUM)); assertThat(a.getPath().isSymbolicLink()).isTrue(); assertThat(a.getPath().readSymbolicLink()) @@ -376,7 +384,7 @@ public void prefetchFiles_downloadRemoteTrees() throws Exception { AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas); - wait(prefetcher.prefetchFiles(action, children, metadata::get, Priority.MEDIUM)); + wait(prefetcher.prefetchFilesInterruptibly(action, children, metadata::get, Priority.MEDIUM)); assertThat(FileSystemUtils.readContent(firstChild.getPath(), UTF_8)).isEqualTo("content1"); assertThat(FileSystemUtils.readContent(secondChild.getPath(), UTF_8)).isEqualTo("content2"); @@ -407,7 +415,7 @@ public void prefetchFiles_downloadRemoteTrees_partial() throws Exception { AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas); wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(firstChild, secondChild), metadata::get, Priority.MEDIUM)); assertThat(firstChild.getPath().exists()).isFalse(); @@ -437,7 +445,7 @@ public void prefetchFiles_downloadRemoteTrees_withMaterializationExecPath() thro AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas); - wait(prefetcher.prefetchFiles(action, children, metadata::get, Priority.MEDIUM)); + wait(prefetcher.prefetchFilesInterruptibly(action, children, metadata::get, Priority.MEDIUM)); assertThat(tree.getPath().isSymbolicLink()).isTrue(); assertThat(tree.getPath().readSymbolicLink()) @@ -475,13 +483,13 @@ public void prefetchFiles_downloadRemoteTrees_forActionTemplateExpansion() throw AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas); wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(firstChild), metadata::get, Priority.MEDIUM)); assertTreeReadableWritableAndExecutable(tree.getPath()); wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(secondChild), metadata::get, Priority.MEDIUM)); assertTreeReadableWritableAndExecutable(tree.getPath()); @@ -497,7 +505,7 @@ public void prefetchFiles_missingFiles_fails() throws Exception { Exception.class, () -> wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(a), metadata::get, Priority.MEDIUM))); assertThat(prefetcher.downloadedFiles()).isEmpty(); @@ -515,7 +523,9 @@ public void prefetchFiles_ignoreNonRemoteFiles() throws Exception { ImmutableMap metadata = ImmutableMap.of(a, f); AbstractActionInputPrefetcher prefetcher = createPrefetcher(new HashMap<>()); - wait(prefetcher.prefetchFiles(action, ImmutableList.of(a), metadata::get, Priority.MEDIUM)); + wait( + prefetcher.prefetchFilesInterruptibly( + action, ImmutableList.of(a), metadata::get, Priority.MEDIUM)); assertThat(prefetcher.downloadedFiles()).isEmpty(); assertThat(prefetcher.downloadsInProgress()).isEmpty(); @@ -541,7 +551,7 @@ public void prefetchFiles_ignoreNonRemoteFiles_tree() throws Exception { AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas); - wait(prefetcher.prefetchFiles(action, children, metadata::get, Priority.MEDIUM)); + wait(prefetcher.prefetchFilesInterruptibly(action, children, metadata::get, Priority.MEDIUM)); assertThat(firstChild.getPath().exists()).isFalse(); assertThat(FileSystemUtils.readContent(secondChild.getPath(), UTF_8)).isEqualTo("content2"); @@ -573,7 +583,7 @@ public void prefetchFiles_treeFiles_minimizeFilesystemOperations() throws Except reset(fs); wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(firstChild, secondChild), metadata::get, Priority.MEDIUM)); verify(fs).createWritableDirectory(root); @@ -584,7 +594,7 @@ public void prefetchFiles_treeFiles_minimizeFilesystemOperations() throws Except reset(fs); wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(firstChild, secondChild), metadata::get, Priority.MEDIUM)); verify(fs, never()).createWritableDirectory(root); @@ -623,7 +633,7 @@ public void prefetchFiles_treeFiles_multipleThreads_waitForPermissionsToBeSet() Callable prefetch = () -> { wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(child), metadata::get, Priority.MEDIUM)); assertTreeReadableNonWritableAndExecutable(tree.getPath()); return null; @@ -657,7 +667,7 @@ public void prefetchFiles_multipleThreads_downloadIsCancelled() throws Exception () -> { try { wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM)); } catch (IOException | ExecException | InterruptedException ignored) { // do nothing @@ -669,7 +679,7 @@ public void prefetchFiles_multipleThreads_downloadIsCancelled() throws Exception () -> { try { wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM)); } catch (IOException | ExecException | InterruptedException ignored) { // do nothing @@ -707,7 +717,7 @@ public void prefetchFiles_multipleThreads_downloadIsNotCancelledByOtherThreads() () -> { try { wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM)); } catch (IOException | ExecException | InterruptedException ignored) { // do nothing @@ -720,7 +730,7 @@ public void prefetchFiles_multipleThreads_downloadIsNotCancelledByOtherThreads() () -> { try { wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(artifact), metadata::get, Priority.MEDIUM)); successful.set(true); } catch (IOException | ExecException | InterruptedException ignored) { @@ -764,7 +774,7 @@ public void prefetchFile_interruptingMetadataSupplier_interruptsDownload() throw }; ListenableFuture future = - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(a1), interruptedMetadataSupplier, Priority.MEDIUM); assertThrows(InterruptedException.class, () -> getFromFuture(future)); @@ -792,7 +802,7 @@ public void prefetchFiles_onInterrupt_deletePartialDownloadedFile() throws Excep () -> { try { getFromFuture( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(a1), metadata::get, Priority.MEDIUM)); } catch (IOException ignored) { // Intentionally left empty @@ -831,7 +841,7 @@ public void onLostInputsEvent(LostInputsEvent event) { Exception.class, () -> wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, metadata.keySet(), metadata::get, Priority.MEDIUM))); assertThat(lostInputsEvents).hasSize(1); diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java index ccf5fe551b0b19..dc9515489e6368 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcherTest.java @@ -102,7 +102,7 @@ public void testStagingVirtualActionInput() throws Exception { // act wait( - actionInputFetcher.prefetchFiles( + actionInputFetcher.prefetchFilesInterruptibly( action, ImmutableList.of(a), (ActionInput unused) -> null, Priority.MEDIUM)); // assert @@ -131,7 +131,7 @@ public void testStagingEmptyVirtualActionInput() throws Exception { // act wait( - actionInputFetcher.prefetchFiles( + actionInputFetcher.prefetchFilesInterruptibly( action, ImmutableList.of(VirtualActionInput.EMPTY_MARKER), (ActionInput unused) -> null, @@ -153,7 +153,7 @@ public void prefetchFiles_missingFiles_failsWithSpecificMessage() throws Excepti BulkTransferException.class, () -> wait( - prefetcher.prefetchFiles( + prefetcher.prefetchFilesInterruptibly( action, ImmutableList.of(a), metadata::get, Priority.MEDIUM))); assertThat(prefetcher.downloadedFiles()).isEmpty();