Skip to content

Commit 35011d7

Browse files
tamarinvs19tamarinvs19
committed
Support SIGTERM cancellation
1 parent 7aee092 commit 35011d7

File tree

6 files changed

+71
-33
lines changed

6 files changed

+71
-33
lines changed

‎utbot-cli-python/src/main/kotlin/org/utbot/cli/language/python/sbft/SbftGenerateTestsCommand.kt‎

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ import org.utbot.python.utils.RequirementsInstaller
3737
import org.utbot.python.utils.Success
3838
import org.utbot.python.utils.separateTimeout
3939
import java.io.File
40+
import java.util.concurrent.atomic.AtomicBoolean
4041
import kotlin.system.measureTimeMillis
4142

4243
private const val DEFAULT_TIMEOUT_IN_MILLIS = 60000L
@@ -140,6 +141,39 @@ class SbftGenerateTestsCommand : CliktCommand(
140141
return methods + listOf(functions)
141142
}
142143

144+
private val globalImportCollection = mutableSetOf<PythonImport>()
145+
private val globalCodeCollection = mutableListOf<String>()
146+
147+
private val shutdown: AtomicBoolean = AtomicBoolean(false)
148+
private val alreadySaved: AtomicBoolean = AtomicBoolean(false)
149+
150+
private fun shutdownHook() {
151+
Runtime.getRuntime().addShutdownHook(object : Thread() {
152+
override fun run() {
153+
shutdown.set(true)
154+
try {
155+
if (!alreadySaved.get()) {
156+
saveTestsEndExit()
157+
}
158+
} catch (_: InterruptedException) {
159+
logger.warn { "Interrupted exception" }
160+
}
161+
}
162+
})
163+
}
164+
165+
private fun saveTestsEndExit() {
166+
logger.info("Saving tests...")
167+
val importCode = globalImportCollection
168+
.sortedBy { it.order }
169+
.map { renderPythonImport(it) }
170+
val testCode = (listOf(importCode.joinToString("\n")) + globalCodeCollection).joinToString("\n\n\n")
171+
writeToFileAndSave(output, testCode)
172+
173+
Cleaner.doCleaning()
174+
alreadySaved.set(true)
175+
}
176+
143177
override fun run() {
144178
absPathToSourceFile = sourceFile.toAbsolutePath()
145179
sourceFileContent = File(absPathToSourceFile).readText()
@@ -179,8 +213,8 @@ class SbftGenerateTestsCommand : CliktCommand(
179213
}
180214
logger.info { "Mypy time: $mypyTime" }
181215

182-
val globalImportCollection = mutableSetOf<PythonImport>()
183-
val globalCodeCollection = mutableListOf<String>()
216+
shutdownHook()
217+
184218
val startTime = System.currentTimeMillis()
185219
val countOfFunctions = pythonMethodGroups.sumOf { it.size }
186220
val timeoutAfterMypy = if (includeMypyAnalysisTime) timeout - mypyTime else timeout
@@ -208,7 +242,7 @@ class SbftGenerateTestsCommand : CliktCommand(
208242
testFramework = testFramework,
209243
testSourceRootPath = null,
210244
withMinimization = !doNotMinimize,
211-
isCanceled = { false },
245+
isCanceled = { shutdown.get() },
212246
runtimeExceptionTestsBehaviour = RuntimeExceptionTestsBehaviour.valueOf(runtimeExceptionTestsBehaviour),
213247
coverageMeasureMode = PythonCoverageMode.parse(coverageMeasureMode),
214248
sendCoverageContinuously = !doNotSendCoverageContinuously,
@@ -227,14 +261,6 @@ class SbftGenerateTestsCommand : CliktCommand(
227261
globalImportCollection.addAll(imports)
228262
}
229263
}
230-
logger.info("Saving tests...")
231-
val importCode = globalImportCollection
232-
.sortedBy { it.order }
233-
.map { renderPythonImport(it) }
234-
val testCode = (listOf(importCode.joinToString("\n")) + globalCodeCollection).joinToString("\n\n\n")
235-
writeToFileAndSave(output, testCode)
236-
237-
Cleaner.doCleaning()
238-
System.exit(0)
264+
saveTestsEndExit()
239265
}
240266
}

