Close Spring transaction in the same thread that opened it even on timeout (#2482)

* Improve transaction rollbacking in Spring in case of timeout

* Make `ThreadBasedExecutor` mark timed out threads to call `onPhaseTimeout`

Author: IlyaMuravjov

utbot-core/src/main/kotlin/org/utbot/common/ThreadUtil.kt

@@ -1,5 +1,6 @@
 package org.utbot.common
 
+import java.util.WeakHashMap
 import java.util.concurrent.ArrayBlockingQueue
 import java.util.concurrent.TimeUnit
 import kotlin.concurrent.thread
@@ -24,11 +25,15 @@ class ThreadBasedExecutor {
         val threadLocal by threadLocalLazy { ThreadBasedExecutor() }
     }
 
+    // there's no `WeakHashSet`, so we use `WeakHashMap` with dummy values
+    private val timedOutThreads = WeakHashMap<Thread, Unit>()
     private var thread: Thread? = null
 
     private var requestQueue = ArrayBlockingQueue<() -> Any?>(1)
     private var responseQueue = ArrayBlockingQueue<Result<Any?>>(1)
 
+    fun isCurrentThreadTimedOut(): Boolean =
+        Thread.currentThread() in timedOutThreads
 
     /**
      * Invoke [action] with timeout.
@@ -59,6 +64,7 @@ class ThreadBasedExecutor {
         if (res == null) {
             try {
                 val t = thread ?: return res
+                timedOutThreads[t] = Unit
                 t.interrupt()
                 t.join(10)
                 if (t.isAlive)

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/context/InstrumentationContext.kt

@@ -4,6 +4,7 @@ import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.UtConcreteValue
 import org.utbot.framework.plugin.api.UtModel
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtModelWithCompositeOriginConstructor
+import org.utbot.instrumentation.instrumentation.execution.phases.ExecutionPhase
 import java.lang.reflect.Method
 import java.util.IdentityHashMap
 import org.utbot.instrumentation.instrumentation.mock.computeKeyForMethod
@@ -34,6 +35,13 @@ interface InstrumentationContext {
      */
     fun findUtModelWithCompositeOriginConstructor(classId: ClassId): UtModelWithCompositeOriginConstructor?
 
+    /**
+     * Called when [timedOutedPhase] times out.
+     * This method is executed in the same thread that [timedOutedPhase] was run in.
+     * Implementor is expected to only perform some clean up operations (e.g. rollback transactions in Spring).
+     */
+    fun onPhaseTimeout(timedOutedPhase: ExecutionPhase)
+
     object MockGetter {
         data class MockContainer(private val values: List<*>) {
             private var ptr: Int = 0

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/context/SimpleInstrumentationContext.kt

@@ -6,6 +6,7 @@ import org.utbot.framework.plugin.api.UtModel
 import org.utbot.framework.plugin.api.util.jClass
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtModelWithCompositeOriginConstructor
 import org.utbot.instrumentation.instrumentation.execution.constructors.javaStdLibModelWithCompositeOriginConstructors
+import org.utbot.instrumentation.instrumentation.execution.phases.ExecutionPhase
 
 /**
  * Simple instrumentation context, that is used for pure JVM projects without
@@ -21,4 +22,6 @@ class SimpleInstrumentationContext : InstrumentationContext {
 
     override fun findUtModelWithCompositeOriginConstructor(classId: ClassId): UtModelWithCompositeOriginConstructor? =
         javaStdLibModelWithCompositeOriginConstructors[classId.jClass]?.invoke()
+
+    override fun onPhaseTimeout(timedOutedPhase: ExecutionPhase) = Unit
 }

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/phases/PhasesController.kt

@@ -1,5 +1,7 @@
 package org.utbot.instrumentation.instrumentation.execution.phases
 
+import com.jetbrains.rd.util.error
+import com.jetbrains.rd.util.getLogger
 import org.utbot.common.StopWatch
 import org.utbot.common.ThreadBasedExecutor
 import org.utbot.framework.plugin.api.Coverage
@@ -17,7 +19,7 @@ import org.utbot.instrumentation.instrumentation.execution.context.Instrumentati
 import java.security.AccessControlException
 
 class PhasesController(
-    instrumentationContext: InstrumentationContext,
+    private val instrumentationContext: InstrumentationContext,
     traceHandler: TraceHandler,
     delegateInstrumentation: Instrumentation<Result<*>>,
     private val timeout: Long
@@ -56,13 +58,23 @@ class PhasesController(
         }
     }
 
+    companion object {
+        private val logger = getLogger<PhasesController>()
+    }
+
     fun <T, R : ExecutionPhase> executePhaseInTimeout(phase: R, block: R.() -> T): T = phase.start {
         val stopWatch = StopWatch()
         val context = UtContext(utContext.classLoader, stopWatch)
         val timeoutForCurrentPhase = timeout - currentlyElapsed
-        val result = ThreadBasedExecutor.threadLocal.invokeWithTimeout(timeout - currentlyElapsed, stopWatch) {
+        val executor = ThreadBasedExecutor.threadLocal
+        val result = executor.invokeWithTimeout(timeout - currentlyElapsed, stopWatch) {
             withUtContext(context) {
-                phase.block()
+                try {
+                    phase.block()
+                } finally {
+                    if (executor.isCurrentThreadTimedOut())
+                        instrumentationContext.onPhaseTimeout(phase)
+                }
             }
         } ?: throw TimeoutException("Timeout $timeoutForCurrentPhase ms for phase ${phase.javaClass.simpleName} elapsed, controller timeout - $timeout")
 

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/spring/SpringInstrumentationContext.kt

@@ -11,6 +11,7 @@ import org.utbot.framework.plugin.api.util.isSubtypeOf
 import org.utbot.framework.plugin.api.util.utContext
 import org.utbot.instrumentation.instrumentation.execution.constructors.UtModelWithCompositeOriginConstructor
 import org.utbot.instrumentation.instrumentation.execution.context.InstrumentationContext
+import org.utbot.instrumentation.instrumentation.execution.phases.ExecutionPhase
 import org.utbot.spring.api.SpringApi
 import org.utbot.spring.api.provider.SpringApiProviderFacade
 import org.utbot.spring.api.provider.InstantiationSettings
@@ -52,4 +53,7 @@ class SpringInstrumentationContext(
     override fun findUtModelWithCompositeOriginConstructor(classId: ClassId): UtModelWithCompositeOriginConstructor? =
         if (classId.isSubtypeOf(resultActionsClassId)) UtMockMvcResultActionsModelConstructor()
         else delegateInstrumentationContext.findUtModelWithCompositeOriginConstructor(classId)
+
+    override fun onPhaseTimeout(timedOutedPhase: ExecutionPhase) =
+        springApi.afterTestMethod()
 }

utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/spring/SpringUtExecutionInstrumentation.kt

@@ -65,7 +65,7 @@ class SpringUtExecutionInstrumentation(
         arguments: ArgumentList,
         parameters: Any?,
         phasesWrapper: PhasesController.(invokeBasePhases: () -> UtConcreteExecutionResult) -> UtConcreteExecutionResult
-    ): UtConcreteExecutionResult {
+    ): UtConcreteExecutionResult = synchronized(this) {
         getRelevantBeans(clazz).forEach { beanName -> springApi.resetBean(beanName) }
         return delegateInstrumentation.invoke(clazz, methodSignature, arguments, parameters) { invokeBasePhases ->
             phasesWrapper {

utbot-spring-commons/src/main/kotlin/org/utbot/spring/SpringApiImpl.kt

@@ -152,6 +152,10 @@ class SpringApiImpl(
     }
 
     override fun afterTestMethod() {
+        if (!isInsideTestMethod) {
+            logger.warn { "afterTestMethod() was probably called twice, ignoring second call" }
+            return
+        }
         testContextManager.afterTestExecution(dummyTestClassInstance, dummyTestMethod, null)
         testContextManager.afterTestMethod(dummyTestClassInstance, dummyTestMethod, null)
         isInsideTestMethod = false