ProtoLogTool: Fix thread pool leak on code generation exception

Fixes an issue where the thread pool leaks when code generation
throws an exception (for example because of an invalid log string).

When the thread pool leaks, System.exit() no longer works, because
the JVM waits for all threads to exit - which never happens for the
leaked thread pool.

Test: `ProtoLog.w(LOG_GROUP, new String())` -> make droid should not hang
Change-Id: Ibdae0f1e68441c6fe004398d146f496503ccc4cd

1 files changed, 46 insertions, 42 deletions

diff --git a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
index 99a26dc80288..3c55237ce443 100644
--- a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
+++ b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
@@ -53,36 +53,38 @@ object ProtoLogTool {
 
         val executor = newThreadPool()
 
-        command.javaSourceArgs.map { path ->
-            executor.submitCallable {
-                val transformer = SourceTransformer(command.protoLogImplClassNameArg,
-                        command.protoLogCacheClassNameArg, processor)
-                val file = File(path)
-                val text = injector.readText(file)
-                val outSrc = try {
-                    val code = tryParse(text, path)
-                    if (containsProtoLogText(text, command.protoLogClassNameArg)) {
-                        transformer.processClass(text, path, packagePath(file, code), code)
-                    } else {
+        try {
+            command.javaSourceArgs.map { path ->
+                executor.submitCallable {
+                    val transformer = SourceTransformer(command.protoLogImplClassNameArg,
+                            command.protoLogCacheClassNameArg, processor)
+                    val file = File(path)
+                    val text = injector.readText(file)
+                    val outSrc = try {
+                        val code = tryParse(text, path)
+                        if (containsProtoLogText(text, command.protoLogClassNameArg)) {
+                            transformer.processClass(text, path, packagePath(file, code), code)
+                        } else {
+                            text
+                        }
+                    } catch (ex: ParsingException) {
+                        // If we cannot parse this file, skip it (and log why). Compilation will
+                        // fail in a subsequent build step.
+                        injector.reportParseError(ex)
                         text
                     }
-                } catch (ex: ParsingException) {
-                    // If we cannot parse this file, skip it (and log why). Compilation will fail
-                    // in a subsequent build step.
-                    injector.reportParseError(ex)
-                    text
+                    path to outSrc
                 }
-                path to outSrc
+            }.map { future ->
+                val (path, outSrc) = future.get()
+                outJar.putNextEntry(ZipEntry(path))
+                outJar.write(outSrc.toByteArray())
+                outJar.closeEntry()
             }
-        }.map { future ->
-            val (path, outSrc) = future.get()
-            outJar.putNextEntry(ZipEntry(path))
-            outJar.write(outSrc.toByteArray())
-            outJar.closeEntry()
+        } finally {
+            executor.shutdown()
         }
 
-        executor.shutdown()
-
         val cacheSplit = command.protoLogCacheClassNameArg.split(".")
         val cacheName = cacheSplit.last()
         val cachePackage = cacheSplit.dropLast(1).joinToString(".")
@@ -153,30 +155,32 @@ ${updates.replaceIndent("                    ")}
 
         val executor = newThreadPool()
 
-        command.javaSourceArgs.map { path ->
-            executor.submitCallable {
-                val file = File(path)
-                val text = injector.readText(file)
-                if (containsProtoLogText(text, command.protoLogClassNameArg)) {
-                    try {
-                        val code = tryParse(text, path)
-                        builder.findLogCalls(code, path, packagePath(file, code))
-                    } catch (ex: ParsingException) {
-                        // If we cannot parse this file, skip it (and log why). Compilation will fail
-                        // in a subsequent build step.
-                        injector.reportParseError(ex)
+        try {
+            command.javaSourceArgs.map { path ->
+                executor.submitCallable {
+                    val file = File(path)
+                    val text = injector.readText(file)
+                    if (containsProtoLogText(text, command.protoLogClassNameArg)) {
+                        try {
+                            val code = tryParse(text, path)
+                            builder.findLogCalls(code, path, packagePath(file, code))
+                        } catch (ex: ParsingException) {
+                            // If we cannot parse this file, skip it (and log why). Compilation will
+                            // fail in a subsequent build step.
+                            injector.reportParseError(ex)
+                            null
+                        }
+                    } else {
                         null
                     }
-                } else {
-                    null
                 }
+            }.forEach { future ->
+                builder.addLogCalls(future.get() ?: return@forEach)
             }
-        }.forEach { future ->
-            builder.addLogCalls(future.get() ?: return@forEach)
+        } finally {
+            executor.shutdown()
         }
 
-        executor.shutdown()
-
         val out = injector.fileOutputStream(command.viewerConfigJsonArg)
         out.write(builder.build().toByteArray())
         out.close()