Skip to content

Commit

Permalink
Migrate ExampleWorker to use WorkRequestHandler.
Browse files Browse the repository at this point in the history
This means that we're using a real worker implementation in our example and integration tests. It also includes a JSON WorkRequestHandler, but it isn't used except in the tests.

RELNOTES: None.
PiperOrigin-RevId: 345765354
  • Loading branch information
susinmotion authored and larsrc-google committed Jul 30, 2021
1 parent 35c98d0 commit c4e22b9
Show file tree
Hide file tree
Showing 6 changed files with 164 additions and 178 deletions.
4 changes: 4 additions & 0 deletions src/main/java/com/google/devtools/build/lib/worker/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -51,11 +51,15 @@ java_library(
java_library(
name = "work_request_handlers",
srcs = [
"JsonWorkerMessageProcessor.java",
"ProtoWorkerMessageProcessor.java",
"WorkRequestHandler.java",
],
deps = [
"//src/main/protobuf:worker_protocol_java_proto",
"//third_party:gson",
"//third_party:guava",
"//third_party/protobuf:protobuf_java",
"//third_party/protobuf:protobuf_java_util",
],
)
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,7 @@
// limitations under the License.
package com.google.devtools.build.lib.worker;

import static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.collect.ImmutableList;
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;
Expand All @@ -23,43 +22,42 @@
import com.google.protobuf.ByteString;
import com.google.protobuf.util.JsonFormat;
import com.google.protobuf.util.JsonFormat.Printer;
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.List;

/** Sample implementation of the Worker Protocol using JSON to communicate with Bazel. */
public class JsonExampleWorkerProtocolImpl implements ExampleWorkerProtocol {
private final Printer jsonPrinter =
JsonFormat.printer().omittingInsignificantWhitespace().includingDefaultValueFields();
/** Implementation of the Worker Protocol using JSON to communicate with Bazel. */
public final class JsonWorkerMessageProcessor implements WorkRequestHandler.WorkerMessageProcessor {
/** Reader for reading the WorkResponse. */
private final JsonReader reader;
/** Printer for printing the WorkRequest. */
private final Printer jsonPrinter;
/** Writer for writing the WorkRequest to the worker. */
private final BufferedWriter jsonWriter;

public JsonExampleWorkerProtocolImpl(InputStream stdin, OutputStream stdout) {
reader = new JsonReader(new BufferedReader(new InputStreamReader(stdin, UTF_8)));
/** Constructs a {@code WorkRequestHandler} that reads and writes JSON. */
public JsonWorkerMessageProcessor(JsonReader reader, BufferedWriter jsonWriter) {
this.reader = reader;
reader.setLenient(true);
jsonWriter = new BufferedWriter(new OutputStreamWriter(stdout, UTF_8));
this.jsonWriter = jsonWriter;
jsonPrinter =
JsonFormat.printer().omittingInsignificantWhitespace().includingDefaultValueFields();
}

private static ArrayList<String> readArguments(JsonReader reader) throws IOException {
private static ImmutableList<String> readArguments(JsonReader reader) throws IOException {
reader.beginArray();
ArrayList<String> arguments = new ArrayList<>();
ImmutableList.Builder<String> argumentsBuilder = ImmutableList.builder();
while (reader.hasNext()) {
arguments.add(reader.nextString());
argumentsBuilder.add(reader.nextString());
}
reader.endArray();
return arguments;
return argumentsBuilder.build();
}

private static ArrayList<Input> readInputs(JsonReader reader) throws IOException {
private static ImmutableList<Input> readInputs(JsonReader reader) throws IOException {
reader.beginArray();
ArrayList<Input> inputs = new ArrayList<>();
ImmutableList.Builder<Input> inputsBuilder = ImmutableList.builder();
while (reader.hasNext()) {
String digest = null;
String path = null;
Expand All @@ -81,19 +79,25 @@ private static ArrayList<Input> readInputs(JsonReader reader) throws IOException
path = reader.nextString();
break;
default:
throw new IOException(name + " is an incorrect field in input");
continue;
}
}
reader.endObject();
inputs.add(
Input.newBuilder().setDigest(ByteString.copyFromUtf8(digest)).setPath(path).build());
Input.Builder inputBuilder = Input.newBuilder();
if (digest != null) {
inputBuilder.setDigest(ByteString.copyFromUtf8(digest));
}
if (path != null) {
inputBuilder.setPath(path);
}
inputsBuilder.add(inputBuilder.build());
}
reader.endArray();
return inputs;
return inputsBuilder.build();
}

@Override
public WorkRequest readRequest() throws IOException {
public WorkRequest readWorkRequest() throws IOException {
List<String> arguments = null;
List<Input> inputs = null;
Integer requestId = null;
Expand All @@ -104,24 +108,24 @@ public WorkRequest readRequest() throws IOException {
switch (name) {
case "arguments":
if (arguments != null) {
throw new IOException("Work request cannot have more than one list of arguments");
throw new IOException("WorkRequest cannot have more than one 'arguments' field");
}
arguments = readArguments(reader);
break;
case "inputs":
if (inputs != null) {
throw new IOException("Work request cannot have more than one list of inputs");
throw new IOException("WorkRequest cannot have more than one 'inputs' field");
}
inputs = readInputs(reader);
break;
case "requestId":
if (requestId != null) {
throw new IOException("Work request cannot have more than one requestId");
throw new IOException("WorkRequest cannot have more than one requestId");
}
requestId = reader.nextInt();
break;
default:
throw new IOException(name + " is an incorrect field in work request");
break;
}
}
reader.endObject();
Expand All @@ -143,17 +147,13 @@ public WorkRequest readRequest() throws IOException {
}

@Override
public void writeResponse(WorkResponse response) throws IOException {
public void writeWorkResponse(WorkResponse response) throws IOException {
jsonPrinter.appendTo(response, jsonWriter);
jsonWriter.flush();
}

@Override
public void close() {
try {
jsonWriter.close();
} catch (IOException e) {
System.err.printf("Could not close json writer. %s", e);
}
public void close() throws IOException {
jsonWriter.close();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ public interface WorkerMessageProcessor {
/** This worker's stderr. */
private final PrintStream stderr;

private final WorkerMessageProcessor messageProcessor;
final WorkerMessageProcessor messageProcessor;

private final CpuTimeBasedGcScheduler gcScheduler;

Expand Down
1 change: 1 addition & 0 deletions src/test/java/com/google/devtools/build/lib/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -605,6 +605,7 @@ java_library(
],
deps = [
"//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
"//src/main/java/com/google/devtools/build/lib/worker:work_request_handlers",
"//src/main/java/com/google/devtools/common/options",
"//src/main/protobuf:worker_protocol_java_proto",
"//third_party:gson",
Expand Down
Loading

0 comments on commit c4e22b9

Please sign in to comment.