Java tutorial
// 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; /** * An example implementation of a worker process that is used for integration tests. */ public class ExampleWorker { // 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) { 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 { if (args.size() == 1 && args.get(0).startsWith("@")) { args = Files.readAllLines(Paths.get(args.get(0).substring(1)), UTF_8); } OptionsParser parser = OptionsParser.newOptionsParser(ExampleWorkOptions.class); parser.setAllowResidue(true); parser.parse(args); 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); } } } }