Kapt: Handle exceptions during annotation processing gracefully

(cherry picked from commit ee26c19)

plugins/annotation-processing/src/org/jetbrains/kotlin/annotation/processing/AnnotationProcessingExtension.kt

@@ -41,6 +41,8 @@ import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisCompletedHandlerExtension
 import java.io.File
 import java.io.IOException
+import java.io.PrintWriter
+import java.io.StringWriter
 import java.net.URLClassLoader
 import java.util.*
 import javax.annotation.processing.Processor
@@ -132,8 +134,18 @@ abstract class AbstractAnnotationProcessingExtension(
                 elements, types, messager, options, filer, processors,
                 project, psiManager, javaPsiFacade, projectScope, bindingTrace.bindingContext, appendJavaSourceRootsHandler)
 
-        val processingResult = processingEnvironment.doAnnotationProcessing(files)
+        var processingResult: ProcessingResult
-        processingEnvironment.dispose()
+        try {
+            processingResult = processingEnvironment.doAnnotationProcessing(files)
+            processingEnvironment.dispose()
+        }
+        catch (thr: Throwable) {
+            messager.printMessage(
+                    Diagnostic.Kind.ERROR,
+                    "An exception occurred during annotation processing. Stacktrace: \n" + thr.getStackTraceAsString())
+
+            processingResult = ProcessingResult(errorCount = 1, warningCount = 0, wasAnythingGenerated = false)
+        }
 
         annotationProcessingComplete = true
         log {
@@ -163,6 +175,18 @@ abstract class AbstractAnnotationProcessingExtension(
                 listOf(generatedSourcesOutputDir),
                 addToEnvironment = false)
     }
+
+    private fun Throwable.getStackTraceAsString(): String {
+        val out = StringWriter(1024)
+        val printWriter = PrintWriter(out)
+        try {
+            printStackTrace(printWriter)
+            return out.toString().replace("\r", "")
+        }
+        finally {
+            printWriter.close()
+        }
+    }
     
     private fun KotlinProcessingEnvironment.createTypeMapper(): KotlinTypeMapper {
         return KotlinTypeMapper(bindingContext(), ClassBuilderMode.full(false), NoResolveFileClassesProvider,

plugins/plugins-tests/tests/org/jetbrains/kotlin/annotation/processing/test/processor/AbstractProcessorTest.kt

@@ -16,12 +16,15 @@
 
 package org.jetbrains.kotlin.annotation.processing.test.processor
 
+import com.intellij.openapi.extensions.Extensions
 import com.intellij.testFramework.registerServiceInstance
 import org.jetbrains.kotlin.annotation.AbstractAnnotationProcessingExtension
+import org.jetbrains.kotlin.annotation.processing.diagnostic.DefaultErrorMessagesAnnotationProcessing
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.codegen.AbstractBytecodeTextTest
 import org.jetbrains.kotlin.codegen.CodegenTestUtil
+import org.jetbrains.kotlin.diagnostics.rendering.DefaultErrorMessages
 import org.jetbrains.kotlin.incremental.SourceRetentionAnnotationHandlerImpl
 import org.jetbrains.kotlin.java.model.elements.JeAnnotationMirror
 import org.jetbrains.kotlin.java.model.elements.JeMethodExecutableElement
@@ -74,6 +77,9 @@ abstract class AbstractProcessorTest : AbstractBytecodeTextTest() {
 
         project.registerServiceInstance(JeElementRegistry::class.java, JeElementRegistry())
 
+        Extensions.getRootArea().getExtensionPoint(DefaultErrorMessages.Extension.EP_NAME)
+                .registerExtension(DefaultErrorMessagesAnnotationProcessing())
+
         return environment
     }
 

plugins/plugins-tests/tests/org/jetbrains/kotlin/annotation/processing/test/processor/ProcessorTests.kt

@@ -362,4 +362,22 @@ class ProcessorTests : AbstractProcessorTest() {
         test("c", "java.util.Map<java.lang.Integer,Test>")
         test("d", "java.util.Map<java.lang.Integer,Test>")
     }
+
+    fun testExceptionDuringAp() {
+        class HiThere : RuntimeException()
+
+        var kotlinEnv: KotlinProcessingEnvironment? = null
+        try {
+            test("MapMutableMap", "*") { set, roundEnv, env ->
+                kotlinEnv = env as KotlinProcessingEnvironment
+                throw HiThere()
+            }
+        } catch (e: IllegalStateException) {
+            assertTrue(e.message!!.startsWith("ANNOTATION_PROCESSING_ERROR"))
+        }
+        val env = kotlinEnv!!
+
+        assertEquals(1, env.messager.errorCount)
+        assertEquals(0, env.messager.warningCount)
+    }
 }