// Copyright 2015 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package com.google.devtools.build.lib.worker;
import static java.nio.charset.StandardCharsets.UTF_8;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.devtools.build.lib.worker.ExampleWorkerOptions.ExampleWorkOptions;
import com.google.devtools.build.lib.worker.WorkerProtocol.Input;
import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
import com.google.devtools.common.options.OptionsParser;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.Random;
import java.util.UUID;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
/**
* An example implementation of a worker process that is used for integration tests.
*/
public class ExampleWorker {
static final Pattern FLAG_FILE_PATTERN = Pattern.compile("(?:@|--?flagfile=)(.+)");
// A UUID that uniquely identifies this running worker process.
static final UUID workerUuid = UUID.randomUUID();
// A counter that increases with each work unit processed.
static int workUnitCounter = 1;
// If true, returns corrupt responses instead of correct protobufs.
static boolean poisoned = false;
// Keep state across multiple builds.
static final LinkedHashMap<String, String> inputs = new LinkedHashMap<>();
public static void main(String[] args) throws Exception {
if (ImmutableSet.copyOf(args).contains("--persistent_worker")) {
OptionsParser parser = OptionsParser.newOptionsParser(ExampleWorkerOptions.class);
parser.setAllowResidue(false);
parser.parse(args);
ExampleWorkerOptions workerOptions = parser.getOptions(ExampleWorkerOptions.class);
Preconditions.checkState(workerOptions.persistentWorker);
runPersistentWorker(workerOptions);
} else {
// This is a single invocation of the example that exits after it processed the request.
processRequest(ImmutableList.copyOf(args));
}
}
private static void runPersistentWorker(ExampleWorkerOptions workerOptions) throws IOException {
PrintStream originalStdOut = System.out;
PrintStream originalStdErr = System.err;
while (true) {
try {
WorkRequest request = WorkRequest.parseDelimitedFrom(System.in);
if (request == null) {
break;
}
inputs.clear();
for (Input input : request.getInputsList()) {
inputs.put(input.getPath(), input.getDigest().toStringUtf8());
}
ByteArrayOutputStream baos = new ByteArrayOutputStream();
int exitCode = 0;
try (PrintStream ps = new PrintStream(baos)) {
System.setOut(ps);
System.setErr(ps);
if (poisoned) {
if (workerOptions.hardPoison) {
throw new IllegalStateException("I'm a very poisoned worker and will just crash.");
}
System.out.println("I'm a poisoned worker and this is not a protobuf.");
System.out.println("Here's a fake stack trace for you:");
System.out.println(" at com.example.Something(Something.java:83)");
System.out.println(" at java.lang.Thread.run(Thread.java:745)");
System.out.print("And now, 8k of random bytes: ");
byte[] b = new byte[8192];
new Random().nextBytes(b);
System.out.write(b);
} else {
try {
processRequest(request.getArgumentsList());
} catch (Exception e) {
e.printStackTrace();
exitCode = 1;
}
}
} finally {
System.setOut(originalStdOut);
System.setErr(originalStdErr);
}
if (poisoned) {
baos.writeTo(System.out);
} else {
WorkResponse.newBuilder()
.setOutput(baos.toString())
.setExitCode(exitCode)
.build()
.writeDelimitedTo(System.out);
}
System.out.flush();
if (workerOptions.exitAfter > 0 && workUnitCounter > workerOptions.exitAfter) {
return;
}
if (workerOptions.poisonAfter > 0 && workUnitCounter > workerOptions.poisonAfter) {
poisoned = true;
}
} finally {
// Be a good worker process and consume less memory when idle.
System.gc();
}
}
}
private static void processRequest(List<String> args) throws Exception {
ImmutableList.Builder<String> expandedArgs = ImmutableList.builder();
for (String arg : args) {
Matcher flagFileMatcher = FLAG_FILE_PATTERN.matcher(arg);
if (flagFileMatcher.matches()) {
expandedArgs.addAll(Files.readAllLines(Paths.get(flagFileMatcher.group(1)), UTF_8));
} else {
expandedArgs.add(arg);
}
}
OptionsParser parser = OptionsParser.newOptionsParser(ExampleWorkOptions.class);
parser.setAllowResidue(true);
parser.parse(expandedArgs.build());
ExampleWorkOptions options = parser.getOptions(ExampleWorkOptions.class);
List<String> outputs = new ArrayList<>();
if (options.writeUUID) {
outputs.add("UUID " + workerUuid.toString());
}
if (options.writeCounter) {
outputs.add("COUNTER " + workUnitCounter++);
}
String residueStr = Joiner.on(' ').join(parser.getResidue());
if (options.uppercase) {
residueStr = residueStr.toUpperCase();
}
outputs.add(residueStr);
if (options.printInputs) {
for (Entry<String, String> input : inputs.entrySet()) {
outputs.add("INPUT " + input.getKey() + " " + input.getValue());
}
}
if (options.printEnv) {
for (Entry<String, String> entry : System.getenv().entrySet()) {
outputs.add(entry.getKey() + "=" + entry.getValue());
}
}
String outputStr = Joiner.on('\n').join(outputs);
if (options.outputFile.isEmpty()) {
System.out.println(outputStr);
} else {
try (PrintStream outputFile = new PrintStream(options.outputFile)) {
outputFile.println(outputStr);
}
}
}
}