clean up workers

src/java/io/bazel/rulesscala/worker/BUILD

@@ -12,8 +12,12 @@ java_test(
     srcs = [
         "WorkerTest.java",
     ],
+    jvm_flags = [
+        "-Dcom.google.testing.junit.runner.shouldInstallTestSecurityManager=false",
+    ],
     test_class = "io.bazel.rulesscala.worker.WorkerTest",
     deps = [
         ":worker",
+        "//third_party/bazel/src/main/protobuf:worker_protocol_java_proto",
     ],
 )

src/java/io/bazel/rulesscala/worker/Worker.java

@@ -62,6 +62,10 @@ public void checkPermission(Permission permission) {
 		WorkerProtocol.WorkRequest request =
 		    WorkerProtocol.WorkRequest.parseDelimitedFrom(stdin);
 
+		if (request == null) {
+		    break;
+		}
+
 		int code = 0;
 
 		try {
@@ -136,7 +140,7 @@ public void reset() {
 	}
     }
 
-    private static class ExitTrapped extends RuntimeException {
+    static class ExitTrapped extends RuntimeException {
 	final int code;
 	ExitTrapped(int code) {
 	    super();

src/java/io/bazel/rulesscala/worker/WorkerTest.java

@@ -1,23 +1,105 @@
 package io.bazel.rulesscala.worker;
 
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 import java.io.ByteArrayOutputStream;
+import java.io.InputStream;
+import java.io.PipedInputStream;
+import java.io.PipedOutputStream;
 import java.io.PrintStream;
+import java.lang.SecurityManager;
+
+import com.google.devtools.build.lib.worker.WorkerProtocol;
 
 @RunWith(JUnit4.class)
 public class WorkerTest {
 
+    @Test
+    public void testEphemeralWorkerSystemExit() throws Exception {
+
+	// An ephemeral worker behaves like a regular main method,
+	// so we expect the worker to system exit normally
+
+	Worker.Interface worker = new Worker.Interface() {
+            @Override
+	    public void work(String[] args) {
+		System.exit(99);
+	    }
+	};
+
+	// we expect ephemeral workers to just exit normally
+	int code = assertThrows(Worker.ExitTrapped.class, () ->
+	    Worker.workerMain(new String[]{}, worker)).code;
+
+	assert(code == 99);
+    }
+
+    @Test
+    public void testPersistentWorkerSystemExit() throws Exception {
+
+	// We're going to spin up a persistent worker and run a single
+	// work request. We expect System exists to impact the worker
+	// request lifecycle without exiting the overall worker
+	// process.
+
+	Worker.Interface worker = new Worker.Interface() {
+            @Override
+	    public void work(String[] args) {
+		// we should see this print statement
+		System.out.println("before exit");
+		System.exit(99);
+		// we should not see this print statement
+		System.out.println("after exit");
+	    }
+	};
+
+	try (
+	    PipedInputStream in = new PipedInputStream();
+	    PipedOutputStream outToIn = new PipedOutputStream(in);
+
+	    ByteArrayOutputStream out = new ByteArrayOutputStream();
+	) {
+
+	    InputStream stdin = System.in;
+	    PrintStream stdout = System.out;
+	    PrintStream stderr = System.err;
+
+	    System.setIn(in);
+	    System.setOut(new PrintStream(out));
+
+	    WorkerProtocol.WorkRequest.newBuilder()
+	        .build()
+		.writeDelimitedTo(outToIn);
+
+	    // otherwise the worker will poll indefinitely
+	    outToIn.close();
+
+	    Worker.workerMain(new String[]{"--persistent_worker"}, worker);
+
+	    System.setIn(stdin);
+	    System.setOut(stdout);
+	    System.setErr(stderr);
+
+	    String outString = out.toString("UTF-8");
+	    // check to make sure the before statement printed
+	    assert(outString.contains("before"));
+	    // check to make sure the after statement did not print
+	    assert(!outString.contains("after"));
+	}
+    }
+
     private static void fill(ByteArrayOutputStream baos, int amount) {
 	for (int i = 0; i < amount; i++) {
 	    baos.write(0);
 	}
     }
 
     @Test
-    public void testWriteReadAndReset() throws Exception {
+    public void testBufferWriteReadAndReset() throws Exception {
 	Worker.SmartByteArrayOutputStream baos = new Worker.SmartByteArrayOutputStream();
 	PrintStream out = new PrintStream(baos);
 
@@ -33,4 +115,46 @@ public void testWriteReadAndReset() throws Exception {
 	assert(baos.toString("UTF-8").equals("goodbye, world"));
 	assert(!baos.isOversized());
     }
+
+    @AfterClass
+    public static void teardown() {
+	// Persistent workers install a security manager. We need to
+	// reset it here so that our own process can exit!
+	System.setSecurityManager(null);
+    }
+
+    // Copied/modified from Bazel's MoreAsserts
+    //
+    // Note: this goes away soon-ish, as JUnit 4.13 was recently
+    // released and includes assertThrows
+    public static <T extends Throwable> T assertThrows(
+	Class<T> expectedThrowable,
+	ThrowingRunnable runnable)
+    {
+	try {
+	    runnable.run();
+	} catch (Throwable actualThrown) {
+	    if (expectedThrowable.isInstance(actualThrown)) {
+		@SuppressWarnings("unchecked")
+		    T retVal = (T) actualThrown;
+		return retVal;
+	    } else {
+		throw new AssertionError(
+		    String.format(
+		        "expected %s to be thrown, but %s was thrown",
+			expectedThrowable.getSimpleName(),
+		        actualThrown.getClass().getSimpleName()),
+		    actualThrown);
+	    }
+	}
+	String mismatchMessage = String.format(
+	   "expected %s to be thrown, but nothing was thrown",
+	   expectedThrowable.getSimpleName());
+	throw new AssertionError(mismatchMessage);
+    }
+
+    // see note on assertThrows
+    public interface ThrowingRunnable {
+	void run() throws Throwable;
+    }
 }