Skip to content

Commit

Permalink
Make WorkerExecRoot not be re-created on each createFileSystem() call…
Browse files Browse the repository at this point in the history
…. Preparation for holding a map of existing links, but also just nicer.

RELNOTES: None
PiperOrigin-RevId: 356465646
  • Loading branch information
larsrc-google authored and copybara-github committed Feb 9, 2021
1 parent 4b68532 commit 31db460
Show file tree
Hide file tree
Showing 7 changed files with 88 additions and 91 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand All @@ -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();
Expand All @@ -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());
Expand All @@ -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)) {
Expand All @@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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(
Expand All @@ -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"));
Expand All @@ -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();
Expand All @@ -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());
Expand All @@ -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(
Expand Down

0 comments on commit 31db460

Please sign in to comment.