From 5455c6e155da9b6a741b90f5519849a1ef761107 Mon Sep 17 00:00:00 2001
From: Ao Li <5557706+aoli-al@users.noreply.github.com>
Date: Mon, 2 Sep 2024 10:10:33 -0400
Subject: [PATCH] Instrument sleep statements. (#38)

---
 .../org/pastalab/fray/core/RuntimeDelegate.kt    |  7 +++++++
 .../base/visitors/SleepInstrumenter.kt           | 16 +++++++++++++++-
 .../java/org/pastalab/fray/runtime/Delegate.java | 13 +++++++++++++
 .../java/org/pastalab/fray/runtime/Runtime.java  | 12 ++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/core/src/main/kotlin/org/pastalab/fray/core/RuntimeDelegate.kt b/core/src/main/kotlin/org/pastalab/fray/core/RuntimeDelegate.kt
index b09e2c6a..2d84c75e 100644
--- a/core/src/main/kotlin/org/pastalab/fray/core/RuntimeDelegate.kt
+++ b/core/src/main/kotlin/org/pastalab/fray/core/RuntimeDelegate.kt
@@ -1,5 +1,6 @@
 package org.pastalab.fray.core
 
+import java.time.Duration
 import java.util.*
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.ForkJoinPool
@@ -780,4 +781,10 @@ class RuntimeDelegate(val context: RunContext) : org.pastalab.fray.runtime.Deleg
     entered.set(false)
     return probe
   }
+
+  override fun onThreadSleepDuration(duration: Duration?) {}
+
+  override fun onThreadSleepMillis(millis: Long) {}
+
+  override fun onThreadSleepMillisNanos(millis: Long, nanos: Int) {}
 }
diff --git a/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/visitors/SleepInstrumenter.kt b/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/visitors/SleepInstrumenter.kt
index b46b3a6b..4ab17f58 100644
--- a/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/visitors/SleepInstrumenter.kt
+++ b/instrumentation/base/src/main/kotlin/org/pastalab/fray/instrumentation/base/visitors/SleepInstrumenter.kt
@@ -3,7 +3,9 @@ package org.pastalab.fray.instrumentation.base.visitors
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes.ASM9
+import org.objectweb.asm.Type
 import org.objectweb.asm.commons.GeneratorAdapter
+import org.pastalab.fray.runtime.Runtime
 
 class SleepInstrumenter(cv: ClassVisitor) : ClassVisitor(ASM9, cv) {
   override fun visitMethod(
@@ -28,7 +30,19 @@ class SleepInstrumenter(cv: ClassVisitor) : ClassVisitor(ASM9, cv) {
           isInterface: Boolean
       ) {
         if (owner == "java/lang/Thread" && name == "sleep") {
-          pop2()
+          if (descriptor == "(J)V") {
+            invokeStatic(
+                Type.getObjectType(Runtime::class.java.name.replace(".", "/")),
+                Utils.kFunctionToASMMethod(Runtime::onThreadSleepMillis))
+          } else if (descriptor == "(JI)V") {
+            invokeStatic(
+                Type.getObjectType(Runtime::class.java.name.replace(".", "/")),
+                Utils.kFunctionToASMMethod(Runtime::onThreadSleepMillisNanos))
+          } else {
+            invokeStatic(
+                Type.getObjectType(Runtime::class.java.name.replace(".", "/")),
+                Utils.kFunctionToASMMethod(Runtime::onThreadSleepDuration))
+          }
         } else {
           super.visitMethodInsn(opcode, owner, name, descriptor, isInterface)
         }
diff --git a/runtime/src/main/java/org/pastalab/fray/runtime/Delegate.java b/runtime/src/main/java/org/pastalab/fray/runtime/Delegate.java
index c7032c10..56aa2701 100644
--- a/runtime/src/main/java/org/pastalab/fray/runtime/Delegate.java
+++ b/runtime/src/main/java/org/pastalab/fray/runtime/Delegate.java
@@ -1,5 +1,6 @@
 package org.pastalab.fray.runtime;
 
+import java.time.Duration;
 import java.util.Date;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ForkJoinPool;
@@ -281,5 +282,17 @@ public ForkJoinPool onForkJoinPoolCommonPool(ForkJoinPool pool) {
     public int onThreadLocalRandomGetProbe(int probe) {
         return probe;
     }
+
+    public void onThreadSleepMillis(long millis) throws InterruptedException {
+        Thread.sleep(millis);
+    }
+
+    public void onThreadSleepDuration(Duration duration) throws InterruptedException {
+        Thread.sleep(duration);
+    }
+
+    public void onThreadSleepMillisNanos(long millis, int nanos) throws InterruptedException {
+        Thread.sleep(millis, nanos);
+    }
 }
 
diff --git a/runtime/src/main/java/org/pastalab/fray/runtime/Runtime.java b/runtime/src/main/java/org/pastalab/fray/runtime/Runtime.java
index afa32768..d6aba60d 100644
--- a/runtime/src/main/java/org/pastalab/fray/runtime/Runtime.java
+++ b/runtime/src/main/java/org/pastalab/fray/runtime/Runtime.java
@@ -1,5 +1,6 @@
 package org.pastalab.fray.runtime;
 
+import java.time.Duration;
 import java.util.Date;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ForkJoinPool;
@@ -372,4 +373,15 @@ public static int onThreadLocalRandomGetProbe(int probe) {
         return DELEGATE.onThreadLocalRandomGetProbe(probe);
     }
 
+    public static void onThreadSleepMillis(long millis) throws InterruptedException {
+        DELEGATE.onThreadSleepMillis(millis);
+    }
+
+    public static void onThreadSleepDuration(Duration duration) throws InterruptedException {
+        DELEGATE.onThreadSleepDuration(duration);
+    }
+
+    public static void onThreadSleepMillisNanos(long millis, int nanos) throws InterruptedException {
+        DELEGATE.onThreadSleepMillisNanos(millis, nanos);
+    }
 }