Add socket channel for coverage collection (#2195)

* Coverage algo update

* Update coverage

* Update coverage socket algorithm

Author: tamarinvs19

utbot-python/samples/easy_samples/long_function_coverage.py

@@ -0,0 +1,11 @@
+import time
+
+
+def long_function(x: int):
+    x += 4
+    x /= 2
+    x += 100
+    x *= 3
+    x -= 15
+    time.sleep(2000)
+    return x

utbot-python/src/main/kotlin/org/utbot/python/PythonEngine.kt

@@ -11,6 +11,8 @@ import org.utbot.fuzzing.utils.Trie
 import org.utbot.python.evaluation.*
 import org.utbot.python.evaluation.serialiation.MemoryDump
 import org.utbot.python.evaluation.serialiation.toPythonTree
+import org.utbot.python.evaluation.utils.CoverageIdGenerator
+import org.utbot.python.evaluation.utils.coveredLinesToInstructions
 import org.utbot.python.framework.api.python.PythonTree
 import org.utbot.python.framework.api.python.PythonTreeModel
 import org.utbot.python.framework.api.python.PythonTreeWrapper
@@ -84,6 +86,7 @@ class PythonEngine(
     private fun handleTimeoutResult(
         arguments: List<PythonFuzzedValue>,
         methodUnderTestDescription: PythonMethodDescription,
+        coveredLines: Collection<Int>,
     ): FuzzingExecutionFeedback {
         val summary = arguments
             .zip(methodUnderTest.arguments)
@@ -100,12 +103,16 @@ class PythonEngine(
         val beforeThisObject = beforeThisObjectTree?.let { PythonTreeModel(it.tree) }
         val beforeModelList = beforeModelListTree.map { PythonTreeModel(it.tree) }
 
+        val coveredInstructions = coveredLinesToInstructions(coveredLines, methodUnderTest)
+        val coverage = Coverage(coveredInstructions)
+
         val utFuzzedExecution = PythonUtExecution(
             stateInit = EnvironmentModels(beforeThisObject, beforeModelList, emptyMap()),
             stateBefore = EnvironmentModels(beforeThisObject, beforeModelList, emptyMap()),
             stateAfter = EnvironmentModels(beforeThisObject, beforeModelList, emptyMap()),
             diffIds = emptyList(),
             result = executionResult,
+            coverage = coverage,
             testMethodName = testMethodName.testName?.camelToSnakeCase(),
             displayName = testMethodName.displayName,
             summary = summary.map { DocRegularStmt(it) }
@@ -218,7 +225,9 @@ class PythonEngine(
                     methodUnderTest.argumentsNames
                 )
                 try {
-                    return when (val evaluationResult = manager.run(functionArguments, localAdditionalModules)) {
+                    val coverageId = CoverageIdGenerator.createId()
+                    return when (val evaluationResult =
+                        manager.runWithCoverage(functionArguments, localAdditionalModules, coverageId)) {
                         is PythonEvaluationError -> {
                             val utError = UtError(
                                 "Error evaluation: ${evaluationResult.status}, ${evaluationResult.message}",
@@ -229,8 +238,20 @@ class PythonEngine(
                         }
 
                         is PythonEvaluationTimeout -> {
-                            val utTimeoutException = handleTimeoutResult(arguments, description)
-                            PythonExecutionResult(utTimeoutException, PythonFeedback(control = Control.PASS))
+                            val coveredLines =
+                                manager.coverageReceiver.coverageStorage.getOrDefault(coverageId, mutableSetOf())
+                            val utTimeoutException = handleTimeoutResult(arguments, description, coveredLines)
+                            val coveredInstructions = coveredLinesToInstructions(coveredLines, methodUnderTest)
+                            val trieNode: Trie.Node<Instruction> =
+                                if (coveredInstructions.isEmpty())
+                                    Trie.emptyNode()
+                                else description.tracer.add(
+                                    coveredInstructions
+                                )
+                            PythonExecutionResult(
+                                utTimeoutException,
+                                PythonFeedback(control = Control.PASS, result = trieNode)
+                            )
                         }
 
                         is PythonEvaluationSuccess -> {

utbot-python/src/main/kotlin/org/utbot/python/evaluation/CodeEvaluationApi.kt

@@ -17,6 +17,12 @@ interface PythonCodeExecutor {
         additionalModulesToImport: Set<String>
     ): PythonEvaluationResult
 
+    fun runWithCoverage(
+        fuzzedValues: FunctionArguments,
+        additionalModulesToImport: Set<String>,
+        coverageId: String,
+    ): PythonEvaluationResult
+
     fun stop()
 }
 

utbot-python/src/main/kotlin/org/utbot/python/evaluation/PythonCodeSocketExecutor.kt

@@ -12,6 +12,7 @@ import org.utbot.python.evaluation.serialiation.FailExecution
 import org.utbot.python.evaluation.serialiation.PythonExecutionResult
 import org.utbot.python.evaluation.serialiation.SuccessExecution
 import org.utbot.python.evaluation.serialiation.serializeObjects
+import org.utbot.python.evaluation.utils.CoverageIdGenerator
 import org.utbot.python.framework.api.python.util.pythonAnyClassId
 import org.utbot.python.newtyping.pythonTypeName
 import org.utbot.python.newtyping.pythonTypeRepresentation
@@ -48,6 +49,15 @@ class PythonCodeSocketExecutor(
     override fun run(
         fuzzedValues: FunctionArguments,
         additionalModulesToImport: Set<String>
+    ): PythonEvaluationResult {
+        val coverageId = CoverageIdGenerator.createId()
+        return runWithCoverage(fuzzedValues, additionalModulesToImport, coverageId)
+    }
+
+    override fun runWithCoverage(
+        fuzzedValues: FunctionArguments,
+        additionalModulesToImport: Set<String>,
+        coverageId: String
     ): PythonEvaluationResult {
         val (arguments, memory) = serializeObjects(fuzzedValues.allArguments.map { it.tree })
 
@@ -69,6 +79,7 @@ class PythonCodeSocketExecutor(
             emptyMap(),  // here can be only-kwargs arguments
             memory,
             method.moduleFilename,
+            coverageId,
         )
         val message = ExecutionRequestSerializer.serializeRequest(request) ?: error("Cannot serialize request to python executor")
         try {
@@ -81,6 +92,7 @@ class PythonCodeSocketExecutor(
         val (status, response) = UtExecutorThread.run(pythonWorker, executionTimeout)
         return when (status) {
             UtExecutorThread.Status.TIMEOUT -> {
+
                 PythonEvaluationTimeout()
             }
 

utbot-python/src/main/kotlin/org/utbot/python/evaluation/PythonCoverageReceiver.kt

@@ -0,0 +1,37 @@
+package org.utbot.python.evaluation
+
+import mu.KotlinLogging
+import java.io.IOException
+import java.net.DatagramPacket
+import java.net.DatagramSocket
+import java.net.SocketException
+
+class PythonCoverageReceiver(
+    private val until: Long,
+) : Thread() {
+    val coverageStorage = mutableMapOf<String, MutableSet<Int>>()
+    private val socket = DatagramSocket()
+    private val logger = KotlinLogging.logger {}
+
+    fun address(): Pair<String, String> {
+        return "localhost" to socket.localPort.toString()
+    }
+
+    override fun run() {
+        try {
+            while (System.currentTimeMillis() < until) {
+                val buf = ByteArray(256)
+                val request = DatagramPacket(buf, buf.size)
+                socket.receive(request)
+                val (id, line) = request.data.decodeToString().take(request.length).split(":")
+                logger.debug { "Got coverage: $id, $line" }
+                val lineNumber = line.toInt()
+                coverageStorage.getOrPut(id) { mutableSetOf() } .add(lineNumber)
+            }
+        } catch (ex: SocketException) {
+            logger.error("Socket error: " + ex.message)
+        } catch (ex: IOException) {
+            logger.error("IO error: " + ex.message)
+        }
+    }
+}

utbot-python/src/main/kotlin/org/utbot/python/evaluation/PythonWorkerManager.kt

@@ -26,7 +26,10 @@ class PythonWorkerManager(
     private lateinit var workerSocket: Socket
     private lateinit var codeExecutor: PythonCodeExecutor
 
+    val coverageReceiver = PythonCoverageReceiver(until)
+
     init {
+        coverageReceiver.start()
         connect()
     }
 
@@ -37,6 +40,8 @@ class PythonWorkerManager(
             "-m", "utbot_executor",
             "localhost",
             serverSocket.localPort.toString(),
+            coverageReceiver.address().first,
+            coverageReceiver.address().second,
             "--logfile", logfile.absolutePath,
             "--loglevel", "INFO",  // "DEBUG", "INFO", "WARNING", "ERROR"
         ))
@@ -60,13 +65,32 @@ class PythonWorkerManager(
     fun disconnect() {
         workerSocket.close()
         process.destroy()
+        coverageReceiver.interrupt()
     }
 
     fun reconnect() {
         disconnect()
         connect()
     }
 
+    fun runWithCoverage(
+        fuzzedValues: FunctionArguments,
+        additionalModulesToImport: Set<String>,
+        coverageId: String
+    ): PythonEvaluationResult {
+        val evaluationResult = try {
+            codeExecutor.runWithCoverage(fuzzedValues, additionalModulesToImport, coverageId)
+        } catch (_: SocketTimeoutException) {
+            logger.debug { "Socket timeout" }
+            reconnect()
+            PythonEvaluationTimeout()
+        }
+        if (evaluationResult is PythonEvaluationError || evaluationResult is PythonEvaluationTimeout) {
+            reconnect()
+        }
+        return evaluationResult
+    }
+
     fun run(
         fuzzedValues: FunctionArguments,
         additionalModulesToImport: Set<String>

utbot-python/src/main/kotlin/org/utbot/python/evaluation/serialiation/ExecutionRequestSerializer.kt

@@ -24,4 +24,5 @@ data class ExecutionRequest(
     val kwargumentsIds: Map<String, String>,
     val serializedMemory: String,
     val filepath: String,
+    val coverageId: String,
 )

utbot-python/src/main/kotlin/org/utbot/python/evaluation/utils/CoverageIdGenerator.kt

@@ -0,0 +1,13 @@
+package org.utbot.python.evaluation.utils
+
+import java.util.concurrent.atomic.AtomicLong
+
+object CoverageIdGenerator {
+    private const val lower_bound: Long = 1500_000_000
+
+    private val lastId: AtomicLong = AtomicLong(lower_bound)
+
+    fun createId(): String {
+        return lastId.incrementAndGet().toString(radix = 16)
+    }
+}

utbot-python/src/main/kotlin/org/utbot/python/evaluation/utils/Utils.kt

@@ -0,0 +1,17 @@
+package org.utbot.python.evaluation.utils
+
+import org.utbot.framework.plugin.api.Instruction
+import org.utbot.python.PythonMethod
+import org.utbot.python.framework.api.python.util.pythonAnyClassId
+import org.utbot.python.newtyping.pythonTypeRepresentation
+
+fun coveredLinesToInstructions(coveredLines: Collection<Int>, method: PythonMethod): List<Instruction> {
+    return coveredLines.map {
+        Instruction(
+            method.containingPythonClass?.pythonTypeRepresentation() ?: pythonAnyClassId.name,
+            method.methodSignature(),
+            it,
+            it.toLong()
+        )
+    }
+}

utbot-python/src/main/kotlin/org/utbot/python/utils/RequirementsUtils.kt

@@ -3,7 +3,7 @@ package org.utbot.python.utils
 object RequirementsUtils {
     val requirements: List<String> = listOf(
         "mypy==1.0.0",
-        "utbot-executor==1.4.21",
+        "utbot-executor==1.4.26",
         "utbot-mypy-runner==0.2.8",
     )