Refactor byte-code transformation

build.gradle.kts

@@ -43,6 +43,7 @@ kotlin {
             val kotlinxCoroutinesVersion: String by project
             val asmVersion: String by project
             val reflectionsVersion: String by project
+            val byteBuddyVersion: String by project
             dependencies {
                 api("org.jetbrains.kotlin:kotlin-stdlib:$kotlinVersion")
                 api("org.jetbrains.kotlin:kotlin-stdlib-common:$kotlinVersion")
@@ -51,6 +52,7 @@ kotlin {
                 api("org.ow2.asm:asm-commons:$asmVersion")
                 api("org.ow2.asm:asm-util:$asmVersion")
                 api("org.reflections:reflections:$reflectionsVersion")
+                api("net.bytebuddy:byte-buddy-agent:$byteBuddyVersion")
             }
         }
 

gradle.properties

@@ -24,6 +24,7 @@ asmVersion=9.6
 atomicfuVersion=0.20.2
 # For java.util.* remapping
 reflectionsVersion=0.10.2
+byteBuddyVersion=1.14.0
 
 junitVersion=4.13.1
 jctoolsVersion=3.3.0

src/jvm/main/org/jetbrains/kotlinx/lincheck/Reporter.kt

@@ -12,9 +12,9 @@ package org.jetbrains.kotlinx.lincheck
 
 import org.jetbrains.kotlinx.lincheck.LoggingLevel.*
 import org.jetbrains.kotlinx.lincheck.execution.*
-import org.jetbrains.kotlinx.lincheck.runner.*
 import org.jetbrains.kotlinx.lincheck.strategy.*
 import org.jetbrains.kotlinx.lincheck.strategy.managed.*
+import sun.nio.ch.lincheck.TestThread
 import java.io.*
 import kotlin.math.max
 
@@ -573,7 +573,7 @@ private fun StringBuilder.appendDeadlockWithDumpFailure(failure: DeadlockWithDum
     failure.threadDump?.let { threadDump ->
         // Sort threads to produce same output for the same results
         for ((t, stackTrace) in threadDump.entries.sortedBy { it.key.id }) {
-            val threadNumber = if (t is FixedActiveThreadsExecutor.TestThread) t.iThread.toString() else "?"
+            val threadNumber = (t as? TestThread)?.threadNumber?.toString() ?: "?"
             appendLine("Thread-$threadNumber:")
             stackTrace.map {
                 StackTraceElement(

src/jvm/main/org/jetbrains/kotlinx/lincheck/TransformationClassLoader.java

@@ -10,21 +10,22 @@
 
 package org.jetbrains.kotlinx.lincheck;
 
+import kotlin.Unit;
 import org.jetbrains.kotlinx.lincheck.runner.*;
 import org.jetbrains.kotlinx.lincheck.strategy.*;
 import org.jetbrains.kotlinx.lincheck.strategy.managed.*;
 import org.jetbrains.kotlinx.lincheck.strategy.managed.modelchecking.*;
 import org.jetbrains.kotlinx.lincheck.strategy.stress.*;
+import org.jetbrains.kotlinx.lincheck.transformation.LincheckClassVisitor;
+import org.jetbrains.kotlinx.lincheck.transformation.TransformationMode;
 import org.objectweb.asm.*;
 import org.objectweb.asm.commons.*;
-import org.objectweb.asm.util.*;
 
 import java.io.*;
 import java.net.*;
 import java.util.*;
 import java.util.concurrent.*;
 import java.util.function.*;
-import java.util.stream.Collectors;
 
 import static org.jetbrains.kotlinx.lincheck.TransformationClassLoader.*;
 import static org.jetbrains.kotlinx.lincheck.UtilsKt.getCanonicalClassName;
@@ -49,18 +50,12 @@ public class TransformationClassLoader extends ExecutionClassLoader {
     private static final Map<String, byte[]> stressStrategyBytecodeCache = new ConcurrentHashMap<>();
     private static final Map<String, byte[]> modelCheckingStrategyBytecodeCache = new ConcurrentHashMap<>();
 
-    public TransformationClassLoader(Strategy strategy, Runner runner) {
+    public TransformationClassLoader(Strategy strategy) {
         classTransformers = new ArrayList<>();
         // Apply the strategy's transformer at first, then the runner's one.
-        if (strategy.needsTransformation()) classTransformers.add(strategy::createTransformer);
-        if (runner.needsTransformation()) classTransformers.add(runner::createTransformer);
-        remapper = UtilsKt.getRemapperByTransformers(
-                // create transformers just for their class information
-                classTransformers.stream()
-                        // pass any parameter to Transformer constructors - it won't be used
-                        .map(constructor -> constructor.apply(new TraceClassVisitor(null)))
-                        .collect(Collectors.toList())
-        );
+        TransformationMode transformationMode = strategy instanceof ModelCheckingStrategy ? TransformationMode.MODEL_CHECKING : TransformationMode.STRESS;
+        classTransformers.add((cw) -> new LincheckClassVisitor(transformationMode, cw));
+        remapper = UtilsKt.getRemapperByTransformers(transformationMode);
         if (strategy instanceof StressStrategy) {
             bytecodeCache = stressStrategyBytecodeCache;
         } else if (strategy instanceof ManagedStrategy && ((ManagedStrategy) strategy).useBytecodeCache()) {
@@ -78,7 +73,7 @@ public TransformationClassLoader(Function<ClassVisitor, ClassVisitor> classTrans
      */
     private static boolean doNotTransform(String className) {
         if (className.startsWith(REMAPPED_PACKAGE_CANONICAL_NAME)) return false;
-        if (ManagedStrategyTransformerKt.isImpossibleToTransformApiClass(className)) return true;
+        if (isImpossibleToTransformApiClass(className)) return true;
         return className.startsWith("sun.") ||
                className.startsWith("java.") ||
                className.startsWith("jdk.internal.") ||
@@ -89,8 +84,7 @@ private static boolean doNotTransform(String className) {
                ) ||
                className.startsWith("com.intellij.rt.coverage.") ||
                (className.startsWith("org.jetbrains.kotlinx.lincheck.") &&
-                   !className.startsWith("org.jetbrains.kotlinx.lincheck_test.") &&
-                   !className.equals(ManagedStrategyStateHolder.class.getName())
+                   !className.startsWith("org.jetbrains.kotlinx.lincheck_test.")
                ) ||
                className.equals(kotlinx.coroutines.DebugKt.class.getName()) ||
                className.equals(kotlin.coroutines.CoroutineContext.class.getName()) ||
@@ -104,6 +98,16 @@ private static boolean doNotTransform(String className) {
                className.equals("kotlinx.coroutines.CancelHandlerBase");
     }
 
+    /**
+     * Some API classes cannot be transformed due to the [sun.reflect.CallerSensitive] annotation.
+     */
+    public static boolean isImpossibleToTransformApiClass(String className) {
+        return Objects.equals(className, "sun.misc.Unsafe") ||
+                Objects.equals(className, "jdk.internal.misc.Unsafe") ||
+                Objects.equals(className, "java.lang.invoke.VarHandle") ||
+                className.startsWith("java.util.concurrent.atomic.Atomic") && className.endsWith("FieldUpdater");
+    }
+
     /**
      * Returns `true` if the specified class should not be transformed.
      * Note that this method takes into consideration whether class name will be remapped.

src/jvm/main/org/jetbrains/kotlinx/lincheck/Utils.kt

@@ -13,9 +13,10 @@ import kotlinx.coroutines.*
 import org.jetbrains.kotlinx.lincheck.execution.*
 import org.jetbrains.kotlinx.lincheck.runner.*
 import org.jetbrains.kotlinx.lincheck.strategy.managed.*
+import org.jetbrains.kotlinx.lincheck.transformation.TransformationMode
 import org.jetbrains.kotlinx.lincheck.verifier.*
-import org.objectweb.asm.*
 import org.objectweb.asm.commons.*
+import sun.nio.ch.lincheck.TestThread
 import java.io.*
 import java.lang.ref.*
 import java.lang.reflect.*
@@ -143,7 +144,16 @@ internal class StoreExceptionHandler :
 internal fun <T> CancellableContinuation<T>.cancelByLincheck(promptCancellation: Boolean): CancellationResult {
     val exceptionHandler = context[CoroutineExceptionHandler] as StoreExceptionHandler
     exceptionHandler.exception = null
-    val cancelled = cancel(cancellationByLincheckException)
+
+    val currentThread = Thread.currentThread() as? TestThread
+    val inIgnoredSection = currentThread?.inIgnoredSection ?: false
+    currentThread?.inIgnoredSection = false
+    val cancelled = try {
+        cancel(cancellationByLincheckException)
+    } finally {
+        currentThread?.inIgnoredSection = inIgnoredSection
+    }
+
     exceptionHandler.exception?.let {
         throw it.cause!! // let's throw the original exception, ignoring the internal coroutines details
     }
@@ -172,8 +182,8 @@ object CancellableContinuationHolder {
 
 fun storeCancellableContinuation(cont: CancellableContinuation<*>) {
     val t = Thread.currentThread()
-    if (t is FixedActiveThreadsExecutor.TestThread) {
-        t.cont = cont
+    if (t is TestThread) {
+        t.suspendedContinuation = cont
     } else {
         CancellableContinuationHolder.storedLastCancellableCont = cont
     }
@@ -282,17 +292,17 @@ private class CustomObjectInputStream(val loader: ClassLoader, inputStream: Inpu
  * threads that are related to the specified [runner].
  */
 internal fun collectThreadDump(runner: Runner) = Thread.getAllStackTraces().filter { (t, _) ->
-    t is FixedActiveThreadsExecutor.TestThread && t.runnerHash == runner.hashCode()
+    t is TestThread && t.runnerHash == runner.hashCode()
 }
 
 /**
  * This method helps to encapsulate remapper logic from strategy interface.
  * The remapper is determined based on the used transformers.
  */
-internal fun getRemapperByTransformers(classTransformers: List<ClassVisitor>): Remapper? =
-    when {
-        classTransformers.any { it is ManagedStrategyTransformer } -> JavaUtilRemapper()
-        else -> null
+internal fun getRemapperByTransformers(transformationMode: TransformationMode): Remapper? =
+    when(transformationMode) {
+        TransformationMode.MODEL_CHECKING -> JavaUtilRemapper()
+        TransformationMode.STRESS -> null
     }
 
 internal val String.canonicalClassName get() = this.replace('/', '.')
@@ -336,13 +346,20 @@ internal object InternalLincheckTestUnexpectedException : Exception()
  */
 internal class LincheckInternalBugException(cause: Throwable): Exception(cause)
 
-internal fun stackTraceRepresentation(stackTrace: Array<StackTraceElement>): List<String> {
-    return transformStackTraceBackFromRemapped(stackTrace).map { it.toString() }.filter { line ->
-        "org.jetbrains.kotlinx.lincheck.strategy" !in line
-                && "org.jetbrains.kotlinx.lincheck.runner" !in line
-                && "org.jetbrains.kotlinx.lincheck.UtilsKt" !in line
+internal inline fun <R> runInIgnoredSection(block: () -> R): R =
+    runInIgnoredSection(Thread.currentThread(), block)
+
+private inline fun <R> runInIgnoredSection(currentThread: Thread, block: () -> R): R =
+    if (currentThread is TestThread && !currentThread.inIgnoredSection) {
+        currentThread.inIgnoredSection = true
+        try {
+            block()
+        } finally {
+            currentThread.inIgnoredSection = false
+        }
+    } else {
+        block()
     }
-}
 
 internal fun transformStackTraceBackFromRemapped(stackTrace: Array<StackTraceElement>) = stackTrace.map {
     StackTraceElement(it.className.removePrefix(TransformationClassLoader.REMAPPED_PACKAGE_CANONICAL_NAME), it.methodName, it.fileName, it.lineNumber)

src/jvm/main/org/jetbrains/kotlinx/lincheck/bootstrap/CodeInformation.kt

@@ -0,0 +1,127 @@
+/*
+ * Lincheck
+ *
+ * Copyright (C) 2019 - 2024 JetBrains s.r.o.
+ *
+ * This Source Code Form is subject to the terms of the
+ * Mozilla Public License, v. 2.0. If a copy of the MPL was not distributed
+ * with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ */
+
+// we need to use some "legal" package for the bootstrap class loader
+@file:Suppress("PackageDirectoryMismatch")
+
+package sun.nio.ch.lincheck
+
+import sun.misc.*
+import java.lang.reflect.Modifier.*
+import java.util.concurrent.atomic.*
+
+/**
+ * [CodeLocations] object is used to maintain the mapping between unique IDs and code locations.
+ * When Lincheck detects an error in the model checking mode, it provides a detailed interleaving trace.
+ * This trace includes a list of all shared memory events that occurred during the execution of the program,
+ * along with their corresponding code locations. To minimize overhead, Lincheck assigns unique IDs to all
+ * code locations it analyses, and stores more detailed information necessary for trace generation in this object.
+ */
+internal object CodeLocations {
+    private val codeLocations = ArrayList<StackTraceElement>()
+
+    /**
+     * Registers a new code location and returns its unique ID.
+     *
+     * @param stackTraceElement Stack trace element representing the new code location.
+     * @return Unique ID of the new code location.
+     */
+    @JvmStatic
+    @Synchronized
+    fun newCodeLocation(stackTraceElement: StackTraceElement): Int {
+        val id = codeLocations.size
+        codeLocations.add(stackTraceElement)
+        return id
+    }
+
+    /**
+     * Returns the [StackTraceElement] associated with the specified code location ID.
+     *
+     * @param codeLocationId ID of the code location.
+     * @return [StackTraceElement] corresponding to the given ID.
+     */
+    @JvmStatic
+    @Synchronized
+    fun stackTrace(codeLocationId: Int): StackTraceElement {
+        return codeLocations[codeLocationId]
+    }
+}
+
+/**
+ * [AtomicFields] is used to map atomic objects to their names (usually, the corresponding field names)
+ * and class owners. The weak identity hash map ensures that atomic objects are compared using reference
+ * equality and does not prevent them from being garbage collected.
+ */
+internal object AtomicFields {
+    private val unsafe: Unsafe = try {
+        val unsafeField = Unsafe::class.java.getDeclaredField("theUnsafe")
+        unsafeField.isAccessible = true
+        unsafeField.get(null) as Unsafe
+    } catch (ex: Exception) {
+        throw RuntimeException(ex)
+    }
+
+    fun getAtomicFieldName(updater: Any): String? {
+        if (updater !is AtomicIntegerFieldUpdater<*> && updater !is AtomicLongFieldUpdater<*> && updater !is AtomicReferenceFieldUpdater<*, *>) {
+            throw IllegalArgumentException("Provided object is not a recognized Atomic*FieldUpdater type.")
+        }
+        // Extract the private offset value and find the matching field.
+        try {
+            // Cannot use neither reflection not MethodHandles.Lookup, as they lead to a warning.
+            val tclassField = updater.javaClass.getDeclaredField("tclass")
+            val targetType = unsafe.getObject(updater, unsafe.objectFieldOffset(tclassField)) as Class<*>
+
+            val offsetField = updater.javaClass.getDeclaredField("offset")
+            val offset = unsafe.getLong(updater, unsafe.objectFieldOffset(offsetField))
+
+            for (field in targetType.declaredFields) {
+                try {
+                    if (isNative(field.modifiers)) continue
+                    val fieldOffset = if (isStatic(field.modifiers)) unsafe.staticFieldOffset(field)
+                    else unsafe.objectFieldOffset(field)
+                    if (fieldOffset == offset) return field.name
+                } catch (t: Throwable) {
+                    t.printStackTrace()
+                }
+            }
+        } catch (t: Throwable) {
+            t.printStackTrace()
+        }
+
+        return null // Field not found
+    }
+}
+
+/**
+ * [FinalFields] object is used to track final fields across different classes.
+ * It is used to maintain a set of all final fields and
+ * provides helper functions to add and check final fields.
+ */
+internal object FinalFields {
+    private val finalFields = HashSet<String>() // className + SEPARATOR + fieldName
+    private const val SEPARATOR = "$^&*-#"
+
+    /**
+     * Adds a new final field to the set.
+     */
+    fun addFinalField(className: String, fieldName: String) {
+        finalFields.add(className + SEPARATOR + fieldName)
+    }
+
+    /**
+     * Checks if the given field of a class is final.
+     *
+     * @param className Name of the class that contains the field.
+     * @param fieldName Name of the field to be checked.
+     * @return `true` if the field is final, `false` otherwise.
+     */
+    fun isFinalField(className: String, fieldName: String) =
+        finalFields.contains(className + SEPARATOR + fieldName)
+}