update.

core/src/main/kotlin/cmu/pasta/fray/core/GlobalContext.kt

@@ -1,5 +1,6 @@
 package cmu.pasta.fray.core
 
+import cmu.pasta.fray.core.command.Configuration
 import cmu.pasta.fray.core.concurrency.HelperThread
 import cmu.pasta.fray.core.concurrency.SynchronizationManager
 import cmu.pasta.fray.core.concurrency.locks.CountDownLatchManager
@@ -65,7 +66,7 @@ object GlobalContext {
   }
 
   fun reportError(e: Throwable) {
-    if (!bugFound && !config!!.ignoreUnhandledExceptions) {
+    if (!bugFound && !config!!.executionInfo.ignoreUnhandledExceptions) {
       bugFound = true
       val sw = StringWriter()
       sw.append("Error found: ${e}\n")
@@ -616,7 +617,8 @@ object GlobalContext {
   }
 
   fun fieldOperation(obj: Any?, owner: String, name: String, type: MemoryOpType) {
-    if (!config!!.interleaveMemoryOps && !volatileManager.isVolatile(owner, name)) return
+    if (!config!!.executionInfo.interleaveMemoryOps && !volatileManager.isVolatile(owner, name))
+        return
     val objIds = mutableListOf<Int>()
     if (obj != null) {
       objIds.add(System.identityHashCode(obj))
@@ -633,7 +635,7 @@ object GlobalContext {
   }
 
   fun arrayOperation(obj: Any, index: Int, type: MemoryOpType) {
-    if (!config!!.interleaveMemoryOps) return
+    if (!config!!.executionInfo.interleaveMemoryOps) return
     val objId = System.identityHashCode(obj)
     memoryOperation((31 * objId) + index, type)
   }
@@ -769,7 +771,7 @@ object GlobalContext {
     }
 
     step += 1
-    if (step > config!!.maxScheduledStep &&
+    if (step > config!!.executionInfo.maxScheduledStep &&
         !currentThread.isExiting &&
         Thread.currentThread() !is HelperThread &&
         !(mainExiting && currentThreadId == mainThreadId)) {

core/src/main/kotlin/cmu/pasta/fray/core/Main.kt

@@ -1,5 +1,7 @@
 package cmu.pasta.fray.core
 
+import cmu.pasta.fray.core.command.MainCommand
+
 fun main(args: Array<String>) {
   val config = MainCommand().apply { main(args) }.toConfiguration()
   val runner = TestRunner(config)

core/src/main/kotlin/cmu/pasta/fray/core/RuntimeDelegate.kt

@@ -545,6 +545,16 @@ class RuntimeDelegate : Delegate() {
     }
   }
 
+  override fun onLatchAwaitTimeout(latch: CountDownLatch, timeout: Long, unit: TimeUnit): Boolean {
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+      return false
+    } else {
+      latch.await()
+      return true
+    }
+  }
+
   override fun onLatchAwaitDone(latch: CountDownLatch) {
     onSkipMethodDone("Latch.await")
     if (checkEntered()) return
@@ -608,33 +618,64 @@ class RuntimeDelegate : Delegate() {
   }
 
   override fun onThreadParkNanos(nanos: Long) {
-    LockSupport.park()
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+    } else {
+      LockSupport.park()
+    }
   }
 
   override fun onThreadParkUntil(nanos: Long) {
-    LockSupport.park()
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+    } else {
+      LockSupport.park()
+    }
   }
 
   override fun onThreadParkNanosWithBlocker(blocker: Any?, nanos: Long) {
-    LockSupport.park(blocker)
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+    } else {
+      LockSupport.park(blocker)
+    }
   }
 
   override fun onThreadParkUntilWithBlocker(blocker: Any?, nanos: Long) {
-    LockSupport.park(blocker)
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+    } else {
+      LockSupport.park(blocker)
+    }
   }
 
   override fun onConditionAwaitTime(o: Condition, time: Long, unit: TimeUnit): Boolean {
-    o.await()
-    return true
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+      return false
+    } else {
+      o.await()
+      return true
+    }
   }
 
   override fun onConditionAwaitNanos(o: Condition, nanos: Long): Long {
-    o.await()
-    return 0
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+      return 0
+    } else {
+      o.await()
+      return nanos
+    }
   }
 
   override fun onConditionAwaitUntil(o: Condition, deadline: Date): Boolean {
-    o.await()
-    return true
+    if (GlobalContext.config!!.executionInfo.timedOpAsYield) {
+      onYield()
+      return false
+    } else {
+      o.await()
+      return true
+    }
   }
 }

core/src/main/kotlin/cmu/pasta/fray/core/TestRunner.kt

@@ -1,5 +1,6 @@
 package cmu.pasta.fray.core
 
