Static memory snapshot collecting and restoring

bootstrap/src/sun/nio/ch/lincheck/EventTracker.java

@@ -32,6 +32,10 @@ public interface EventTracker {
     void beforeNewObjectCreation(String className);
     void afterNewObjectCreation(Object obj);
 
+    void updateSnapshotOnFieldAccess(Object obj, String className, String fieldName, int codeLocation);
+    void updateSnapshotOnArrayElementAccess(Object array, int index, int codeLocation);
+    void updateSnapshotWithEnergeticTracking(Object[] objs);
+
     boolean beforeReadField(Object obj, String className, String fieldName, int codeLocation,
                             boolean isStatic, boolean isFinal);
     boolean beforeReadArrayElement(Object array, int index, int codeLocation);

bootstrap/src/sun/nio/ch/lincheck/Injections.java

@@ -294,6 +294,27 @@ public static void afterNewObjectCreation(Object obj) {
         getEventTracker().afterNewObjectCreation(obj);
     }
 
+    /**
+     * Called from instrumented code on field access.
+     */
+    public static void updateSnapshotOnFieldAccess(Object obj, String className, String fieldName, int codeLocation) {
+        getEventTracker().updateSnapshotOnFieldAccess(obj, className, fieldName, codeLocation);
+    }
+
+    /**
+     * Called from instrumented code on array element access.
+     */
+    public static void updateSnapshotOnArrayElementAccess(Object array, int index, int codeLocation) {
+        getEventTracker().updateSnapshotOnArrayElementAccess(array, index, codeLocation);
+    }
+
+    /**
+     * Called from instrumented code on constructor invocation, where passed objects are subtypes of the constructor class type.
+     */
+    public static void updateSnapshotWithEnergeticTracking(Object[] objs) {
+        getEventTracker().updateSnapshotWithEnergeticTracking(objs);
+    }
+
     /**
      * Called from the instrumented code to replace [java.lang.Object.hashCode] method call with some
      * deterministic value.

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

@@ -95,6 +95,7 @@ internal fun createFromTestClassAnnotations(testClass: Class<*>): List<CTestConf
                     checkObstructionFreedom = ann.checkObstructionFreedom,
                     hangingDetectionThreshold = ann.hangingDetectionThreshold,
                     invocationsPerIteration = ann.invocationsPerIteration,
+                    restoreStaticMemory = ann.restoreStaticMemory,
                     guarantees = ManagedCTestConfiguration.DEFAULT_GUARANTEES,
                     minimizeFailedScenario = ann.minimizeFailedScenario,
                     sequentialSpecification = chooseSequentialSpecification(

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

@@ -22,7 +22,7 @@ package org.jetbrains.kotlinx.lincheck
 
 import org.jetbrains.kotlinx.lincheck.strategy.managed.ObjectLabelFactory.getObjectNumber
 import org.jetbrains.kotlinx.lincheck.util.*
-import java.lang.reflect.Modifier
+import org.jetbrains.kotlinx.lincheck.util.readFieldViaUnsafe
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.coroutines.Continuation

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

@@ -10,19 +10,23 @@
 package org.jetbrains.kotlinx.lincheck
 
 import kotlinx.coroutines.*
-import sun.nio.ch.lincheck.*
 import org.jetbrains.kotlinx.lincheck.runner.*
 import org.jetbrains.kotlinx.lincheck.strategy.managed.*
 import org.jetbrains.kotlinx.lincheck.transformation.LincheckClassFileTransformer
 import org.jetbrains.kotlinx.lincheck.util.UnsafeHolder
+import org.jetbrains.kotlinx.lincheck.util.isAtomicArray
+import org.jetbrains.kotlinx.lincheck.util.isAtomicFUArray
 import org.jetbrains.kotlinx.lincheck.verifier.*
-import org.jetbrains.kotlinx.lincheck.util.*
+import sun.nio.ch.lincheck.*
 import java.io.PrintWriter
 import java.io.StringWriter
 import java.lang.ref.*
 import java.lang.reflect.*
 import java.lang.reflect.Method
 import java.util.*
+import java.util.concurrent.atomic.AtomicIntegerArray
+import java.util.concurrent.atomic.AtomicLongArray
+import java.util.concurrent.atomic.AtomicReferenceArray
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 
@@ -213,6 +217,96 @@ internal val Throwable.text: String get() {
     return writer.buffer.toString()
 }
 
+/**
+ * Returns all found fields in the hierarchy.
+ * Multiple fields with the same name and the same type may be returned
+ * if they appear in the subclass and a parent class.
+ */
+internal val Class<*>.allDeclaredFieldWithSuperclasses get(): List<Field> {
+    val fields: MutableList<Field> = ArrayList<Field>()
+    var currentClass: Class<*>? = this
+    while (currentClass != null) {
+        val declaredFields: Array<Field> = currentClass.declaredFields
+        fields.addAll(declaredFields)
+        currentClass = currentClass.superclass
+    }
+    return fields
+}
+
+/**
+ * Finds a public/protected/private/internal field in the class and its superclasses/interfaces by name.
+ *
+ * @param fieldName the name of the field to find.
+ * @return the [java.lang.reflect.Field] object if found, or `null` if not found.
+ */
+fun Class<*>.findField(fieldName: String): Field {
+    // Search in the class hierarchy
+    var clazz: Class<*>? = this
+    while (clazz != null) {
+        // Check class itself
+        try {
+            return clazz.getDeclaredField(fieldName)
+        }
+        catch (_: NoSuchFieldException) {}
+
+        // Check interfaces
+        for (interfaceClass in clazz.interfaces) {
+            try {
+                return interfaceClass.getDeclaredField(fieldName)
+            }
+            catch (_: NoSuchFieldException) {}
+        }
+
+        // Move up the hierarchy
+        clazz = clazz.superclass
+    }
+
+    throw NoSuchFieldException("Class '${this.name}' does not have field '$fieldName'")
+}
+
+@Suppress("DEPRECATION")
+fun getFieldOffset(field: Field): Long {
+    return if (Modifier.isStatic(field.modifiers)) {
+        UnsafeHolder.UNSAFE.staticFieldOffset(field)
+    }
+    else {
+        UnsafeHolder.UNSAFE.objectFieldOffset(field)
+    }
+}
+
+internal fun getArrayElementOffset(arr: Any, index: Int): Long {
+    val clazz = arr::class.java
+    val baseOffset = UnsafeHolder.UNSAFE.arrayBaseOffset(clazz).toLong()
+    val indexScale = UnsafeHolder.UNSAFE.arrayIndexScale(clazz).toLong()
+
+    return baseOffset + index * indexScale
+}
+
+internal fun getArrayLength(arr: Any): Int {
+    return when {
+        arr is Array<*>     -> arr.size
+        arr is IntArray     -> arr.size
+        arr is DoubleArray  -> arr.size
+        arr is FloatArray   -> arr.size
+        arr is LongArray    -> arr.size
+        arr is ShortArray   -> arr.size
+        arr is ByteArray    -> arr.size
+        arr is BooleanArray -> arr.size
+        arr is CharArray    -> arr.size
+        isAtomicArray(arr)  -> getAtomicArrayLength(arr)
+        else -> error("Argument is not an array")
+    }
+}
+
+internal fun getAtomicArrayLength(arr: Any): Int {
+    return when {
+        arr is AtomicReferenceArray<*> -> arr.length()
+        arr is AtomicIntegerArray -> arr.length()
+        arr is AtomicLongArray -> arr.length()
+        isAtomicFUArray(arr) -> arr.javaClass.getMethod("getSize").invoke(arr) as Int
+        else -> error("Argument is not atomic array")
+    }
+}
 
 @Suppress("DEPRECATION")
 internal fun findFieldNameByOffset(targetType: Class<*>, offset: Long): String? {

src/jvm/main/org/jetbrains/kotlinx/lincheck/strategy/Strategy.kt

@@ -90,6 +90,11 @@ abstract class Strategy protected constructor(
     override fun close() {
         runner.close()
     }
+
+    /**
+     * Restores recorded values of all memory reachable from static state.
+     */
+    open fun restoreStaticMemorySnapshot() {}
 }
 
 /**
@@ -101,15 +106,20 @@ abstract class Strategy protected constructor(
  * @return the failure, if detected, null otherwise.
  */
 fun Strategy.runIteration(invocations: Int, verifier: Verifier): LincheckFailure? {
+    var failure: LincheckFailure? = null
+
     for (invocation in 0 until invocations) {
         if (!nextInvocation())
-            return null
+            break
         val result = runInvocation()
-        val failure = verify(result, verifier)
+        restoreStaticMemorySnapshot()
+        failure = verify(result, verifier)
         if (failure != null)
-            return failure
+            break
     }
-    return null
+
+    restoreStaticMemorySnapshot()
+    return failure
 }
 
 /**

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

@@ -28,6 +28,7 @@ abstract class ManagedCTestConfiguration(
     val checkObstructionFreedom: Boolean,
     val hangingDetectionThreshold: Int,
     val invocationsPerIteration: Int,
+    val restoreStaticMemory: Boolean,
     val guarantees: List<ManagedStrategyGuarantee>,
     minimizeFailedScenario: Boolean,
     sequentialSpecification: Class<*>,
@@ -52,6 +53,7 @@ abstract class ManagedCTestConfiguration(
         const val DEFAULT_CHECK_OBSTRUCTION_FREEDOM = false
         const val DEFAULT_HANGING_DETECTION_THRESHOLD = 101
         const val LIVELOCK_EVENTS_THRESHOLD = 10001
+        const val DEFAULT_RESTORE_STATIC_MEMORY = true
         val DEFAULT_GUARANTEES = listOf<ManagedStrategyGuarantee>()
     }
 }

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlinx.lincheck.strategy.managed.ManagedCTestConfiguration
 import org.jetbrains.kotlinx.lincheck.strategy.managed.ManagedCTestConfiguration.Companion.DEFAULT_GUARANTEES
 import org.jetbrains.kotlinx.lincheck.strategy.managed.ManagedCTestConfiguration.Companion.DEFAULT_HANGING_DETECTION_THRESHOLD
 import org.jetbrains.kotlinx.lincheck.strategy.managed.ManagedCTestConfiguration.Companion.DEFAULT_INVOCATIONS
+import org.jetbrains.kotlinx.lincheck.strategy.managed.ManagedCTestConfiguration.Companion.DEFAULT_RESTORE_STATIC_MEMORY
 import java.util.*
 
 /**
@@ -23,6 +24,7 @@ abstract class ManagedOptions<OPT : Options<OPT, CTEST>, CTEST : CTestConfigurat
     protected var invocationsPerIteration = DEFAULT_INVOCATIONS
     protected var checkObstructionFreedom = DEFAULT_CHECK_OBSTRUCTION_FREEDOM
     protected var hangingDetectionThreshold = DEFAULT_HANGING_DETECTION_THRESHOLD
+    protected var restoreStaticMemory = DEFAULT_RESTORE_STATIC_MEMORY
     protected val guarantees: MutableList<ManagedStrategyGuarantee> = ArrayList(DEFAULT_GUARANTEES)
 
     /**
@@ -50,6 +52,13 @@ abstract class ManagedOptions<OPT : Options<OPT, CTEST>, CTEST : CTestConfigurat
         this.hangingDetectionThreshold = hangingDetectionThreshold
     }
 
+    /**
+     * Set to `true` to activate the static memory snapshot tracking and restoring algorithm.
+     */
+    fun restoreStaticMemory(restoreStaticMemory: Boolean): OPT = applyAndCast {
+        this.restoreStaticMemory = restoreStaticMemory
+    }
+
     /**
      * Add a guarantee that methods in some classes are either correct in terms of concurrent execution or irrelevant.
      * These guarantees can be used for optimization. For example, we can add a guarantee that all the methods

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

@@ -82,6 +82,9 @@ abstract class ManagedStrategy(
     // Tracker of the thread parking.
     protected abstract val parkingTracker: ParkingTracker
 
+    // Snapshot of the memory, reachable from static fields
+    protected val staticMemorySnapshot = SnapshotTracker()
+
     // InvocationResult that was observed by the strategy during the execution (e.g., a deadlock).
     @Volatile
     protected var suddenInvocationResult: InvocationResult? = null
@@ -216,12 +219,20 @@ abstract class ManagedStrategy(
         parkingTracker.reset()
     }
 
+    override fun restoreStaticMemorySnapshot() {
+        if (testCfg.restoreStaticMemory) {
+            staticMemorySnapshot.restoreValues()
+            super.restoreStaticMemorySnapshot()
+        }
+    }
+
     /**
      * Runs the current invocation.
      */
     override fun runInvocation(): InvocationResult {
         while (true) {
             initializeInvocation()
+            restoreStaticMemorySnapshot()
             val result = runner.run()
             // In case the runner detects a deadlock, some threads can still manipulate the current strategy,
             // so we're not interested in suddenInvocationResult in this case
@@ -909,6 +920,35 @@ abstract class ManagedStrategy(
         }
     }
 
+    /**
+     * Tracks a specific field of an [obj], if the [obj] is either `null` (which means that field is static),
+     * or one this objects which contains it is already stored.
+     */
+    override fun updateSnapshotOnFieldAccess(obj: Any?, className: String, fieldName: String, codeLocation: Int) = runInIgnoredSection {
+        if (testCfg.restoreStaticMemory) {
+            staticMemorySnapshot.trackField(obj, className, fieldName)
+        }
+    }
+
+    /**
+     * Tracks a specific [array] element at [index], if the [array] is already tracked.
+     */
+    override fun updateSnapshotOnArrayElementAccess(array: Any, index: Int, codeLocation: Int) = runInIgnoredSection {
+        if (testCfg.restoreStaticMemory) {
+            staticMemorySnapshot.trackArrayCell(array, index)
+        }
+    }
+
+    /**
+     * Tracks all objects in [objs] energetically.
+     * Required as a trick to overcome issue with leaking this in constructors, see https://github.com/JetBrains/lincheck/issues/424.
+     */
+    override fun updateSnapshotWithEnergeticTracking(objs: Array<Any?>) = runInIgnoredSection {
+        if (testCfg.restoreStaticMemory) {
+            staticMemorySnapshot.trackObjects(objs)
+        }
+    }
+
     private fun methodGuaranteeType(owner: Any?, className: String, methodName: String): ManagedGuaranteeType? = runInIgnoredSection {
         userDefinedGuarantees?.forEach { guarantee ->
             val ownerName = owner?.javaClass?.canonicalName ?: className