Skip to content

Commit

Permalink
Add test utils and tests for ProtoCoordinator
Browse files Browse the repository at this point in the history
Co-authored-by: Shane Delmore <[email protected]>
  • Loading branch information
wiwa and ShaneDelmore committed Nov 14, 2023
1 parent de43727 commit 3237c63
Show file tree
Hide file tree
Showing 16 changed files with 488 additions and 102 deletions.
8 changes: 2 additions & 6 deletions src/main/java/build/buildfarm/worker/Executor.java
Original file line number Diff line number Diff line change
Expand Up @@ -452,12 +452,7 @@ private Code executeCommand(
Tree execTree = workerContext.getQueuedOperation(operationContext.queueEntry).getTree();

WorkFilesContext filesContext =
new WorkFilesContext(
execDir,
execTree,
ImmutableList.copyOf(operationContext.command.getOutputPathsList()),
ImmutableList.copyOf(operationContext.command.getOutputFilesList()),
ImmutableList.copyOf(operationContext.command.getOutputDirectoriesList()));
WorkFilesContext.fromContext(execDir, execTree, operationContext.command);

return PersistentExecutor.runOnPersistentWorker(
limits.persistentWorkerCommand,
Expand All @@ -467,6 +462,7 @@ private Code executeCommand(
ImmutableMap.copyOf(environment),
limits,
timeout,
PersistentExecutor.defaultWorkRootsDir,
resultBuilder);
}

Expand Down
2 changes: 1 addition & 1 deletion src/main/java/build/buildfarm/worker/OperationContext.java
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import com.google.longrunning.Operation;
import java.nio.file.Path;

final class OperationContext {
public final class OperationContext {
final ExecuteResponse.Builder executeResponse;
final Operation operation;
final Poller poller;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,23 +3,41 @@
import static java.nio.file.StandardCopyOption.COPY_ATTRIBUTES;
import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;

import com.google.common.collect.ImmutableSet;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.attribute.PosixFilePermission;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;
import java.util.logging.Logger;

// Utility for concurrent move/copy/link of files
/**
* Utility for concurrent move/copy of files Can be extended in the future to (sym)linking if we
* need performance
*/
public final class FileAccessUtils {
// singleton class with only static methods
private FileAccessUtils() {}

private static final Logger logger = Logger.getLogger(FileAccessUtils.class.getName());

public static Path addPosixOwnerWrite(Path absPath) throws IOException {
Set<PosixFilePermission> perms = Files.getPosixFilePermissions(absPath);

ImmutableSet<PosixFilePermission> permsWithWrite =
ImmutableSet.<PosixFilePermission>builder()
.addAll(perms)
.add(PosixFilePermission.OWNER_WRITE)
.build();

return Files.setAttribute(absPath, "posix:permissions", permsWithWrite);
}

private static final ConcurrentHashMap<Path, PathLock> fileLocks = new ConcurrentHashMap<>();

// Used here for locking "files"
// Used here as a simple lock for locking "files" (paths)
private static class PathLock {
// Not used elsewhere
private PathLock() {}
Expand All @@ -46,13 +64,10 @@ public static void copyFile(Path from, Path to) throws IOException {
() -> {
try {
Files.copy(from, absTo, REPLACE_EXISTING, COPY_ATTRIBUTES);
boolean writeable = absTo.toFile().setWritable(true);
if (!writeable) {
return new IOException("copyFile() could not set writeable: " + absTo);
}
addPosixOwnerWrite(absTo);
return null;
} catch (IOException e) {
return e;
return new IOException("copyFile() could not set writeable: " + absTo, e);
}
});
if (ioException != null) {
Expand Down Expand Up @@ -81,44 +96,10 @@ public static void moveFile(Path from, Path to) throws IOException {
() -> {
try {
Files.move(from, absTo, REPLACE_EXISTING);
boolean writeable = absTo.toFile().setWritable(true);
if (!writeable) {
return new IOException("moveFile() could not set writeable: " + absTo);
}
return null;
} catch (IOException e) {
return e;
}
});
if (ioException != null) {
throw ioException;
}
}

/**
* Creates a symlink, creating necessary directories. Deletes pre-existing files/links which have
* the same path as the specified link, effectively overwriting any existing files/links.
*
* @param from
* @param to
* @throws IOException
*/
public static void linkFile(Path from, Path to) throws IOException {
Path absTo = to.toAbsolutePath();
logger.finer("linkFile: " + from + " to " + absTo);
if (!Files.exists(from)) {
throw new IOException("linkFile: source file doesn't exist: " + from);
}
IOException ioException =
writeFileSafe(
absTo,
() -> {
try {
Files.deleteIfExists(absTo);
Files.createSymbolicLink(absTo, from);
addPosixOwnerWrite(absTo);
return null;
} catch (IOException e) {
return e;
return new IOException("copyFile() could not set writeable: " + absTo, e);
}
});
if (ioException != null) {
Expand Down
13 changes: 11 additions & 2 deletions src/main/java/build/buildfarm/worker/persistent/Keymaker.java
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,12 @@
import persistent.bazel.client.PersistentWorker;
import persistent.bazel.client.WorkerKey;

/** Much of the logic (hashing) is from Bazel itself (private library/methods, i.e. WorkerKey). */
public class Keymaker {
// Constructs a key with its worker tool input files being relative paths
public static WorkerKey make(
Path opRoot,
Path workRootsDir,
ImmutableList<String> workerInitCmd,
ImmutableList<String> workerInitArgs,
ImmutableMap<String, String> workerEnv,
Expand All @@ -29,7 +31,13 @@ public static WorkerKey make(

Path workRoot =
calculateWorkRoot(
workerInitCmd, workerInitArgs, workerEnv, executionName, sandboxed, cancellable);
workRootsDir,
workerInitCmd,
workerInitArgs,
workerEnv,
executionName,
sandboxed,
cancellable);
Path toolsRoot = workRoot.resolve(PersistentWorker.TOOL_INPUT_SUBDIR);

SortedMap<Path, HashCode> hashedTools = workerFilesWithHashes(workerFiles);
Expand All @@ -49,6 +57,7 @@ public static WorkerKey make(

// Hash of a subset of the WorkerKey
private static Path calculateWorkRoot(
Path workRootsDir,
ImmutableList<String> workerInitCmd,
ImmutableList<String> workerInitArgs,
ImmutableMap<String, String> workerEnv,
Expand All @@ -57,7 +66,7 @@ private static Path calculateWorkRoot(
boolean cancellable) {
int workRootId = Objects.hash(workerInitCmd, workerInitArgs, workerEnv, sandboxed, cancellable);
String workRootDirName = "work-root_" + executionName + "_" + workRootId;
return PersistentExecutor.workRootsDir.resolve(workRootDirName);
return workRootsDir.resolve(workRootDirName);
}

private static ImmutableSortedMap<Path, HashCode> workerFilesWithHashes(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,23 +24,19 @@
* Executes an Action like Executor/DockerExecutor, writing to ActionResult.
*
* <p>Currently has special code for discriminating between Javac/Scalac, and other persistent
* workers.
* workers, likely for debugging purposes, but need to revisit. (Can't remember fully since it was
* so long ago!)
*/
public class PersistentExecutor {
private static final Logger logger = Logger.getLogger(PersistentExecutor.class.getName());

// How many workers can exist at once for a given WorkerKey
// There may be multiple WorkerKeys per mnemonic,
// e.g. if builds are run with different tool fingerprints
private static final int defaultMaxWorkersPerKey = 6;

private static final ProtoCoordinator coordinator =
ProtoCoordinator.ofCommonsPool(getMaxWorkersPerKey());

// TODO load from config (i.e. {worker_root}/persistent)
static final Path workRootsDir = Paths.get("/tmp/worker/persistent/");
public static final Path defaultWorkRootsDir = Paths.get("/tmp/worker/persistent/");

static final String PERSISTENT_WORKER_FLAG = "--persistent_worker";
public static final String PERSISTENT_WORKER_FLAG = "--persistent_worker";

// TODO Revisit hardcoded actions
static final String JAVABUILDER_JAR =
Expand All @@ -49,6 +45,11 @@ public class PersistentExecutor {
private static final String SCALAC_EXEC_NAME = "Scalac";
private static final String JAVAC_EXEC_NAME = "JavaBuilder";

// How many workers can exist at once for a given WorkerKey
// There may be multiple WorkerKeys per mnemonic,
// e.g. if builds are run with different tool fingerprints
private static final int defaultMaxWorkersPerKey = 6;

private static int getMaxWorkersPerKey() {
try {
return Integer.parseInt(System.getenv("BUILDFARM_MAX_WORKERS_PER_KEY"));
Expand All @@ -73,6 +74,7 @@ public static Code runOnPersistentWorker(
ImmutableMap<String, String> envVars,
ResourceLimits limits,
Duration timeout,
Path workRootsDir,
ActionResult.Builder resultBuilder)
throws IOException {
//// Pull out persistent worker start command from the overall action request
Expand All @@ -87,6 +89,7 @@ public static Code runOnPersistentWorker(
return Code.INVALID_ARGUMENT;
}

// TODO revisit why this was necessary in the first place
ImmutableMap<String, String> env;
if (executionName.equals(JAVAC_EXEC_NAME)) {
env = ImmutableMap.of();
Expand All @@ -112,7 +115,13 @@ public static Code runOnPersistentWorker(

WorkerKey key =
Keymaker.make(
context.opRoot, workerExecCmd, workerInitArgs, env, executionName, workerFiles);
context.opRoot,
workRootsDir,
workerExecCmd,
workerInitArgs,
env,
executionName,
workerFiles);

coordinator.copyToolInputsIntoWorkerToolRoot(key, workerFiles);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ private void copyNontoolInputs(WorkerInputs workerInputs, Path workerExecRoot)

// Make outputs visible to the rest of Worker machinery
// see DockerExecutor::copyOutputsOutOfContainer
private void moveOutputsToOperationRoot(WorkFilesContext context, Path workerExecRoot)
void moveOutputsToOperationRoot(WorkFilesContext context, Path workerExecRoot)
throws IOException {
Path opRoot = context.opRoot;

Expand All @@ -216,9 +216,11 @@ private void moveOutputsToOperationRoot(WorkFilesContext context, Path workerExe
}

for (String relOutput : context.outputFiles) {
Path relPath = Paths.get(relOutput);
Path opOutputPath = opRoot.resolve(relPath);
Path execOutputPath = workerExecRoot.resolve(relPath);
System.out.println(relOutput);
Path execOutputPath = workerExecRoot.resolve(relOutput);
System.out.println(execOutputPath);
Path opOutputPath = opRoot.resolve(relOutput);
System.out.println(opOutputPath);

FileAccessUtils.moveFile(execOutputPath, opOutputPath);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package build.buildfarm.worker.persistent;

import build.bazel.remote.execution.v2.Command;
import build.buildfarm.v1test.Tree;
import build.buildfarm.worker.util.InputsIndexer;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -37,14 +38,23 @@ public WorkFilesContext(
this.outputFiles = outputFiles;
this.outputDirectories = outputDirectories;

this.inputsIndexer = new InputsIndexer(execTree);
this.inputsIndexer = new InputsIndexer(execTree, this.opRoot);
}

public static WorkFilesContext fromContext(Path opRoot, Tree inputsTree, Command opCommand) {
return new WorkFilesContext(
opRoot,
inputsTree,
ImmutableList.copyOf(opCommand.getOutputPathsList()),
ImmutableList.copyOf(opCommand.getOutputFilesList()),
ImmutableList.copyOf(opCommand.getOutputDirectoriesList()));
}

// Paths are absolute paths from the opRoot; same as the Input.getPath();
public ImmutableMap<Path, Input> getPathInputs() {
synchronized (this) {
if (pathInputs == null) {
pathInputs = inputsIndexer.getAllInputs(opRoot);
pathInputs = inputsIndexer.getAllInputs();
}
}
return pathInputs;
Expand All @@ -53,7 +63,7 @@ public ImmutableMap<Path, Input> getPathInputs() {
public ImmutableMap<Path, Input> getToolInputs() {
synchronized (this) {
if (toolInputs == null) {
toolInputs = inputsIndexer.getToolInputs(opRoot);
toolInputs = inputsIndexer.getToolInputs();
}
}
return toolInputs;
Expand Down
12 changes: 2 additions & 10 deletions src/main/java/build/buildfarm/worker/persistent/WorkerInputs.java
Original file line number Diff line number Diff line change
Expand Up @@ -57,16 +57,6 @@ public void copyInputFile(Path from, Path to) throws IOException {
FileAccessUtils.copyFile(from, to);
}

public void moveInputFile(Path from, Path to) throws IOException {
checkFileIsInput("moveInputFile()", from);
FileAccessUtils.moveFile(from, to);
}

public void linkInputFile(Path from, Path to) throws IOException {
checkFileIsInput("linkInputFile()", from);
FileAccessUtils.linkFile(from, to);
}

public void deleteInputFileIfExists(Path workerExecRoot, Path opPathInput) throws IOException {
checkFileIsInput("deleteInputFile()", opPathInput);
Path execPathInput = relativizeInput(workerExecRoot, opPathInput);
Expand Down Expand Up @@ -109,6 +99,8 @@ public static WorkerInputs from(WorkFilesContext workFilesContext, List<String>

logger.fine(inputsDebugMsg);

System.out.println(inputsDebugMsg);

return new WorkerInputs(workFilesContext.opRoot, absToolInputs, toolInputs, pathInputs);
}
}
Loading

0 comments on commit 3237c63

Please sign in to comment.