fix: replace ThreadLocal with TransmittableThreadLocal in ThreadCover…

test2code/build.gradle.kts

@@ -45,6 +45,7 @@ dependencies {
     implementation(project(":test2code-common"))
     implementation(project(":test2code-jacoco"))
     implementation(project(":konform"))
+    implementation(project(":transmittable-thread-local"))
 
     testImplementation(kotlin("test-junit"))
 }

test2code/src/main/kotlin/com/epam/drill/agent/test2code/Test2Code.kt

@@ -63,7 +63,8 @@ class Test2Code(
         intervalMs = configuration.parameters[Test2CodeParameterDefinitions.COVERAGE_SEND_INTERVAL],
         pageSize = configuration.parameters[Test2CodeParameterDefinitions.COVERAGE_SEND_PAGE_SIZE],
         sender = sender,
-        collectProbes = { coverageManager.pollRecorded() }
+        collectReleasedProbes = { coverageManager.pollRecorded() },
+        collectUnreleasedProbes = { coverageManager.getUnreleased() }
     )
     private val coverageCollectionEnabled = configuration.parameters[COVERAGE_COLLECTION_ENABLED]
     private val classScanningEnabled = configuration.parameters[Test2CodeParameterDefinitions.CLASS_SCANNING_ENABLED]

test2code/src/main/kotlin/com/epam/drill/agent/test2code/coverage/CoverageRecorder.kt

@@ -20,6 +20,7 @@ interface ICoverageRecorder {
     fun stopRecording(sessionId: String?, testId: String?)
     fun getContext(): ContextCoverage?
     fun pollRecorded(): Sequence<ExecDatum>
+    fun getUnreleased(): Sequence<ExecDatum> = emptySequence()
 }
 
 data class ContextCoverage(val context: ContextKey, val execData: ExecData)

test2code/src/main/kotlin/com/epam/drill/agent/test2code/coverage/CoverageSender.kt

@@ -37,15 +37,16 @@ class IntervalCoverageSender(
     private val intervalMs: Long,
     private val pageSize: Int,
     private val sender: AgentMessageSender = StubSender(),
-    private val collectProbes: () -> Sequence<ExecDatum> = { emptySequence() }
+    private val collectReleasedProbes: () -> Sequence<ExecDatum> = { emptySequence() },
+    private val collectUnreleasedProbes: () -> Sequence<ExecDatum> = { emptySequence() }
 ) : CoverageSender {
     private val scheduledThreadPool = Executors.newSingleThreadScheduledExecutor()
     private val destination = AgentMessageDestination("POST", "coverage")
     private val logger = KotlinLogging.logger {}
 
     override fun startSendingCoverage() {
         scheduledThreadPool.scheduleAtFixedRate(
-            Runnable { sendProbes(collectProbes()) },
+            Runnable { sendProbes(collectReleasedProbes()) },
             0,
             intervalMs,
             TimeUnit.MILLISECONDS
@@ -59,7 +60,8 @@ class IntervalCoverageSender(
             logger.error("Failed to send some coverage data prior to shutdown")
             scheduledThreadPool.shutdownNow();
         }
-        sendProbes(collectProbes())
+        sendProbes(collectReleasedProbes())
+        sendProbes(collectUnreleasedProbes())
         sender.shutdown()
         logger.info { "Coverage sending job is stopped." }
     }

test2code/src/main/kotlin/com/epam/drill/agent/test2code/coverage/DataPool.kt

@@ -39,6 +39,12 @@ interface DataPool<K: Any, V> {
      */
     fun get(key: K): V?
 
+    /**
+     * Get all objects from the pool
+     * @return a map of all objects in the pool
+     */
+    fun getAll(): Map<K, V>
+
     /**
      * Release the used object by key and move it into the released queue
      * @param key a key
@@ -74,6 +80,10 @@ class ConcurrentDataPool<K : Any, V> : DataPool<K, V> {
         return dataMap[key]
     }
 
+    override fun getAll(): Map<K, V> {
+        return dataMap.toMap()
+    }
+
     override fun pollReleased(): Sequence<V> {
         return generateSequence { released.poll() }
     }

test2code/src/main/kotlin/com/epam/drill/agent/test2code/coverage/ThreadCoverageRecorder.kt

@@ -15,14 +15,17 @@
  */
 package com.epam.drill.agent.test2code.coverage
 
+import com.epam.drill.agent.ttl.TransmittableThreadLocal
 import mu.KotlinLogging
+import kotlin.sequences.filter
+import kotlin.sequences.flatMap
 
 class ThreadCoverageRecorder(
     private val execDataPool: DataPool<ContextKey, ExecData> = ConcurrentDataPool()
 ) : ICoverageRecorder {
     private val logger = KotlinLogging.logger {}
-    private val context: ThreadLocal<ContextKey?> = ThreadLocal()
-    private val execData: ThreadLocal<ExecData?> = ThreadLocal()
+    private val context: ThreadLocal<ContextKey?> = TransmittableThreadLocal()
+    private val execData: ThreadLocal<ExecData?> = TransmittableThreadLocal()
 
     override fun startRecording(sessionId: String?, testId: String?) {
         stopRecording(context.get()?.sessionId, context.get()?.testId)
@@ -50,6 +53,10 @@ class ThreadCoverageRecorder(
             .filter { it.probes.containCovered() }
     }
 
+    override fun getUnreleased(): Sequence<ExecDatum> {
+        return execDataPool.getAll().values.flatMap { it.values }.filter { it.probes.containCovered() }.asSequence()
+    }
+
     override fun getContext(): ContextCoverage? {
         return context.get()?.let { ctx -> execData.get()?.let { ContextCoverage(ctx, it) } }
     }