Add cancellation support

Fixes #17

lincheck/pom.xml

@@ -31,7 +31,7 @@
         <dependency>
             <groupId>org.jetbrains.kotlinx</groupId>
             <artifactId>kotlinx-coroutines-core</artifactId>
-            <version>1.2.2</version>
+            <version>1.3.1</version>
         </dependency>
         <dependency>
             <groupId>org.jetbrains.kotlin</groupId>

lincheck/src/main/java/org/jetbrains/kotlinx/lincheck/Actor.kt

@@ -31,12 +31,14 @@ import kotlin.coroutines.Continuation
  *
  * @see Operation
  */
-data class Actor(
+data class Actor @JvmOverloads constructor(
     val method: Method,
     val arguments: List<Any?>,
-    val handledExceptions: List<Class<out Throwable>>
+    val handledExceptions: List<Class<out Throwable>>,
+    val cancelOnSuspension: Boolean = false
 ) {
-    override fun toString() = method.name + arguments.joinToString(prefix = "(", postfix = ")", separator = ",") { it.toString() }
+    override fun toString() = cancellableMark + method.name + arguments.joinToString(prefix = "(", postfix = ")", separator = ",") { it.toString() }
+    private val cancellableMark get() = (if (cancelOnSuspension) "*" else "")
 
     val handlesExceptions = handledExceptions.isNotEmpty()
 

lincheck/src/main/java/org/jetbrains/kotlinx/lincheck/CTestStructure.java

@@ -104,7 +104,7 @@ private static void readTestStructureFromClass(Class<?> clazz, Map<String, Param
                 }
                 // Get list of handled exceptions if they are presented
                 List<Class<? extends Throwable>> handledExceptions = Arrays.asList(operationAnn.handleExceptionsAsResult());
-                ActorGenerator actorGenerator = new ActorGenerator(m, gens, handledExceptions, operationAnn.runOnce());
+                ActorGenerator actorGenerator = new ActorGenerator(m, gens, handledExceptions, operationAnn.runOnce(), operationAnn.cancellableOnSuspension());
                 actorGenerators.add(actorGenerator);
                 // Get list of groups and add this operation to specified ones
                 String opGroup = operationAnn.group();

lincheck/src/main/java/org/jetbrains/kotlinx/lincheck/CancellabilitySupportTransformer.kt

@@ -0,0 +1,60 @@
+/*-
+ * #%L
+ * Lincheck
+ * %%
+ * Copyright (C) 2019 JetBrains s.r.o.
+ * %%
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ * 
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Lesser Public License for more details.
+ * 
+ * You should have received a copy of the GNU General Lesser Public
+ * License along with this program.  If not, see
+ * <http://www.gnu.org/licenses/lgpl-3.0.html>.
+ * #L%
+ */
+package org.jetbrains.kotlinx.lincheck
+
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import org.objectweb.asm.Type
+import org.objectweb.asm.commons.AdviceAdapter
+import org.objectweb.asm.commons.Method
+import kotlin.reflect.jvm.javaMethod
+
+internal class CancellabilitySupportClassTransformer(cv: ClassVisitor) : ClassVisitor(Opcodes.ASM5, cv) {
+    private var className: String? = null
+
+    override fun visit(version: Int, access: Int, name: String?, signature: String?, superName: String?, interfaces: Array<String>?) {
+        className = name
+        super.visit(version, access, name, signature, superName, interfaces)
+    }
+
+    override fun visitMethod(access: Int, methodName: String?, desc: String?, signature: String?, exceptions: Array<String>?): MethodVisitor {
+        var mv = super.visitMethod(access, methodName, desc, signature, exceptions)
+        val transform = "kotlinx/coroutines/CancellableContinuationImpl" == className && "getResult" == methodName
+        if (transform) {
+            mv = CancellabilitySupportMethodTransformer(access, methodName, desc, mv)
+        }
+        return mv
+    }
+}
+
+private class CancellabilitySupportMethodTransformer(access: Int, methodName: String?, desc: String?, mv: MethodVisitor)
+    : AdviceAdapter(Opcodes.ASM5, mv, access, methodName, desc)
+{
+    override fun onMethodEnter() {
+        this.loadThis()
+        this.invokeStatic(storeCancellableContOwnerType, storeCancellableContMethod)
+    }
+}
+
+private val storeCancellableContMethod = Method.getMethod(::storeCancellableContinuation.javaMethod)
+private val storeCancellableContOwnerType = Type.getType(::storeCancellableContinuation.javaMethod!!.declaringClass)

lincheck/src/main/java/org/jetbrains/kotlinx/lincheck/Result.kt

@@ -1,6 +1,6 @@
 package org.jetbrains.kotlinx.lincheck
 
-import kotlin.coroutines.Continuation
+import kotlin.coroutines.*
 
 /*
  * #%L
@@ -38,7 +38,7 @@ import kotlin.coroutines.Continuation
  */
 sealed class Result {
     abstract val wasSuspended: Boolean
-    protected val wasSuspendedPrefix: String get() = (if (wasSuspended) "S + " else "")
+    protected val wasSuspendedPrefix: String get() = (if (wasSuspended) "SUSPENDED + " else "")
 }
 
 /**
@@ -63,12 +63,27 @@ object SuspendedVoidResult : Result() {
 
 private const val VOID = "void"
 
+object Cancelled : Result() {
+    override val wasSuspended get() = true
+    override fun toString() = wasSuspendedPrefix + "CANCELLED"
+}
+
 /**
  * Type of result used if the actor invocation fails with the specified in {@link Operation#handleExceptionsAsResult()} exception [tClazz].
  */
-data class ExceptionResult @JvmOverloads constructor(val tClazz: Class<out Throwable>?, override val wasSuspended: Boolean = false) : Result() {
-    override fun toString() = wasSuspendedPrefix + "${tClazz?.simpleName}"
+@Suppress("DataClassPrivateConstructor")
+data class ExceptionResult private constructor(val tClazz: Class<out Throwable>, override val wasSuspended: Boolean) : Result() {
+    override fun toString() = wasSuspendedPrefix + tClazz.simpleName
+
+    companion object {
+        @Suppress("UNCHECKED_CAST")
+        @JvmOverloads
+        fun create(tClazz: Class<out Throwable>, wasSuspended: Boolean = false) = ExceptionResult(tClazz.normalize(), wasSuspended)
+    }
 }
+// for byte-code generation
+@JvmSynthetic
+fun createExceptionResult(tClazz: Class<out Throwable>) = ExceptionResult.create(tClazz, false)
 
 /**
  * Type of result used if the actor invocation suspended the thread and did not get the final result yet

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

@@ -31,23 +31,32 @@
 
 import java.io.IOException;
 import java.net.URL;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.ArrayList;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 
 /**
  * This transformer applies required for {@link Strategy} and {@link Runner}
  * class transformations and hines them from others.
  */
 public class TransformationClassLoader extends ExecutionClassLoader {
-    // Strategy and runner provide class transformers
-    private final Strategy strategy;
-    private final Runner runner;
+    private final List<Function<ClassVisitor, ClassVisitor>> classTransformers;
     // Cache for classloading and frames computing during the transformation
     private final Map<String, Class<?>> cache = new ConcurrentHashMap<>();
 
     public TransformationClassLoader(Strategy strategy, Runner runner) {
-        this.strategy = strategy;
-        this.runner = runner;
+        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);
+    }
+
+    public TransformationClassLoader(Function<ClassVisitor, ClassVisitor> classTransformer) {
+        this.classTransformers = Collections.singletonList(classTransformer);
     }
 
     /**
@@ -60,10 +69,12 @@ private static boolean doNotTransform(String className) {
         return className == null ||
             (className.startsWith("org.jetbrains.kotlinx.lincheck.") &&
                 !className.startsWith("org.jetbrains.kotlinx.lincheck.test.") &&
-                !className.equals("ManagedStrategyHolder")) ||
+                !className.endsWith("ManagedStrategyHolder")) ||
             className.startsWith("sun.") ||
             className.startsWith("java.") ||
-            className.startsWith("jdk.internal.");
+            className.startsWith("jdk.internal.") ||
+            className.startsWith("kotlin.") ||
+            (className.equals("kotlinx.coroutines.CancellableContinuation"));
         // TODO let's transform java.util.concurrent
     }
 
@@ -106,12 +117,8 @@ private byte[] instrument(String className) throws IOException {
         // then the runner's one.
         ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
         ClassVisitor cv = new CheckClassAdapter(cw, false); // for debug
-        if (runner.needsTransformation()) {
-            cv = runner.createTransformer(cv);
-        }
-        if (strategy.needsTransformation()) {
-            cv = strategy.createTransformer(cv);
-        }
+        for (Function<ClassVisitor, ClassVisitor> ct : classTransformers)
+            cv = ct.apply(cv);
         // Get transformed bytecode
         cr.accept(cv, ClassReader.EXPAND_FRAMES);
         return cw.toByteArray();

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

@@ -21,15 +21,16 @@
  */
 package org.jetbrains.kotlinx.lincheck
 
-import org.jetbrains.kotlinx.lincheck.execution.ExecutionResult
-import org.jetbrains.kotlinx.lincheck.execution.ExecutionScenario
-import org.jetbrains.kotlinx.lincheck.verifier.DummySequentialSpecification
+import kotlinx.coroutines.*
+import org.jetbrains.kotlinx.lincheck.CancellableContinuationHolder.storedLastCancellableCont
+import org.jetbrains.kotlinx.lincheck.execution.*
+import org.jetbrains.kotlinx.lincheck.runner.*
+import org.jetbrains.kotlinx.lincheck.verifier.*
 import java.lang.ref.WeakReference
-import java.lang.reflect.InvocationTargetException
-import java.lang.reflect.Method
+import java.lang.reflect.*
 import java.util.*
-import kotlin.coroutines.Continuation
-import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
 
 
 @Volatile
@@ -81,12 +82,12 @@ internal fun executeActor(
         val res = m.invoke(specification, *args.toTypedArray())
         return if (m.returnType.isAssignableFrom(Void.TYPE)) VoidResult else createLinCheckResult(res)
     } catch (invE: Throwable) {
-        val eClass = (invE.cause ?: invE)::class.java
+        val eClass = (invE.cause ?: invE).javaClass.normalize()
         for (ec in actor.handledExceptions) {
             if (ec.isAssignableFrom(eClass))
-                return ExceptionResult(eClass)
+                return ExceptionResult.create(eClass)
         }
-        throw IllegalStateException("Invalid exception as a result", invE)
+        throw IllegalStateException("Invalid exception as a result of $actor", invE)
     } catch (e: Exception) {
         e.catch(
             NoSuchMethodException::class.java,
@@ -97,6 +98,8 @@ internal fun executeActor(
     }
 }
 
+internal fun <T> Class<T>.normalize() = LinChecker::class.java.classLoader.loadClass(name) as Class<T>
+
 private val methodsCache = WeakHashMap<Class<*>, WeakHashMap<Method, WeakReference<Method>>>()
 
 private fun getMethod(instance: Any, actor: Actor): Method {
@@ -123,7 +126,7 @@ private fun getMethod(instance: Any, actor: Actor): Method {
  */
 internal fun createLinCheckResult(res: Any?, wasSuspended: Boolean = false) = when {
     (res != null && res.javaClass.isAssignableFrom(Void.TYPE)) || res is Unit -> if (wasSuspended) SuspendedVoidResult else VoidResult
-    res != null && res is Throwable -> ExceptionResult(res.javaClass, wasSuspended)
+    res != null && res is Throwable -> ExceptionResult.create(res.javaClass, wasSuspended)
     res === COROUTINE_SUSPENDED -> Suspended
     res is kotlin.Result<Any?> -> res.toLinCheckResult(wasSuspended)
     else -> ValueResult(res, wasSuspended)
@@ -137,10 +140,10 @@ private fun kotlin.Result<Any?>.toLinCheckResult(wasSuspended: Boolean) =
             is Throwable -> ValueResult(value::class.java, wasSuspended)
             else -> ValueResult(value, wasSuspended)
         }
-    } else ExceptionResult(exceptionOrNull()?.let { it::class.java }, wasSuspended)
+    } else ExceptionResult.create(exceptionOrNull()!!.let { it::class.java }, wasSuspended)
 
 
-fun <R> Throwable.catch(vararg exceptions: Class<*>, block: () -> R): R {
+inline fun <R> Throwable.catch(vararg exceptions: Class<*>, block: () -> R): R {
     if (exceptions.any { this::class.java.isAssignableFrom(it) }) {
         return block()
     } else throw this
@@ -171,3 +174,21 @@ internal operator fun ExecutionResult.set(threadId: Int, actorId: Int, value: Re
     parallelResults.size + 1 -> postResults[actorId] = value
     else -> parallelResults[threadId - 1][actorId] = value
 }
+
+/**
+ * Returns `true` if the continuation was cancelled by [CancellableContinuation.cancel].
+ */
+fun <T> kotlin.Result<T>.cancelled() = isFailure && exceptionOrNull() is CancellationException
+
+object CancellableContinuationHolder {
+    var storedLastCancellableCont: CancellableContinuation<*>? = null
+}
+
+fun storeCancellableContinuation(cont: CancellableContinuation<*>) {
+    val t = Thread.currentThread()
+    if (t is ParallelThreadsRunner.TestThread) {
+        t.cont = cont
+    } else {
+        storedLastCancellableCont = cont
+    }
+}

lincheck/src/main/java/org/jetbrains/kotlinx/lincheck/annotations/Operation.java

@@ -22,10 +22,8 @@
  * #L%
  */
 
-import java.lang.annotation.ElementType;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
+import kotlinx.coroutines.*;
+import java.lang.annotation.*;
 
 /**
  * Mark your method with this annotation in order
@@ -56,4 +54,10 @@
      * Handle the specified exceptions as a result of this operation invocation.
      */
     Class<? extends Throwable>[] handleExceptionsAsResult() default {};
+
+    /**
+     * Specified whether the operation can be cancelled if it suspends,
+     * see {@link CancellableContinuation#cancel}; {@code true} by default.
+     */
+    boolean cancellableOnSuspension() default true;
 }

lincheck/src/main/java/org/jetbrains/kotlinx/lincheck/execution/ActorGenerator.java

@@ -22,6 +22,7 @@
  * #L%
  */
 
+import kotlin.random.Random;
 import org.jetbrains.kotlinx.lincheck.Actor;
 
 import org.jetbrains.kotlinx.lincheck.ActorKt;
@@ -40,19 +41,21 @@ public class ActorGenerator {
     private final List<ParameterGenerator<?>> parameterGenerators;
     private final List<Class<? extends Throwable>> handledExceptions;
     private final boolean useOnce;
+    private final boolean cancellableOnSuspension;
 
     public ActorGenerator(Method method, List<ParameterGenerator<?>> parameterGenerators,
-        List<Class<? extends Throwable>> handledExceptions, boolean useOnce)
+        List<Class<? extends Throwable>> handledExceptions, boolean useOnce, boolean cancellableOnSuspension)
     {
         this.method = method;
         this.parameterGenerators = parameterGenerators;
         this.handledExceptions = handledExceptions;
         this.useOnce = useOnce;
+        this.cancellableOnSuspension = cancellableOnSuspension && isSuspendable();
     }
 
     public Actor generate() {
         List<Object> parameters = parameterGenerators.stream().map(ParameterGenerator::generate).collect(Collectors.toList());
-        return new Actor(method, parameters, handledExceptions);
+        return new Actor(method, parameters, handledExceptions, cancellableOnSuspension & Random.Default.nextBoolean());
     }
 
     public boolean useOnce() {