Skip to content

Commit

Permalink
Fix data race in prefetcher.
Browse files Browse the repository at this point in the history 
Individual file prefetches within a single prefetchFiles() can race against each other, so they must synchronize when writing to the DirectoryContext.

Closes bazelbuild#17678.

PiperOrigin-RevId: 515024483
Change-Id: Ic8097979d06ab143b4d63f5e90f871f8cbf83959
  • Loading branch information
@tjgq
tjgq authored and Copybara-Service committed Mar 8, 2023
1 parent f9008f6 commit e0cdace
Show file tree
Hide file tree
Showing 4 changed files with 140 additions and 45 deletions.
145 changes: 105 additions & 40 deletions src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Deque;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
Expand All @@ -87,53 +86,124 @@ public abstract class AbstractActionInputPrefetcher implements ActionInputPrefet

private final Set<ActionInput> missingActionInputs = Sets.newConcurrentHashSet();

// Tracks the number of ongoing prefetcher calls temporarily making an output directory writable.
// Since concurrent calls may write to the same directory, it's not safe to make it non-writable
// until no other ongoing calls are writing to it.
private final ConcurrentHashMap<Path, Integer> temporarilyWritableDirectories =
private static final Object dummyValue = new Object();

/**
* Tracks output directories temporarily made writable for prefetching. Since concurrent calls may
* write to the same directory, it's not safe to make it non-writable until no other ongoing
* prefetcher calls are writing to it.
*/
private final ConcurrentHashMap<Path, DirectoryState> temporarilyWritableDirectories =
new ConcurrentHashMap<>();

/** Keeps track of output directories written to by a single prefetcher call. */
/** The state of a single temporarily writable directory. */
private static final class DirectoryState {
/** The number of ongoing prefetcher calls touching this directory. */
int numCalls;
/** Whether the output permissions must be set on the directory when prefetching completes. */
boolean mustSetOutputPermissions;
}