‎utbot-python/src/main/kotlin/org/utbot/python/PythonTestGenerationProcessor.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,9 @@ abstract class PythonTestGenerationProcessor {
6868
val startTime = System.currentTimeMillis()
6969

7070
val tests = configuration.testedMethods.mapIndexedNotNull { index, methodHeader ->
71+
if (configuration.isCanceled()) {
72+
return emptyList()
73+
}
7174
val usedTime = System.currentTimeMillis() - startTime
7275
val expectedTime = index * oneFunctionTimeout
7376
val localOneFunctionTimeout = if (usedTime < expectedTime) {

‎utbot-python/src/main/kotlin/org/utbot/python/engine/GlobalPythonEngine.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ class GlobalPythonEngine(
101101
fun run() {
102102
val fuzzing = thread(
103103
start = true,
104-
isDaemon = true,
104+
isDaemon = false,
105105
name = "Fuzzer"
106106
) {
107107
logger.info { " >>>>>>> Start fuzzer >>>>>>> " }
@@ -110,7 +110,7 @@ class GlobalPythonEngine(
110110
}
111111
val symbolic = thread(
112112
start = true,
113-
isDaemon = true,
113+
isDaemon = false,
114114
name = "Symbolic"
115115
) {
116116
logger.info { " ------- Start symbolic ------- " }

‎utbot-python/src/main/kotlin/org/utbot/python/engine/fuzzing/FuzzingEngine.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ class FuzzingEngine(
218218
}
219219
} else {
220220
try {
221-
PythonFuzzing(typeStorage, typeInferenceAlgorithm) { description, arguments ->
221+
PythonFuzzing(typeStorage, typeInferenceAlgorithm, isCancelled) { description, arguments ->
222222
if (isCancelled()) {
223223
logger.debug { "Fuzzing process was interrupted" }
224224
manager.disconnect()

‎utbot-python/src/main/kotlin/org/utbot/python/engine/symbolic/USVMPythonAnalysisResultReceiverImpl.kt‎

Lines changed: 22 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -48,25 +48,29 @@ class USVMPythonAnalysisResultReceiverImpl(
4848
connect()
4949
}
5050

51-
fun connect() {
52-
serverSocket = ServerSocket(0)
53-
manager =
54-
PythonWorkerManager(
55-
serverSocket,
56-
configuration.pythonPath,
57-
until,
58-
configuration.coverageMeasureMode,
59-
configuration.sendCoverageContinuously,
60-
) {
61-
PythonCodeSocketExecutor(
62-
method,
63-
configuration.testFileInformation.moduleName,
51+
private fun connect() {
52+
try {
53+
serverSocket = ServerSocket(0)
54+
manager =
55+
PythonWorkerManager(
56+
serverSocket,
6457
configuration.pythonPath,
65-
configuration.sysPathDirectories,
66-
configuration.timeoutForRun,
67-
it,
68-
)
69-
}
58+
until,
59+
configuration.coverageMeasureMode,
60+
configuration.sendCoverageContinuously,
61+
) {
62+
PythonCodeSocketExecutor(
63+
method,
64+
configuration.testFileInformation.moduleName,
65+
configuration.pythonPath,
66+
configuration.sysPathDirectories,
67+
configuration.timeoutForRun,
68+
it,
69+
)
70+
}
71+
} catch (_: TimeoutException) {
72+
close()
73+
}
7074
}
7175

7276
fun close() {

‎utbot-python/src/main/kotlin/org/utbot/python/fuzzing/PythonApi.kt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ fun pythonAnyTypeValueProviders() = listOf(
104104
class PythonFuzzing(
105105
private val pythonTypeStorage: PythonTypeHintsStorage,
106106
private val typeInferenceAlgorithm: BaselineAlgorithm,
107+
private val globalIsCancelled: () -> Boolean,
107108
val execute: suspend (description: PythonMethodDescription, values: List<PythonFuzzedValue>) -> PythonFeedback,
108109
) : Fuzzing<UtType, PythonFuzzedValue, PythonMethodDescription, PythonFeedback> {
109110

@@ -165,6 +166,10 @@ class PythonFuzzing(
165166
description: PythonMethodDescription,
166167
stats: Statistic<UtType, PythonFuzzedValue>
167168
): Boolean {
169+
if (globalIsCancelled()) {
170+
logger.warn { "Cancellation in fuzzing" }
171+
return true
172+
}
168173
if (description.limitManager.isCancelled() || description.parameters.any { it.isAny() }) {
169174
forkType(description, stats)
170175
if (description.limitManager.isRootManager) {

0 commit comments

Comments
 (0)