workers: Use the new worker testbed to check that workers get reused correctly, restart after a clean exit and that workers returning junk are being dealt with.

--
MOS_MIGRATED_REVID=103542544

3 files changed, 152 insertions, 22 deletions

diff --git a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java
index b9146b2fba..3364c8847b 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java
@@ -31,12 +31,22 @@ import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
+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;
+
   public static void main(String[] args) throws Exception {
     if (ImmutableSet.copyOf(args).contains("--persistent_worker")) {
       OptionsParser parser = OptionsParser.newOptionsParser(ExampleWorkerOptions.class);
@@ -81,12 +91,24 @@ public class ExampleWorker {
           System.setErr(originalStdErr);
         }
 
-        WorkResponse.newBuilder()
-            .setOutput(baos.toString())
-            .setExitCode(exitCode)
-            .build()
-            .writeDelimitedTo(System.out);
+        if (poisoned) {
+          System.out.println("I'm a poisoned worker and this is not a protobuf.");
+        } 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();
@@ -102,25 +124,29 @@ public class ExampleWorker {
     OptionsParser parser = OptionsParser.newOptionsParser(ExampleWorkOptions.class);
     parser.setAllowResidue(true);
     parser.parse(args);
-    ExampleWorkOptions workOptions = parser.getOptions(ExampleWorkOptions.class);
-
-    List<String> residue = parser.getResidue();
-    List<String> outputs = new ArrayList(residue.size());
-    for (String arg : residue) {
-      String output = arg;
-      if (workOptions.uppercase) {
-        output = output.toUpperCase();
-      }
-      outputs.add(output);
+    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);
 
-    String outputStr = Joiner.on(' ').join(outputs);
-    if (workOptions.outputFile.isEmpty()) {
-      System.err.println("ExampleWorker: Writing to stdout!");
+    String outputStr = Joiner.on('\n').join(outputs);
+    if (options.outputFile.isEmpty()) {
       System.out.println(outputStr);
     } else {
-      System.err.println("ExampleWorker: Writing to file " + workOptions.outputFile);
-      try (PrintStream outputFile = new PrintStream(workOptions.outputFile)) {
+      try (PrintStream outputFile = new PrintStream(options.outputFile)) {
         outputFile.println(outputStr);
       }
     }
diff --git a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java
index 3b5878c8f4..48d9e8f159 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java
@@ -34,8 +34,34 @@ public class ExampleWorkerOptions extends OptionsBase {
 
     @Option(name = "uppercase", defaultValue = "false", help = "Uppercase the input.")
     public boolean uppercase;
+
+    @Option(name = "write_uuid", defaultValue = "false", help = "Writes a UUID into the output.")
+    public boolean writeUUID;
+
+    @Option(
+      name = "write_counter",
+      defaultValue = "false",
+      help = "Writes a counter that increases with each work unit processed into the output."
+    )
+    public boolean writeCounter;
   }
 
   @Option(name = "persistent_worker", defaultValue = "false")
   public boolean persistentWorker;
+
+  @Option(
+    name = "exit_after",
+    defaultValue = "0",
+    help = "The worker exits after processing this many work units (default: disabled)."
+  )
+  public int exitAfter;
+
+  @Option(
+    name = "poison_after",
+    defaultValue = "0",
+    help =
+        "Poisons the worker after processing this many work units, so that it returns a "
+            + "corrupt response instead of a response protobuf from then on (default: disabled)."
+  )
+  public int poisonAfter;
 }
diff --git a/src/test/shell/bazel/bazel_worker_test.sh b/src/test/shell/bazel/bazel_worker_test.sh
index 9367bb2f89..7bd410a623 100755
--- a/src/test/shell/bazel/bazel_worker_test.sh
+++ b/src/test/shell/bazel/bazel_worker_test.sh
@@ -143,7 +143,7 @@ function test_workers_quit_after_build() {
 }
 
 function prepare_example_worker() {
-  cp -v ${example_worker} worker_lib.jar
+  cp ${example_worker} worker_lib.jar
 
   cat >work.bzl <<'EOF'
 def _impl(ctx):
@@ -152,7 +152,10 @@ def _impl(ctx):
 
   # Generate the "@"-file containing the command-line args for the unit of work.
   argfile = ctx.new_file(ctx.configuration.bin_dir, "worker_input")
-  ctx.file_action(output=argfile, content="\n".join(["--output_file=" + output.path] + ctx.attr.args))
+  argfile_contents = "\n".join(["--output_file=" + output.path] + ctx.attr.args)
+  ctx.file_action(output=argfile, content=argfile_contents)
+
+  print("Using argfile_contents: " + argfile_contents)
 
   ctx.action(
       inputs=[argfile],
@@ -223,4 +226,79 @@ EOF
   assert_workers_not_running
 }
 
+function test_worker_restarts() {
+  prepare_example_worker
+
+  cat >>BUILD <<'EOF'
+[work(
+  name = "hello_world_%s" % idx,
+  worker = ":worker",
+  worker_args = ["--exit_after=2"],
+  args = ["--write_uuid", "--write_counter"],
+) for idx in range(10)]
+EOF
+
+  bazel --batch clean
+  assert_workers_not_running
+
+  bazel build --strategy=Work=worker --worker_max_instances=1 :hello_world_1 \
+    || fail "build failed"
+  worker_uuid_1=$(cat bazel-bin/hello_world_1.out | grep UUID | cut -d' ' -f2)
+  work_count=$(cat bazel-bin/hello_world_1.out | grep COUNTER | cut -d' ' -f2)
+  assert_equals "1" $work_count
+  assert_workers_running
+
+  bazel build --strategy=Work=worker --worker_max_instances=1 :hello_world_2 \
+    || fail "build failed"
+  worker_uuid_2=$(cat bazel-bin/hello_world_2.out | grep UUID | cut -d' ' -f2)
+  work_count=$(cat bazel-bin/hello_world_2.out | grep COUNTER | cut -d' ' -f2)
+  assert_equals "2" $work_count
+  assert_workers_not_running
+
+  # Check that the same worker was used twice.
+  assert_equals "$worker_uuid_1" "$worker_uuid_2"
+
+  bazel build --strategy=Work=worker --worker_max_instances=1 :hello_world_3 \
+    || fail "build failed"
+  worker_uuid_3=$(cat bazel-bin/hello_world_3.out | grep UUID | cut -d' ' -f2)
+  work_count=$(cat bazel-bin/hello_world_3.out | grep COUNTER | cut -d' ' -f2)
+  assert_equals "1" $work_count
+  assert_workers_running
+
+  # Check that we used a new worker.
+  assert_not_equals "$worker_uuid_2" "$worker_uuid_3"
+}
+
+# When a worker does not conform to the protocol and returns a response that is not a parseable
+# protobuf, it must be killed, the output thrown away, a new worker restarted and Bazel has to retry
+# the action without struggling.
+function test_bazel_recovers_from_worker_returning_junk() {
+  prepare_example_worker
+
+  cat >>BUILD <<'EOF'
+[work(
+  name = "hello_world_%s" % idx,
+  worker = ":worker",
+  worker_args = ["--poison_after=1"],
+  args = ["--write_uuid", "--write_counter"],
+) for idx in range(10)]
+EOF
+
+  bazel --batch clean
+  assert_workers_not_running
+
+  bazel build --strategy=Work=worker --worker_max_instances=1 :hello_world_1 \
+    || fail "build failed"
+  worker_uuid_1=$(cat bazel-bin/hello_world_1.out | grep UUID | cut -d' ' -f2)
+  assert_workers_running
+
+  bazel build --strategy=Work=worker --worker_max_instances=1 :hello_world_2 \
+    || fail "build failed"
+  worker_uuid_2=$(cat bazel-bin/hello_world_2.out | grep UUID | cut -d' ' -f2)
+  assert_workers_running
+
+  # Check that the worker failed & was restarted.
+  assert_not_equals "$worker_uuid_1" "$worker_uuid_2"
+}
+
 run_suite "Worker integration tests"