+import cmu.pasta.fray.core.command.Configuration
 import cmu.pasta.fray.core.logger.ConsoleLogger
 import cmu.pasta.fray.runtime.Runtime
 import java.nio.file.Paths
@@ -35,7 +36,7 @@ class TestRunner(val config: Configuration) {
         try {
           Runtime.DELEGATE = RuntimeDelegate()
           Runtime.start()
-          config.exec()
+          config.executionInfo.executor.execute()
           Runtime.onMainExit()
         } catch (e: Throwable) {
           Runtime.onReportError(e)

core/src/main/kotlin/cmu/pasta/fray/core/Configuration.kt → core/src/main/kotlin/cmu/pasta/fray/core/command/Configuration.kt

@@ -1,4 +1,4 @@
-package cmu.pasta.fray.core
+package cmu.pasta.fray.core.command
 
 import cmu.pasta.fray.core.logger.CsvLogger
 import cmu.pasta.fray.core.logger.JsonLogger
@@ -9,26 +9,29 @@ import com.github.ajalt.clikt.parameters.groups.OptionGroup
 import com.github.ajalt.clikt.parameters.groups.defaultByName
 import com.github.ajalt.clikt.parameters.groups.groupChoice
 import com.github.ajalt.clikt.parameters.options.*
-import com.github.ajalt.clikt.parameters.types.boolean
 import com.github.ajalt.clikt.parameters.types.file
 import com.github.ajalt.clikt.parameters.types.int
-import java.net.URI
-import java.net.URLClassLoader
 import java.util.*
+import kotlinx.serialization.Polymorphic
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.Json
+import kotlinx.serialization.modules.SerializersModule
+import kotlinx.serialization.modules.polymorphic
+import kotlinx.serialization.modules.subclass
 
 @Serializable
 data class ExecutionInfo(
-    val clazz: String,
-    val method: String,
-    val args: List<String>,
-    val classpaths: List<String>
-) {}
+    @Polymorphic val executor: Executor,
+    val ignoreUnhandledExceptions: Boolean,
+    val timedOpAsYield: Boolean,
+    val interleaveMemoryOps: Boolean,
+    val maxScheduledStep: Int,
+)
 
 sealed class ExecutionConfig(name: String) : OptionGroup(name) {
   open fun getExecutionInfo(): ExecutionInfo {
-    return ExecutionInfo("", "", emptyList(), emptyList())
+    return ExecutionInfo(
+        MethodExecutor("", "", emptyList(), emptyList()), false, false, false, 10000000)
   }
 }
 
@@ -43,17 +46,33 @@ class CliExecutionConfig : ExecutionConfig("cli") {
       option("-cp", "--classpath", help = "Arguments passed to target application")
           .split(":")
           .default(emptyList())
+  val timedOpAsYield by option("-t", "--timed-op-as-yield").flag()
+  val ignoreUnhandledExceptions by option("-e", "--ignore-unhandled-exceptions").flag()
+  val interleaveMemoryOps by option("-m", "--memory").flag()
+  val maxScheduledStep by option("-s", "--max-scheduled-step").int().default(10000)
 
   override fun getExecutionInfo(): ExecutionInfo {
-    return ExecutionInfo(clazz, method, targetArgs, classpaths)
+    return ExecutionInfo(
+        MethodExecutor(clazz, method, targetArgs, classpaths),
+        ignoreUnhandledExceptions,
+        timedOpAsYield,
+        interleaveMemoryOps,
+        maxScheduledStep)
   }
 }
 
 class JsonExecutionConfig : ExecutionConfig("json") {
   val path by option("--config-path").file().required()
 
   override fun getExecutionInfo(): ExecutionInfo {
-    return Json.decodeFromString<ExecutionInfo>(path.readText())
+    val module = SerializersModule {
+      polymorphic(Executor::class) {
+        subclass(MethodExecutor::class)
+        defaultDeserializer { MethodExecutor.serializer() }
+      }
+    }
+    val json = Json { serializersModule = module }
+    return json.decodeFromString<ExecutionInfo>(path.readText())
   }
 }
 
@@ -118,13 +137,9 @@ class PCT : ScheduleAlgorithm("pct") {
 class MainCommand : CliktCommand() {
   val report by option("-o").default("/tmp/report")
   val iter by option("-i", "--iter", help = "Number of iterations").int().default(1)
-  val fullSchedule by option("-f", "--full").boolean().default(false)
+  val fullSchedule by option("-f", "--full").flag()
   val logger by
       option("-l", "--logger").groupChoice("json" to JsonLoggerOption(), "csv" to CsvLoggerOption())
-  val interleaveMemoryOps by option("-m", "--memory").boolean().default(false)
-  val maxScheduledStep by option("-s", "--max-scheduled-step").int().default(10000)
-  val ignoreUnhandledExceptions by
-      option("-e", "--ignore-unhandled-exceptions").boolean().default(false)
   val scheduler by
       option()
           .groupChoice(
@@ -145,48 +160,21 @@ class MainCommand : CliktCommand() {
 
   fun toConfiguration(): Configuration {
     val executionInfo = runConfig.getExecutionInfo()
-    val exec = {
-      val classLoader =
-          URLClassLoader(
-              executionInfo.classpaths.map { it -> URI("file://$it").toURL() }.toTypedArray(),
-              Thread.currentThread().contextClassLoader)
-      Thread.currentThread().contextClassLoader = classLoader
-      val clazz = Class.forName(executionInfo.clazz, true, classLoader)
-      if (executionInfo.args.isEmpty() && executionInfo.method != "main") {
-        val m = clazz.getMethod(executionInfo.method)
-        if (m.modifiers and java.lang.reflect.Modifier.STATIC == 0) {
-          val obj = clazz.getConstructor().newInstance()
-          m.invoke(obj)
-        } else {
-          m.invoke(null)
-        }
-      } else {
-        val m = clazz.getMethod(executionInfo.method, Array<String>::class.java)
-        m.invoke(null, executionInfo.args.toTypedArray())
-      }
-      Unit
-    }
     return Configuration(
-        exec,
+        executionInfo,
         report,
         iter,
         scheduler!!.getScheduler(),
         fullSchedule,
-        logger!!.getLogger(report, fullSchedule),
-        interleaveMemoryOps,
-        ignoreUnhandledExceptions,
-        maxScheduledStep)
+        logger!!.getLogger(report, fullSchedule))
   }
 }
 
 data class Configuration(
-    val exec: () -> Unit,
+    val executionInfo: ExecutionInfo,
     val report: String,
     val iter: Int,
     val scheduler: Scheduler,
     val fullSchedule: Boolean,
     val logger: LoggerBase,
-    val interleaveMemoryOps: Boolean,
-    val ignoreUnhandledExceptions: Boolean,
-    val maxScheduledStep: Int,
 ) {}

core/src/main/kotlin/cmu/pasta/fray/core/command/Executor.kt

@@ -0,0 +1,47 @@
+package cmu.pasta.fray.core.command
+
+import java.net.URI
+import java.net.URLClassLoader
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+@Serializable
+sealed interface Executor {
+  fun execute()
+}
+
+@Serializable
+@SerialName("executor")
+data class MethodExecutor(
+    val clazz: String,
+    val method: String,
+    val args: List<String>,
+    val classpaths: List<String>
+) : Executor {
+  override fun execute() {
+    val classLoader =
+        URLClassLoader(
+            classpaths.map { it -> URI("file://$it").toURL() }.toTypedArray(),
+            Thread.currentThread().contextClassLoader)
+    Thread.currentThread().contextClassLoader = classLoader
+    val clazz = Class.forName(clazz, true, classLoader)
+    if (args.isEmpty() && method != "main") {
+      val m = clazz.getMethod(method)
+      if (m.modifiers and java.lang.reflect.Modifier.STATIC == 0) {
+        val obj = clazz.getConstructor().newInstance()
+        m.invoke(obj)
+      } else {
+        m.invoke(null)
+      }
+    } else {
+      val m = clazz.getMethod(method, Array<String>::class.java)
+      m.invoke(null, args.toTypedArray())
+    }
+  }
+}
+
+class LambdaExecutor(val lambda: () -> Unit) : Executor {
+  override fun execute() {
+    lambda()
+  }
+}

instrumentation/src/main/kotlin/cmu/pasta/fray/instrumentation/visitors/CountDownLatchInstrumenter.kt

@@ -15,7 +15,7 @@ class CountDownLatchInstrumenter(cv: ClassVisitor) :
       signature: String?,
       exceptions: Array<out String>?
   ): MethodVisitor {
-    if (name == "await") {
+    if (name == "await" && descriptor == "()V") {
       val eMv = MethodEnterVisitor(mv, Runtime::onLatchAwait, access, name, descriptor, true, false)
       return MethodExitVisitor(
           eMv, Runtime::onLatchAwaitDone, access, name, descriptor, true, false, true)

instrumentation/src/main/kotlin/cmu/pasta/fray/instrumentation/visitors/TimedWaitInstrumenter.kt

@@ -38,6 +38,11 @@ class TimedWaitInstrumenter(cv: ClassVisitor) : ClassVisitor(ASM9, cv) {
         Runtime::onConditionAwaitUntil
       }
     }
+    if (owner == "java/util/concurrent/CountDownLatch" &&
+        name == "await" &&
+        descriptor == "(JLjava/util/concurrent/TimeUnit;)Z") {
+      return Runtime::onLatchAwaitTimeout
+    }
     return null
   }