From 97c46ca1c3df82c7a6612075f4a4fcac6ab2260c Mon Sep 17 00:00:00 2001
From: Jaden Peterson <jpeterson@lucidchart.com>
Date: Fri, 7 Feb 2025 18:14:30 -0500
Subject: [PATCH] Return the proper exit code when a worker is run directly

If a worker (subclass of `WorkerMain`) is run directly and throws an
`AnnexWorkerError`, the worker should exit with that error's code, since
there's no concept of a "work request" in this context.
---
 .../common/worker/WorkerMain.scala            | 43 +++++++++++--------
 tests/worker-error/BUILD                      | 11 ++++-
 .../WorkerThrowingAnnexWorkerError.scala      | 17 ++++++++
 tests/worker-error/test                       |  7 +++
 4 files changed, 60 insertions(+), 18 deletions(-)
 create mode 100644 tests/worker-error/WorkerThrowingAnnexWorkerError.scala

diff --git a/src/main/scala/higherkindness/rules_scala/common/worker/WorkerMain.scala b/src/main/scala/higherkindness/rules_scala/common/worker/WorkerMain.scala
index 960e4a44..ecf21d60 100644
--- a/src/main/scala/higherkindness/rules_scala/common/worker/WorkerMain.scala
+++ b/src/main/scala/higherkindness/rules_scala/common/worker/WorkerMain.scala
@@ -218,28 +218,37 @@ abstract class WorkerMain[S](stdin: InputStream = System.in, stdout: PrintStream
         }
 
       case args =>
-        Using.Manager { use =>
+        val returnCode = Using.Manager { use =>
           val outStream = use(new ByteArrayOutputStream())
           val out = use(new PrintStream(outStream))
-          try {
-            work(
-              init(args = None),
-              args.toArray,
-              out,
-              workDir = Path.of(""),
-              verbosity = 0,
-            )
-          } catch {
-            // This error means the work function encountered an error that we want to not be caught
-            // inside that function. That way it stops work and exits the function. However, we
-            // also don't want to crash the whole program.
-            case e: AnnexWorkerError => e.print(out)
-          } finally {
-            out.flush()
-          }
+          val returnCode =
+            try {
+              work(
+                init(args = None),
+                args.toArray,
+                out,
+                workDir = Path.of(""),
+                verbosity = 0,
+              )
+
+              0
+            } catch {
+              // This error means the work function encountered an error that we want to not be caught
+              // inside that function. That way it stops work and exits the function. However, we
+              // also don't want to crash the whole program.
+              case e: AnnexWorkerError =>
+                e.print(out)
+                e.code
+            } finally {
+              out.flush()
+            }
 
           outStream.writeTo(System.err)
+
+          returnCode
         }.get
+
+        sys.exit(returnCode)
     }
   }
 }
diff --git a/tests/worker-error/BUILD b/tests/worker-error/BUILD
index 0c2cbf34..45c91e59 100644
--- a/tests/worker-error/BUILD
+++ b/tests/worker-error/BUILD
@@ -1,4 +1,4 @@
-load("@rules_scala_annex//rules:scala.bzl", "scala_library", "scala_test")
+load("@rules_scala_annex//rules:scala.bzl", "scala_binary", "scala_library", "scala_test")
 
 scala_library(
     name = "fatal-error-spec-workers",
@@ -48,3 +48,12 @@ scala_test(
         "@rules_scala_annex//third_party/bazel/src/main/protobuf:worker_protocol_java_proto",
     ],
 )
+
+scala_binary(
+    name = "worker-throwing-annexworkererror",
+    srcs = ["WorkerThrowingAnnexWorkerError.scala"],
+    deps = [
+        "@rules_scala_annex//src/main/scala/higherkindness/rules_scala/common/error",
+        "@rules_scala_annex//src/main/scala/higherkindness/rules_scala/common/worker",
+    ],
+)
diff --git a/tests/worker-error/WorkerThrowingAnnexWorkerError.scala b/tests/worker-error/WorkerThrowingAnnexWorkerError.scala
new file mode 100644
index 00000000..176d1568
--- /dev/null
+++ b/tests/worker-error/WorkerThrowingAnnexWorkerError.scala
@@ -0,0 +1,17 @@
+package anx.cancellation
+
+import higherkindness.rules_scala.common.error.AnnexWorkerError
+import higherkindness.rules_scala.common.worker.WorkerMain
+import java.io.PrintStream
+import java.nio.file.Path
+
+object WorkerThrowingAnnexWorkerError extends WorkerMain[Unit] {
+  override def init(arguments: Option[Array[String]]): Unit = {}
+  override def work(
+    context: Unit,
+    arguments: Array[String],
+    output: PrintStream,
+    workingDirectory: Path,
+    verbosity: Int,
+  ): Unit = throw new AnnexWorkerError(1)
+}
diff --git a/tests/worker-error/test b/tests/worker-error/test
index 1ce3c9af..786d06c3 100755
--- a/tests/worker-error/test
+++ b/tests/worker-error/test
@@ -1,5 +1,12 @@
 #!/bin/bash -e
 . "$(dirname "$0")"/../common.sh
 
+worker_run_directly_returns_proper_exit_code() {
+    bazel run :worker-throwing-annexworkererror ||
+
+    [ $? -eq 1 ]
+}
+
 bazel test :error-spec
 bazel test :fatal-error-spec |& grep -q 'java.lang.OutOfMemoryError'
+worker_run_directly_returns_proper_exit_code