Static memory snapshot: add support for Unsafe, VarHandle, AFU, kotlinx.atomicfu

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

@@ -25,6 +25,57 @@ import java.util.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 
+fun shouldTransformClass(className: String): Boolean {
+    // We do not need to instrument most standard Java classes.
+    // It is fine to inject the Lincheck analysis only into the
+    // `java.util.*` ones, ignored the known atomic constructs.
+    if (className.startsWith("java.")) {
+        if (className.startsWith("java.util.concurrent.") && className.contains("Atomic")) return false
+        if (className.startsWith("java.util.")) return true
+        if (className.startsWith("com.sun.")) return false
+        return false
+    }
+    if (className.startsWith("sun.")) return false
+    if (className.startsWith("javax.")) return false
+    if (className.startsWith("jdk.")) return false
+    // We do not need to instrument most standard Kotlin classes.
+    // However, we need to inject the Lincheck analysis into the classes
+    // related to collections, iterators, random and coroutines.
+    if (className.startsWith("kotlin.")) {
+        if (className.startsWith("kotlin.collections.")) return true
+        if (className.startsWith("kotlin.jvm.internal.Array") && className.contains("Iterator")) return true
+        if (className.startsWith("kotlin.ranges.")) return true
+        if (className.startsWith("kotlin.random.")) return true
+        if (className.startsWith("kotlin.coroutines.jvm.internal.")) return false
+        if (className.startsWith("kotlin.coroutines.")) return true
+        return false
+    }
+    if (className.startsWith("kotlinx.atomicfu.")) return false
+    // We need to skip the classes related to the debugger support in Kotlin coroutines.
+    if (className.startsWith("kotlinx.coroutines.debug.")) return false
+    if (className == "kotlinx.coroutines.DebugKt") return false
+    // We should never transform the coverage-related classes.
+    if (className.startsWith("com.intellij.rt.coverage.")) return false
+    // We can also safely do not instrument some libraries for performance reasons.
+    if (className.startsWith("com.esotericsoftware.kryo.")) return false
+    if (className.startsWith("net.bytebuddy.")) return false
+    if (className.startsWith("net.rubygrapefruit.platform.")) return false
+    if (className.startsWith("io.mockk.")) return false
+    if (className.startsWith("it.unimi.dsi.fastutil.")) return false
+    if (className.startsWith("worker.org.gradle.")) return false
+    if (className.startsWith("org.objectweb.asm.")) return false
+    if (className.startsWith("org.gradle.")) return false
+    if (className.startsWith("org.slf4j.")) return false
+    if (className.startsWith("org.apache.commons.lang.")) return false
+    if (className.startsWith("org.junit.")) return false
+    if (className.startsWith("junit.framework.")) return false
+    // Finally, we should never instrument the Lincheck classes.
+    if (className.startsWith("org.jetbrains.kotlinx.lincheck.")) return false
+    if (className.startsWith("sun.nio.ch.lincheck.")) return false
+    // All the classes that were not filtered out are eligible for transformation.
+    return true
+}
+
 fun <T> List<T>.isSuffixOf(list: List<T>): Boolean {
     if (size > list.size) return false
     for (i in indices) {

src/jvm/main/org/jetbrains/kotlinx/lincheck/strategy/managed/AtomicFieldUpdaterNames.kt

@@ -24,7 +24,7 @@ import java.util.concurrent.atomic.AtomicReferenceFieldUpdater
 internal object AtomicFieldUpdaterNames {
 
     @Suppress("DEPRECATION")
-    internal fun getAtomicFieldUpdaterName(updater: Any): String? {
+    internal fun getAtomicFieldUpdaterName(updater: Any): AtomicFieldUpdaterDescriptor? {
         if (updater !is AtomicIntegerFieldUpdater<*> && updater !is AtomicLongFieldUpdater<*> && updater !is AtomicReferenceFieldUpdater<*, *>) {
             throw IllegalArgumentException("Provided object is not a recognized Atomic*FieldUpdater type.")
         }
@@ -37,11 +37,22 @@ internal object AtomicFieldUpdaterNames {
             val offsetField = updater.javaClass.getDeclaredField("offset")
             val offset = UNSAFE.getLong(updater, UNSAFE.objectFieldOffset(offsetField))
 
-            return findFieldNameByOffsetViaUnsafe(targetType, offset)
+            return findFieldNameByOffsetViaUnsafe(targetType, offset).let { fieldName ->
+                if (fieldName != null) AtomicFieldUpdaterDescriptor(targetType, fieldName)
+                else null
+            }
         } catch (t: Throwable) {
             t.printStackTrace()
         }
 
         return null // Field not found
     }
-}
+}
+
+/**
+ * Descriptor of the Java AFU instance.
+ */
+internal data class AtomicFieldUpdaterDescriptor(
+    val targetType: Class<*>,
+    val fieldName: String
+)

src/jvm/main/org/jetbrains/kotlinx/lincheck/strategy/managed/AtomicReferenceNames.kt

@@ -22,7 +22,7 @@ import java.util.concurrent.atomic.AtomicReferenceArray
 
 /**
  * Provides method call type to create a more convenient trace point
- * with a owner of this AtomicReference field and a name if it can be found.
+ * with an owner of this AtomicReference field and a name if it can be found.
  * Recursively scans the test object, trying to find the provided AtomicReference
  * instance as a field. If two or more fields contain this AtomicReference field, then we
  * fall back to the default behavior.

src/jvm/main/org/jetbrains/kotlinx/lincheck/strategy/managed/ManagedStrategy.kt

@@ -950,6 +950,63 @@ abstract class ManagedStrategy(
         staticMemorySnapshot.trackObjects(objs)
     }
 
+
+
+    /**
+     * Tracks fields that are accessed via System.arraycopy, Unsafe API, VarHandle API, Java AFU API, and kotlinx.atomicfu.
+     */
+    private fun processMethodEffectOnStaticSnapshot(owner: Any?, params: Array<Any?>) {
+        when {
+            // Unsafe API
+            isUnsafe(owner) -> {
+                val methodType: UnsafeName = UnsafeNames.getMethodCallType(params)
+                when (methodType) {
+                    is UnsafeInstanceMethod -> {
+                        staticMemorySnapshot.trackField(methodType.owner, methodType.owner.javaClass, methodType.fieldName)
+                    }
+                    is UnsafeStaticMethod -> {
+                        staticMemorySnapshot.trackField(null, methodType.clazz, methodType.fieldName)
+                    }
+                    is UnsafeArrayMethod -> {
+                        staticMemorySnapshot.trackArrayCell(methodType.array, methodType.index)
+                    }
+                    else -> {}
+                }
+            }
+            // VarHandle API
+            isVarHandle(owner) &&
+            (
+                VarHandleNames.isInstanceVarHandle(owner) && staticMemorySnapshot.isTracked(params.firstOrNull()) ||
+                VarHandleNames.isArrayVarHandle(owner) && staticMemorySnapshot.isTracked(params.firstOrNull()) ||
+                VarHandleNames.isStaticVarHandle(owner)
+            ) -> {
+                val methodType: VarHandleMethodType = VarHandleNames.varHandleMethodType(owner, params)
+                when (methodType) {
+                    is InstanceVarHandleMethod -> {
+                        staticMemorySnapshot.trackField(methodType.owner, methodType.owner.javaClass, methodType.fieldName)
+                    }
+                    is StaticVarHandleMethod -> {
+                        staticMemorySnapshot.trackField(null, methodType.ownerClass, methodType.fieldName)
+                    }
+                    is ArrayVarHandleMethod -> {
+                        staticMemorySnapshot.trackArrayCell(methodType.array, methodType.index)
+                    }
+                    else -> {}
+                }
+            }
+            // Java AFU (this also automatically handles the `kotlinx.atomicfu`, since they are compiled to Java AFU + Java atomic arrays)
+            isAtomicFieldUpdater(owner) -> {
+                val obj = params[0]
+                val afuDesc: AtomicFieldUpdaterDescriptor? = AtomicFieldUpdaterNames.getAtomicFieldUpdaterName(owner!!)
+                check(afuDesc != null) { "Cannot extract field name referenced by Java AFU object $owner" }
+
+                staticMemorySnapshot.trackField(obj, afuDesc.targetType, afuDesc.fieldName)
+            }
+            // TODO: System.arraycopy
+            // TODO: reflexivity
+        }
+    }
+
     private fun methodGuaranteeType(owner: Any?, className: String, methodName: String): ManagedGuaranteeType? = runInIgnoredSection {
         userDefinedGuarantees?.forEach { guarantee ->
             val ownerName = owner?.javaClass?.canonicalName ?: className
@@ -969,9 +1026,15 @@ abstract class ManagedStrategy(
         methodId: Int,
         params: Array<Any?>
     ) {
-        val guarantee = runInIgnoredSection {
+        var guarantee: ManagedGuaranteeType? = null
+
+        runInIgnoredSection {
+            // process method effect on the static memory snapshot
+            processMethodEffectOnStaticSnapshot(owner, params)
+
+            // process known method concurrency guarantee
             val atomicMethodDescriptor = getAtomicMethodDescriptor(owner, methodName)
-            val guarantee = when {
+            guarantee = when {
                 (atomicMethodDescriptor != null) -> ManagedGuaranteeType.TREAT_AS_ATOMIC
                 else -> methodGuaranteeType(owner, className, methodName)
             }
@@ -988,8 +1051,8 @@ abstract class ManagedStrategy(
             if (guarantee == null) {
                 loopDetector.beforeMethodCall(codeLocation, params)
             }
-            guarantee
         }
+
         if (guarantee == ManagedGuaranteeType.IGNORE ||
             guarantee == ManagedGuaranteeType.TREAT_AS_ATOMIC) {
             // It's important that this method can't be called inside runInIgnoredSection, as the ignored section
@@ -1346,7 +1409,7 @@ abstract class ManagedStrategy(
         atomicUpdater: Any,
         parameters: Array<Any?>,
     ): MethodCallTracePoint {
-        getAtomicFieldUpdaterName(atomicUpdater)?.let { tracePoint.initializeOwnerName(it) }
+        getAtomicFieldUpdaterName(atomicUpdater)?.let { tracePoint.initializeOwnerName(it.fieldName) }
         tracePoint.initializeParameters(parameters.drop(1).map { adornedStringRepresentation(it) })
         return tracePoint
     }

src/jvm/main/org/jetbrains/kotlinx/lincheck/strategy/managed/SnapshotTracker.kt

@@ -46,10 +46,27 @@ class SnapshotTracker {
         class ArrayCellNode(val index: Int, initialValue: Any?) : MemoryNode(initialValue)
     }
 
-    fun trackField(obj: Any?, className: String, fieldName: String) {
+    fun isTracked(obj: Any?): Boolean = obj != null && obj in trackedObjects
+
+    fun trackField(obj: Any?, accessClass: Class<*>, fieldName: String) {
         if (obj != null && obj !in trackedObjects) return
+        trackFieldImpl(
+            obj = obj,
+            clazz = getDeclaringClass(obj, accessClass, fieldName),
+            fieldName = fieldName
+        )
+    }
+
+    fun trackField(obj: Any?, accessClassName: String, fieldName: String) {
+        if (obj != null && obj !in trackedObjects) return
+        trackFieldImpl(
+            obj = obj,
+            clazz = getDeclaringClass(obj, Class.forName(accessClassName), fieldName),
+            fieldName = fieldName
+        )
+    }
 
-        val clazz = getDeclaringClass(obj, className, fieldName)
+    private fun trackFieldImpl(obj: Any?, clazz: Class<*>, fieldName: String) {
         val field = clazz.findField(fieldName)
         val readResult = readFieldSafely(obj, field)
 
@@ -179,9 +196,8 @@ class SnapshotTracker {
     private fun shouldTrackEagerly(obj: Any?): Boolean {
         if (obj == null) return false
         return (
-            // TODO: in further development of snapshot restoring feature this check should be removed
-            //  (and only check for java atomic classes should be inserted), see https://github.com/JetBrains/lincheck/pull/418#issue-2595977113
-            //  right now it is needed for collections to be restored properly (because of missing support for `System.arrayCopy()` and other similar methods)
+            // TODO: after support for System.arraycopy,
+            //  rewrite to `obj.javaClass.name.startsWith("java.util.concurrent.") && obj.javaClass.name.contains("Atomic")`
             obj.javaClass.name.startsWith("java.util.")
         )
     }
@@ -237,8 +253,8 @@ class SnapshotTracker {
         )
     }
 
-    private fun getDeclaringClass(obj: Any?, className: String, fieldName: String): Class<*> {
-        return Class.forName(className).let {
+    private fun getDeclaringClass(obj: Any?, clazz: Class<*>, fieldName: String): Class<*> {
+        return clazz.let {
             if (obj != null) it
             else it.findField(fieldName).declaringClass
         }