Fix data race on prefetcher.

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

@@ -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;
@@ -87,53 +86,132 @@ 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();
@@ -428,14 +506,14 @@ private Completable downloadFileNoCheckRx(
                 tempPathGenerator::generateTempPath,
                 tempPath ->
                     toCompletable(
-                            () ->
-                                doDownloadFile(
-                                    reporter,
-                                    tempPath,
-                                    finalPath.relativeTo(execRoot),
-                                    metadata,
-                                    priority),
-                            directExecutor())
+                        () ->
+                            doDownloadFile(
+                                reporter,
+                                tempPath,
+                                finalPath.relativeTo(execRoot),
+                                metadata,
+                                priority),
+                        directExecutor())
                         .doOnComplete(
                             () -> {
                               finalizeDownload(dirCtx, treeRoot, tempPath, finalPath);
@@ -520,24 +598,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());

src/main/java/com/google/devtools/build/lib/testing/vfs/SpiedFileSystem.java

@@ -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);
+  }
 }

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

@@ -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;
@@ -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;
@@ -68,18 +67,21 @@
 import org.junit.Before;
 import org.junit.Test;
 
-/** Base test class for {@link AbstractActionInputPrefetcher} implementations. */
+/**
+ * Base test class for {@link AbstractActionInputPrefetcher} implementations.
+ */
 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");
@@ -421,6 +423,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

src/test/java/com/google/devtools/build/lib/remote/BUILD

@@ -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",