Make WorkerExecRoot not be re-created on each createFileSystem() call. Preparation for holding a map of existing links, but also just nicer.

RELNOTES: None
PiperOrigin-RevId: 356465646

Author: larsrc-google

src/main/java/com/google/devtools/build/lib/worker/SandboxedWorker.java

@@ -25,31 +25,27 @@
 /** A {@link SingleplexWorker} that runs inside a sandboxed execution root. */
 final class SandboxedWorker extends SingleplexWorker {
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
-  private WorkerExecRoot workerExecRoot;
+  private final WorkerExecRoot workerExecRoot;
 
   SandboxedWorker(WorkerKey workerKey, int workerId, Path workDir, Path logFile) {
     super(workerKey, workerId, workDir, logFile);
+    workerExecRoot = new WorkerExecRoot(workDir);
   }
 
   @Override
   public void prepareExecution(
       SandboxInputs inputFiles, SandboxOutputs outputs, Set<PathFragment> workerFiles)
       throws IOException {
-    // Note that workerExecRoot isn't necessarily null at this point, so we can't do a Preconditions
-    // check for it: If a WorkerSpawnStrategy gets interrupted, finishExecution is not guaranteed to
-    // be called.
-    workerExecRoot = new WorkerExecRoot(workDir, inputFiles, outputs, workerFiles);
-    workerExecRoot.createFileSystem();
+    workerExecRoot.createFileSystem(workerFiles, inputFiles, outputs);
 
     super.prepareExecution(inputFiles, outputs, workerFiles);
   }
 
   @Override
-  public void finishExecution(Path execRoot) throws IOException {
-    super.finishExecution(execRoot);
+  public void finishExecution(Path execRoot, SandboxOutputs outputs) throws IOException {
+    super.finishExecution(execRoot, outputs);
 
-    workerExecRoot.copyOutputs(execRoot);
-    workerExecRoot = null;
+    workerExecRoot.copyOutputs(execRoot, outputs);
   }
 
   @Override

src/main/java/com/google/devtools/build/lib/worker/SingleplexWorker.java

@@ -124,7 +124,7 @@ WorkResponse getResponse(int requestId) throws IOException {
   }
 
   @Override
-  public void finishExecution(Path execRoot) throws IOException {}
+  public void finishExecution(Path execRoot, SandboxOutputs outputs) throws IOException {}
 
   @Override
   void destroy() {

src/main/java/com/google/devtools/build/lib/worker/Worker.java

@@ -101,7 +101,7 @@ public abstract void prepareExecution(
   abstract WorkResponse getResponse(int requestId) throws IOException, InterruptedException;
 
   /** Does whatever cleanup may be required after execution is done. */
-  public abstract void finishExecution(Path execRoot) throws IOException;
+  public abstract void finishExecution(Path execRoot, SandboxOutputs outputs) throws IOException;
 
   /**
    * Destroys this worker. Once this has been called, we assume it's safe to clean up related

src/main/java/com/google/devtools/build/lib/worker/WorkerExecRoot.java

@@ -30,42 +30,41 @@
 /** Creates and manages the contents of a working directory of a persistent worker. */
 final class WorkerExecRoot {
   private final Path workDir;
-  private final SandboxInputs inputs;
-  private final SandboxOutputs outputs;
-  private final Set<PathFragment> workerFiles;
 
-  public WorkerExecRoot(
-      Path workDir, SandboxInputs inputs, SandboxOutputs outputs, Set<PathFragment> workerFiles) {
+  public WorkerExecRoot(Path workDir) {
     this.workDir = workDir;
-    this.inputs = inputs;
-    this.outputs = outputs;
-    this.workerFiles = workerFiles;
   }
 
-  public void createFileSystem() throws IOException {
+  public void createFileSystem(
+      Set<PathFragment> workerFiles, SandboxInputs inputs, SandboxOutputs outputs)
+      throws IOException {
     workDir.createDirectoryAndParents();
 
     // First compute all the inputs and directories that we need. This is based only on
     // `workerFiles`, `inputs` and `outputs` and won't do any I/O.
     Set<PathFragment> inputsToCreate = new LinkedHashSet<>();
     LinkedHashSet<PathFragment> dirsToCreate = new LinkedHashSet<>();
-    populateInputsAndDirsToCreate(inputsToCreate, dirsToCreate);
+    populateInputsAndDirsToCreate(inputs, workerFiles, outputs, inputsToCreate, dirsToCreate);
 
     // Then do a full traversal of the `workDir`. This will use what we computed above, delete
     // anything unnecessary and update `inputsToCreate`/`dirsToCreate` if something is can be left
     // without changes (e.g., a symlink that already points to the right destination).
-    cleanExisting(workDir, inputsToCreate, dirsToCreate);
+    cleanExisting(workDir, inputs, inputsToCreate, dirsToCreate);
 
     // Finally, create anything that is still missing.
     createDirectories(dirsToCreate);
-    createInputs(inputsToCreate);
+    createInputs(inputsToCreate, inputs);
 
     inputs.materializeVirtualInputs(workDir);
   }
 
   /** Populates the provided sets with the inputs and directories than need to be created. */
   private void populateInputsAndDirsToCreate(
-      Set<PathFragment> inputsToCreate, LinkedHashSet<PathFragment> dirsToCreate) {
+      SandboxInputs inputs,
+      Set<PathFragment> workerFiles,
+      SandboxOutputs outputs,
+      Set<PathFragment> inputsToCreate,
+      LinkedHashSet<PathFragment> dirsToCreate) {
     // Add all worker files and the ancestor directories.
     for (PathFragment path : workerFiles) {
       inputsToCreate.add(path);
@@ -101,12 +100,16 @@ private void populateInputsAndDirsToCreate(
    * correct and doesn't need any changes.
    */
   private void cleanExisting(
-      Path root, Set<PathFragment> inputsToCreate, Set<PathFragment> dirsToCreate)
+      Path root,
+      SandboxInputs inputs,
+      Set<PathFragment> inputsToCreate,
+      Set<PathFragment> dirsToCreate)
       throws IOException {
     for (Path path : root.getDirectoryEntries()) {
       FileStatus stat = path.stat(Symlinks.NOFOLLOW);
       PathFragment pathRelativeToWorkDir = path.relativeTo(workDir);
-      Optional<PathFragment> destination = getExpectedSymlinkDestination(pathRelativeToWorkDir);
+      Optional<PathFragment> destination =
+          getExpectedSymlinkDestination(pathRelativeToWorkDir, inputs);
       if (destination.isPresent()) {
         if (stat.isSymbolicLink() && path.readSymbolicLink().equals(destination.get())) {
           inputsToCreate.remove(pathRelativeToWorkDir);
@@ -115,7 +118,7 @@ private void cleanExisting(
         }
       } else if (stat.isDirectory()) {
         if (dirsToCreate.contains(pathRelativeToWorkDir)) {
-          cleanExisting(path, inputsToCreate, dirsToCreate);
+          cleanExisting(path, inputs, inputsToCreate, dirsToCreate);
           dirsToCreate.remove(pathRelativeToWorkDir);
         } else {
           path.deleteTree();
@@ -126,7 +129,8 @@ private void cleanExisting(
     }
   }
 
-  private Optional<PathFragment> getExpectedSymlinkDestination(PathFragment fragment) {
+  private Optional<PathFragment> getExpectedSymlinkDestination(
+      PathFragment fragment, SandboxInputs inputs) {
     Path file = inputs.getFiles().get(fragment);
     if (file != null) {
       return Optional.of(file.asFragment());
@@ -140,7 +144,8 @@ private void createDirectories(Iterable<PathFragment> dirsToCreate) throws IOExc
     }
   }
 
-  private void createInputs(Iterable<PathFragment> inputsToCreate) throws IOException {
+  private void createInputs(Iterable<PathFragment> inputsToCreate, SandboxInputs inputs)
+      throws IOException {
     for (PathFragment fragment : inputsToCreate) {
       Path key = workDir.getRelative(fragment);
       if (inputs.getFiles().containsKey(fragment)) {
@@ -159,7 +164,7 @@ private void createInputs(Iterable<PathFragment> inputsToCreate) throws IOExcept
     }
   }
 
-  public void copyOutputs(Path execRoot) throws IOException {
+  public void copyOutputs(Path execRoot, SandboxOutputs outputs) throws IOException {
     SandboxHelpers.moveOutputs(outputs, workDir, execRoot);
   }
 }

src/main/java/com/google/devtools/build/lib/worker/WorkerProxy.java

@@ -94,7 +94,7 @@ WorkResponse getResponse(int requestId) throws InterruptedException {
   }
 
   @Override
-  public void finishExecution(Path execRoot) {}
+  public void finishExecution(Path execRoot, SandboxOutputs outputs) {}
 
   @Override
   boolean diedUnexpectedly() {

src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnRunner.java

@@ -483,7 +483,7 @@ WorkResponse execInWorker(
       try {
         Stopwatch processOutputsStopwatch = Stopwatch.createStarted();
         context.lockOutputFiles();
-        worker.finishExecution(execRoot);
+        worker.finishExecution(execRoot, outputs);
         spawnMetrics.setProcessOutputsTime(processOutputsStopwatch.elapsed());
       } catch (IOException e) {
         restoreInterrupt(e);

src/test/java/com/google/devtools/build/lib/worker/WorkerExecRootTest.java

@@ -62,17 +62,17 @@ public void cleanFileSystem() throws Exception {
     Path workerSh = workspaceDir.getRelative("worker.sh");
     FileSystemUtils.writeContentAsLatin1(workerSh, "#!/bin/bash");
 
-    WorkerExecRoot workerExecRoot =
-        new WorkerExecRoot(
-            execRoot,
-            new SandboxInputs(
-                ImmutableMap.of(PathFragment.create("worker.sh"), workerSh),
-                ImmutableSet.of(),
-                ImmutableMap.of()),
-            SandboxOutputs.create(
-                ImmutableSet.of(PathFragment.create("very/output.txt")), ImmutableSet.of()),
-            ImmutableSet.of(PathFragment.create("worker.sh")));
-    workerExecRoot.createFileSystem();
+    SandboxInputs inputs =
+        new SandboxInputs(
+            ImmutableMap.of(PathFragment.create("worker.sh"), workerSh),
+            ImmutableSet.of(),
+            ImmutableMap.of());
+    SandboxOutputs outputs =
+        SandboxOutputs.create(
+            ImmutableSet.of(PathFragment.create("very/output.txt")), ImmutableSet.of());
+    ImmutableSet<PathFragment> workerFiles = ImmutableSet.of(PathFragment.create("worker.sh"));
+    WorkerExecRoot workerExecRoot = new WorkerExecRoot(execRoot);
+    workerExecRoot.createFileSystem(workerFiles, inputs, outputs);
 
     // Pretend to do some work inside the execRoot.
     execRoot.getRelative("tempdir").createDirectory();
@@ -82,7 +82,7 @@ public void cleanFileSystem() throws Exception {
     FileSystemUtils.writeContentAsLatin1(workerSh, "#!/bin/sh");
 
     // Reuse the same execRoot.
-    workerExecRoot.createFileSystem();
+    workerExecRoot.createFileSystem(workerFiles, inputs, outputs);
 
     assertThat(execRoot.getRelative("worker.sh").exists()).isTrue();
     assertThat(
@@ -102,21 +102,20 @@ public void createsAndCleansInputSymlinks() throws Exception {
     FileSystemUtils.ensureSymbolicLink(execRoot.getRelative("dir/input_symlink_2"), "unchanged");
     FileSystemUtils.ensureSymbolicLink(execRoot.getRelative("dir/input_symlink_3"), "whatever");
 
-    WorkerExecRoot workerExecRoot =
-        new WorkerExecRoot(
-            execRoot,
-            new SandboxInputs(
-                ImmutableMap.of(),
-                ImmutableSet.of(),
-                ImmutableMap.of(
-                    PathFragment.create("dir/input_symlink_1"), PathFragment.create("new_content"),
-                    PathFragment.create("dir/input_symlink_2"), PathFragment.create("unchanged"))),
-            SandboxOutputs.create(ImmutableSet.of(), ImmutableSet.of()),
-            ImmutableSet.of());
+    SandboxInputs inputs =
+        new SandboxInputs(
+            ImmutableMap.of(),
+            ImmutableSet.of(),
+            ImmutableMap.of(
+                PathFragment.create("dir/input_symlink_1"), PathFragment.create("new_content"),
+                PathFragment.create("dir/input_symlink_2"), PathFragment.create("unchanged")));
+    SandboxOutputs outputs = SandboxOutputs.create(ImmutableSet.of(), ImmutableSet.of());
+    ImmutableSet<PathFragment> workerFiles = ImmutableSet.of();
+    WorkerExecRoot workerExecRoot = new WorkerExecRoot(execRoot);
 
     // This should update the `input_symlink_{1,2,3}` according to `SandboxInputs`, i.e., update the
     // first/second (alternatively leave the second unchanged) and delete the third.
-    workerExecRoot.createFileSystem();
+    workerExecRoot.createFileSystem(workerFiles, inputs, outputs);
 
     assertThat(execRoot.getRelative("dir/input_symlink_1").readSymbolicLink())
         .isEqualTo(PathFragment.create("new_content"));
@@ -127,23 +126,23 @@ public void createsAndCleansInputSymlinks() throws Exception {
 
   @Test
   public void createsOutputDirs() throws Exception {
-    WorkerExecRoot workerExecRoot =
-        new WorkerExecRoot(
-            execRoot,
-            new SandboxInputs(ImmutableMap.of(), ImmutableSet.of(), ImmutableMap.of()),
-            SandboxOutputs.create(
-                ImmutableSet.of(
-                    PathFragment.create("dir/foo/bar_kt.jar"),
-                    PathFragment.create("dir/foo/bar_kt.jdeps"),
-                    PathFragment.create("dir/foo/bar_kt-sources.jar")),
-                ImmutableSet.of(
-                    PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_sourcegenfiles"),
-                    PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_classes"),
-                    PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_temp"),
-                    PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_generated_classes"))),
-            ImmutableSet.of());
-
-    workerExecRoot.createFileSystem();
+    SandboxInputs inputs =
+        new SandboxInputs(ImmutableMap.of(), ImmutableSet.of(), ImmutableMap.of());
+    SandboxOutputs outputs =
+        SandboxOutputs.create(
+            ImmutableSet.of(
+                PathFragment.create("dir/foo/bar_kt.jar"),
+                PathFragment.create("dir/foo/bar_kt.jdeps"),
+                PathFragment.create("dir/foo/bar_kt-sources.jar")),
+            ImmutableSet.of(
+                PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_sourcegenfiles"),
+                PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_classes"),
+                PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_temp"),
+                PathFragment.create("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_generated_classes")));
+    ImmutableSet<PathFragment> workerFiles = ImmutableSet.of();
+    WorkerExecRoot workerExecRoot = new WorkerExecRoot(execRoot);
+
+    workerExecRoot.createFileSystem(workerFiles, inputs, outputs);
 
     assertThat(execRoot.getRelative("dir/foo/_kotlinc/bar_kt_jvm/bar_kt_sourcegenfiles").exists())
         .isTrue();
@@ -170,16 +169,15 @@ public void workspaceFilesAreNotDeleted() throws Exception {
     FileSystemUtils.ensureSymbolicLink(execRoot.getRelative("needed_file"), neededWorkspaceFile);
     FileSystemUtils.ensureSymbolicLink(execRoot.getRelative("other_file"), otherWorkspaceFile);
 
-    WorkerExecRoot workerExecRoot =
-        new WorkerExecRoot(
-            execRoot,
-            new SandboxInputs(
-                ImmutableMap.of(PathFragment.create("needed_file"), neededWorkspaceFile),
-                ImmutableSet.of(),
-                ImmutableMap.of()),
-            SandboxOutputs.create(ImmutableSet.of(), ImmutableSet.of()),
-            ImmutableSet.of());
-    workerExecRoot.createFileSystem();
+    SandboxInputs inputs =
+        new SandboxInputs(
+            ImmutableMap.of(PathFragment.create("needed_file"), neededWorkspaceFile),
+            ImmutableSet.of(),
+            ImmutableMap.of());
+    SandboxOutputs outputs = SandboxOutputs.create(ImmutableSet.of(), ImmutableSet.of());
+    ImmutableSet<PathFragment> workerFiles = ImmutableSet.of();
+    WorkerExecRoot workerExecRoot = new WorkerExecRoot(execRoot);
+    workerExecRoot.createFileSystem(workerFiles, inputs, outputs);
 
     assertThat(execRoot.getRelative("needed_file").readSymbolicLink())
         .isEqualTo(neededWorkspaceFile.asFragment());
@@ -199,17 +197,15 @@ public void recreatesEmptyFiles() throws Exception {
 
     HashMap<PathFragment, Path> inputs = new HashMap<>();
     inputs.put(PathFragment.create("some_file"), null);
-    WorkerExecRoot workerExecRoot =
-        new WorkerExecRoot(
-            execRoot,
-            new SandboxInputs(inputs, ImmutableSet.of(), ImmutableMap.of()),
-            SandboxOutputs.create(ImmutableSet.of(), ImmutableSet.of()),
-            ImmutableSet.of());
+    SandboxInputs sandboxInputs = new SandboxInputs(inputs, ImmutableSet.of(), ImmutableMap.of());
+    SandboxOutputs outputs = SandboxOutputs.create(ImmutableSet.of(), ImmutableSet.of());
+    ImmutableSet<PathFragment> workerFiles = ImmutableSet.of();
+    WorkerExecRoot workerExecRoot = new WorkerExecRoot(execRoot);
 
     // This is interesting, because the filepath is a key in `SandboxInputs`, but its value is
     // `null`, which means "create an empty file". So after `createFileSystem` the file should be
     // empty.
-    workerExecRoot.createFileSystem();
+    workerExecRoot.createFileSystem(workerFiles, sandboxInputs, outputs);
 
     assertThat(
             FileSystemUtils.readContent(