/**
* Tracks output directories written to by a single prefetcher call.
*
* <p>This makes it possible to set the output permissions on directories touched by the
* prefetcher call all at once, so that files prefetched within the same call don't repeatedly set
* output permissions on the same directory.
*/
private final class DirectoryContext {
private final HashSet<Path> dirs = new HashSet<>();
private final ConcurrentHashMap<Path, Object> dirs = new ConcurrentHashMap<>();

/**
* Adds to the set of directories written to by the prefetcher call associated with this
* context.
* Makes a directory temporarily writable for the remainder of the prefetcher call associated
* with this context.
*
* @param isDefinitelyTreeDir Whether this directory definitely belongs to a tree artifact.
* Otherwise, whether it belongs to a tree artifact is inferred from its permissions.
*/
void add(Path dir) {
if (dirs.add(dir)) {
temporarilyWritableDirectories.compute(dir, (unused, count) -> count != null ? ++count : 1);
void createOrSetWritable(Path dir, boolean isDefinitelyTreeDir) throws IOException {
AtomicReference<IOException> caughtException = new AtomicReference<>();

dirs.compute(
dir,
(outerUnused, previousValue) -> {
if (previousValue != null) {
return previousValue;
}

temporarilyWritableDirectories.compute(
dir,
(innerUnused, state) -> {
if (state == null) {
state = new DirectoryState();
state.numCalls = 0;

try {
if (isDefinitelyTreeDir) {
state.mustSetOutputPermissions = true;
var ignored = dir.createWritableDirectory();
} else {
// If the directory is writable, it's a package and should be kept writable.
// Otherwise, it must belong to a tree artifact, since the directory for a
// tree is created in a non-writable state before prefetching begins, and
// this is the first time the prefetcher is seeing it.
state.mustSetOutputPermissions = !dir.isWritable();
if (state.mustSetOutputPermissions) {
dir.setWritable(true);
}
}
} catch (IOException e) {
caughtException.set(e);
return null;
}
}

++state.numCalls;

return state;
});

if (caughtException.get() != null) {
return null;
}

return dummyValue;
});

if (caughtException.get() != null) {
throw caughtException.get();
}
}

/**
* Signals that the prefetcher call associated with this context has finished.
*
* <p>The output permissions will be set on any directories written to by this call that are not
* being written to by other concurrent calls.
* <p>The output permissions will be set on any directories temporarily made writable by this
* call, if this is the last remaining call temporarily making them writable.
*/
void close() throws IOException {
AtomicReference<IOException> caughtException = new AtomicReference<>();

for (Path dir : dirs) {
for (Path dir : dirs.keySet()) {
temporarilyWritableDirectories.compute(
dir,
(unused, count) -> {
checkState(count != null);
if (--count == 0) {
try {
dir.chmod(outputPermissions.getPermissionsMode());
} catch (IOException e) {
// Store caught exceptions, but keep cleaning up the map.
if (caughtException.get() == null) {
caughtException.set(e);
} else {
caughtException.get().addSuppressed(e);
(unused, state) -> {
checkState(state != null);
if (--state.numCalls == 0) {
if (state.mustSetOutputPermissions) {
try {
dir.chmod(outputPermissions.getPermissionsMode());
} catch (IOException e) {
// Store caught exceptions, but keep cleaning up the map.
if (caughtException.get() == null) {
caughtException.set(e);
} else {
caughtException.get().addSuppressed(e);
}
}
}
}
return count > 0 ? count : null;
return state.numCalls > 0 ? state : null;
});
}
dirs.clear();
Expand Down Expand Up @@ -520,24 +590,19 @@ private void finalizeDownload(
}
while (!dirs.isEmpty()) {
Path dir = dirs.pop();
dirCtx.add(dir);
// Create directory or make existing directory writable.
var unused = dir.createWritableDirectory();
// We know with certainty that the directory belongs to a tree artifact.
dirCtx.createOrSetWritable(dir, /* isDefinitelyTreeDir= */ true);
}
} else {
// If the parent directory is not writable, temporarily make it so.
// This is needed when fetching a non-tree artifact nested inside a tree artifact, or a tree
// artifact inside a fileset (see b/254844173 for the latter).
// TODO(tjgq): Fix the TOCTTOU race between isWritable and setWritable. This requires keeping
// track of the original directory permissions. Note that nested artifacts are relatively rare
// and will eventually be disallowed (see issue #16729).
if (!parentDir.isWritable()) {
dirCtx.add(parentDir);
parentDir.setWritable(true);
}
// Temporarily make the parent directory writable if needed.
// We don't know with certainty that the directory does not belong to a tree artifact; it
// could if the fetched file is a non-tree artifact nested inside a tree artifact, or a
// tree artifact inside a fileset (see b/254844173 for the latter).
dirCtx.createOrSetWritable(parentDir, /* isDefinitelyTreeDir= */ false);
}

// Set output permissions on files (tree subdirectories are handled in stopPrefetching),
// Set output permissions on files (tree subdirectories are handled in DirectoryContext#close),
// matching the behavior of SkyframeActionExecutor#checkOutputs for artifacts produced by local
// actions.
tmpPath.chmod(outputPermissions.getPermissionsMode());
Expand Down
5 changes: 5 additions & 0 deletions src/main/java/com/google/devtools/build/lib/testing/vfs/SpiedFileSystem.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,4 +50,9 @@ public static SpiedFileSystem createInMemorySpy() {
public OutputStream getOutputStream(PathFragment path, boolean append) throws IOException {
return super.getOutputStream(path, append);
}

@Override
public boolean createWritableDirectory(PathFragment path) throws IOException {
return super.createWritableDirectory(path);
}
}
34 changes: 29 additions & 5 deletions src/test/java/com/google/devtools/build/lib/remote/ActionInputPrefetcherTestBase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

import com.google.common.collect.ImmutableList;
Expand All @@ -44,19 +45,17 @@
import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
import com.google.devtools.build.lib.actions.MetadataProvider;
import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
import com.google.devtools.build.lib.clock.JavaClock;
import com.google.devtools.build.lib.remote.util.StaticMetadataProvider;
import com.google.devtools.build.lib.remote.util.TempPathGenerator;
import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
import com.google.devtools.build.lib.testing.vfs.SpiedFileSystem;
import com.google.devtools.build.lib.util.Pair;
import com.google.devtools.build.lib.vfs.DigestHashFunction;
import com.google.devtools.build.lib.vfs.Dirent;
import com.google.devtools.build.lib.vfs.FileSystem;
import com.google.devtools.build.lib.vfs.FileSystemUtils;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.build.lib.vfs.PathFragment;
import com.google.devtools.build.lib.vfs.Symlinks;
import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
Expand All @@ -72,14 +71,14 @@
public abstract class ActionInputPrefetcherTestBase {
protected static final DigestHashFunction HASH_FUNCTION = DigestHashFunction.SHA256;

protected FileSystem fs;
protected SpiedFileSystem fs;
protected Path execRoot;
protected ArtifactRoot artifactRoot;
protected TempPathGenerator tempPathGenerator;

@Before
public void setUp() throws IOException {
fs = new InMemoryFileSystem(new JavaClock(), HASH_FUNCTION);
fs = SpiedFileSystem.createInMemorySpy();
execRoot = fs.getPath("/exec");
execRoot.createDirectoryAndParents();
artifactRoot = ArtifactRoot.asDerivedRoot(execRoot, RootType.Output, "root");
Expand Down Expand Up @@ -425,6 +424,31 @@ public void prefetchFiles_ignoreNonRemoteFiles_tree() throws Exception {
assertThat(prefetcher.downloadedFiles()).containsExactly(secondChild.getPath());
}

@Test
public void prefetchFiles_treeFiles_minimizeFilesystemOperations() throws Exception {
Map<ActionInput, FileArtifactValue> metadata = new HashMap<>();
Map<HashCode, byte[]> cas = new HashMap<>();
Pair<SpecialArtifact, ImmutableList<TreeFileArtifact>> treeAndChildren =
createRemoteTreeArtifact(
"dir",
/* localContentMap= */ ImmutableMap.of("subdir/file1", "content1"),
/* remoteContentMap= */ ImmutableMap.of("subdir/file2", "content2"),
metadata,
cas);
SpecialArtifact tree = treeAndChildren.getFirst();
ImmutableList<TreeFileArtifact> children = treeAndChildren.getSecond();
Artifact firstChild = children.get(0);
Artifact secondChild = children.get(1);

MetadataProvider metadataProvider = new StaticMetadataProvider(metadata);
AbstractActionInputPrefetcher prefetcher = createPrefetcher(cas);

wait(prefetcher.prefetchFiles(ImmutableList.of(firstChild, secondChild), metadataProvider));

verify(fs, times(1)).createWritableDirectory(tree.getPath().asFragment());
verify(fs, times(1)).createWritableDirectory(tree.getPath().getChild("subdir").asFragment());
}

@Test
public void prefetchFiles_multipleThreads_downloadIsCancelled() throws Exception {
// Test shared downloads are cancelled if all threads/callers are interrupted
Expand Down
1 change: 1 addition & 0 deletions src/test/java/com/google/devtools/build/lib/remote/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ java_test(
"//src/main/java/com/google/devtools/build/lib/remote/util",
"//src/main/java/com/google/devtools/build/lib/runtime/commands",
"//src/main/java/com/google/devtools/build/lib/skyframe:tree_artifact_value",
"//src/main/java/com/google/devtools/build/lib/testing/vfs:spied_filesystem",
"//src/main/java/com/google/devtools/build/lib/util",
"//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
"//src/main/java/com/google/devtools/build/lib/util:exit_code",
Expand Down

0 comments on commit e0cdace

Please sign in to comment.