Handle all failures in Future

Author: bell-db

bridge/src/main/scala/protocbridge/frontend/PluginFrontend.scala

@@ -47,13 +47,7 @@ object PluginFrontend {
       gen: ProtocCodeGenerator,
       request: Array[Byte]
   ): Array[Byte] = {
-    Try {
-      gen.run(request)
-    }.recover { case throwable =>
-      createCodeGeneratorResponseWithError(
-        throwable.toString + "\n" + getStackTrace(throwable)
-      )
-    }.get
+    gen.run(request)
   }
 
   def createCodeGeneratorResponseWithError(error: String): Array[Byte] = {
@@ -116,9 +110,17 @@ object PluginFrontend {
       gen: ProtocCodeGenerator,
       fsin: InputStream,
       env: ExtraEnv
-  ): Array[Byte] = {
+  ): Array[Byte] = try {
     val bytes = readInputStreamToByteArrayWithEnv(fsin, env)
     runWithBytes(gen, bytes)
+  } catch {
+    // This covers all Throwable including OutOfMemoryError, StackOverflowError, etc.
+    // We need to make a best effort to return a response to protoc,
+    // otherwise protoc can hang indefinitely.
+    case throwable: Throwable =>
+      createCodeGeneratorResponseWithError(
+        throwable.toString + "\n" + getStackTrace(throwable)
+      )
   }
 
   def createTempFile(extension: String, content: String): Path = {

bridge/src/main/scala/protocbridge/frontend/PosixPluginFrontend.scala

@@ -40,6 +40,9 @@ object PosixPluginFrontend extends PluginFrontend {
         val response = PluginFrontend.runWithInputStream(plugin, fsin, env)
         fsin.close()
 
+        // Note that the output pipe must be opened after the input pipe is closed.
+        // Otherwise, there might be a deadlock if the shell script is stuck in input pipe and
+        // doesn't open the other end of the output pipe.
         val fsout = Files.newOutputStream(outputPipe)
         fsout.write(response)
         fsout.close()

bridge/src/test/scala/protocbridge/frontend/OsSpecificFrontendSpec.scala

@@ -11,27 +11,21 @@ import scala.util.Random
 
 class OsSpecificFrontendSpec extends AnyFlatSpec with Matchers {
 
-  protected def testPluginFrontend(frontend: PluginFrontend): Array[Byte] = {
-    val random = new Random()
-    val toSend = Array.fill(123)(random.nextInt(256).toByte)
-    val toReceive = Array.fill(456)(random.nextInt(256).toByte)
-    val env = new ExtraEnv(secondaryOutputDir = "tmp")
-
-    val fakeGenerator = new ProtocCodeGenerator {
-      override def run(request: Array[Byte]): Array[Byte] = {
-        request mustBe (toSend ++ env.toByteArrayAsField)
-        toReceive
-      }
-    }
+  protected def testPluginFrontend(
+      frontend: PluginFrontend,
+      generator: ProtocCodeGenerator,
+      env: ExtraEnv,
+      request: Array[Byte]
+  ): Array[Byte] = {
     val (path, state) = frontend.prepare(
-      fakeGenerator,
+      generator,
       env
     )
     val actualOutput = new ByteArrayOutputStream()
     val process = sys.process
       .Process(path.toAbsolutePath.toString)
       .run(new ProcessIO(writeInput => {
-        writeInput.write(toSend)
+        writeInput.write(request)
         writeInput.close()
       }, processOutput => {
         val buffer = new Array[Byte](4096)
@@ -48,4 +42,34 @@ class OsSpecificFrontendSpec extends AnyFlatSpec with Matchers {
     frontend.cleanup(state)
     actualOutput.toByteArray
   }
+
+  protected def testSuccess(frontend: PluginFrontend): Unit = {
+    val random = new Random()
+    val toSend = Array.fill(123)(random.nextInt(256).toByte)
+    val toReceive = Array.fill(456)(random.nextInt(256).toByte)
+    val env = new ExtraEnv(secondaryOutputDir = "tmp")
+
+    val fakeGenerator = new ProtocCodeGenerator {
+      override def run(request: Array[Byte]): Array[Byte] = {
+        request mustBe (toSend ++ env.toByteArrayAsField)
+        toReceive
+      }
+    }
+    val response = testPluginFrontend(frontend, fakeGenerator, env, toSend)
+    response mustBe toReceive
+  }
+
+  protected def testFailure(frontend: PluginFrontend): Unit = {
+    val random = new Random()
+    val toSend = Array.fill(123)(random.nextInt(256).toByte)
+    val env = new ExtraEnv(secondaryOutputDir = "tmp")
+
+    val fakeGenerator = new ProtocCodeGenerator {
+      override def run(request: Array[Byte]): Array[Byte] = {
+        throw new OutOfMemoryError("test error")
+      }
+    }
+    val response = testPluginFrontend(frontend, fakeGenerator, env, toSend)
+    response.length must be > 0
+  }
 }

bridge/src/test/scala/protocbridge/frontend/PosixPluginFrontendSpec.scala

@@ -3,7 +3,11 @@ package protocbridge.frontend
 class PosixPluginFrontendSpec extends OsSpecificFrontendSpec {
   if (!PluginFrontend.isWindows) {
     it must "execute a program that forwards input and output to given stream" in {
-      testPluginFrontend(PosixPluginFrontend)
+      testSuccess(PosixPluginFrontend)
+    }
+
+    it must "not hang if there is an OOM in generator" in {
+      testFailure(PosixPluginFrontend)
     }
   }
 }

bridge/src/test/scala/protocbridge/frontend/WindowsPluginFrontendSpec.scala

@@ -3,7 +3,11 @@ package protocbridge.frontend
 class WindowsPluginFrontendSpec extends OsSpecificFrontendSpec {
   if (PluginFrontend.isWindows) {
     it must "execute a program that forwards input and output to given stream" in {
-      testPluginFrontend(WindowsPluginFrontend)
+      testSuccess(WindowsPluginFrontend)
+    }
+
+    it must "not hang if there is an OOM in generator" in {
+      testFailure(WindowsPluginFrontend)
     }
   }
 }