[Wisp] Support preempt in interpreter

Summary: Preempt check could be left out when thread fall into
interpreter or unhandled thread_state, this patch added check
in interpreter and cleared preempted flag when needed.

Test Plan: com/alibaba/wisp/thread/PreemptTest.java

Reviewed-by: leiyu, zhengxiaolinX

Issue: dragonwell-project/dragonwell8#204

src/share/vm/interpreter/interpreterRuntime.cpp

@@ -1062,6 +1062,9 @@ IRT_ENTRY(void, InterpreterRuntime::at_safepoint(JavaThread* thread))
     // then we may have JVMTI work to do.
     JvmtiExport::at_single_stepping_point(thread, method(thread), bcp(thread));
   }
+  if (EnableCoroutine) {
+    Coroutine::after_safepoint(thread);
+  }
 IRT_END
 
 IRT_ENTRY(void, InterpreterRuntime::post_field_access(JavaThread *thread, oopDesc* obj,

src/share/vm/prims/unsafe.cpp

@@ -1481,7 +1481,7 @@ JVM_ENTRY(jboolean, CoroutineSupport_stealCoroutine(JNIEnv* env, jclass klass, j
   // The lock will also block coroutine switch operation,
   // so we must finish the steal operation as soon as possible.
   Coroutine* coro = (Coroutine*) coroPtr;
-  if (coro == NULL || coro->enable_steal_count() != coro->java_call_counter()) {
+  if (coro == NULL || coro->enable_steal_count() != coro->java_call_counter()|| coro->is_yielding()) {
       return false;       // an Exception throws and the coroutine being stealed is exited
   }
   assert(coro->thread() != thread, "steal from self");
@@ -1503,14 +1503,11 @@ JVM_ENTRY(jboolean, CoroutineSupport_stealCoroutine(JNIEnv* env, jclass klass, j
   return true;
 JVM_END
 
-JVM_ENTRY (void, CoroutineSupport_checkAndThrowException0(JNIEnv* env, jclass klass, jlong coroPtr))
+JVM_ENTRY (jboolean, CoroutineSupport_shouldThrowException0(JNIEnv* env, jclass klass, jlong coroPtr))
   assert(EnableCoroutine, "pre-condition");
   Coroutine* coro = (Coroutine*)coroPtr;
   assert(coro == thread->current_coroutine(), "coroutine is current");
-  if (!coro->is_yielding() && coro->clinit_call_count() == 0) {
-    ThreadToNativeFromVM ttnfv(thread);
-    throw_new(env, "ThreadDeath");
-  }
+  return !coro->is_yielding() && coro->clinit_call_count() == 0;
 JVM_END
 
 /// JVM_RegisterUnsafeMethods
@@ -1869,7 +1866,7 @@ JNINativeMethod coroutine_support_methods[] = {
     {CC"getNextCoroutine",        CC"(J)"COR,         FN_PTR(CoroutineSupport_getNextCoroutine)},
     {CC"moveCoroutine",           CC"(JJ)V",          FN_PTR(CoroutineSupport_moveCoroutine)},
     {CC"markThreadCoroutine",     CC"(J"COBA")V",     FN_PTR(CoroutineSupport_markThreadCoroutine)},
-    {CC"checkAndThrowException0", CC"(J)V",           FN_PTR(CoroutineSupport_checkAndThrowException0)},
+    {CC"shouldThrowException0", CC"(J)Z",           FN_PTR(CoroutineSupport_shouldThrowException0)},
 };
 
 #define COMPILE_CORO_METHODS_BEFORE (3)

src/share/vm/runtime/coroutine.cpp

@@ -580,6 +580,7 @@ void Coroutine::print_stack_on(outputStream* st) {
 bool WispThread::_wisp_booted = false;
 Method* WispThread::parkMethod = NULL;
 Method* WispThread::unparkMethod = NULL;
+Method* WispThread::yieldMethod = NULL;
 Method* WispThread::runOutsideWispMethod = NULL;
 GrowableArray<int>* WispThread::_proxy_unpark = NULL;
 
@@ -588,6 +589,7 @@ void WispThread::set_wisp_booted(Thread* thread) {
   // The flow should be changed.
   CallInfo callinfo;
   KlassHandle kh = KlassHandle(thread, SystemDictionary::com_alibaba_wisp_engine_WispTask_klass());
+  KlassHandle tkh = KlassHandle(thread, SystemDictionary::Thread_klass());
   LinkResolver::resolve_static_call(callinfo, kh,
       vmSymbols::park_name(), vmSymbols::long_void_signature(), KlassHandle(), false, true, thread);
   methodHandle method = callinfo.selected_method();
@@ -600,6 +602,12 @@ void WispThread::set_wisp_booted(Thread* thread) {
   assert(method.not_null(), "should have thrown exception");
   unparkMethod = method();
 
+  LinkResolver::resolve_static_call(callinfo, tkh,
+        vmSymbols::yield_name(), vmSymbols::void_method_signature(), KlassHandle(), false, true, thread);
+  method = callinfo.selected_method();
+  assert(method.not_null(), "should have thrown exception");
+  yieldMethod = method();
+
   LinkResolver::resolve_static_call(callinfo, kh,
       vmSymbols::runOutsideWisp_name(), vmSymbols::runnable_void_signature(), KlassHandle(), false, true, thread);
   method = callinfo.selected_method();
@@ -953,18 +961,39 @@ const char* WispThread::print_blocking_status(int status) {
 
 void Coroutine::after_safepoint(JavaThread* thread) {
   assert(Thread::current() == thread, "sanity check");
-  guarantee(thread->safepoint_state()->is_running(), "safepoint should finish");
+
+  //The only two entries are:
+  //1. SafepointSynchronize::handle_polling_page_exception(at_safepoint)
+  //2. InterpreterRuntime::at_safepoint(at_call_back)
+  //     call_back state would block in JavaCall -> ThreadState::trans
+  //     current yielding call triggers once safepoint ends.
+  if (thread->safepoint_state()->is_at_safepoint()) {
+      return;
+  }
+  // In InterpreterRuntime::at_safepoint: thread state is _thread_in_vm now.
+  // So VM Thread could make the safepoint state of current thread `running -> call_back`
+  // We should assert at first _running and then _call_back here.
+  assert(thread->safepoint_state()->is_running() ||
+         thread->safepoint_state()->is_at_call_back(), "illegal safepoint state");
 
   Coroutine* coroutine = thread->current_coroutine();
-  if (thread->thread_state() != _thread_in_Java ||
-      // indicates we're inside compiled code or interpreter.
-      // rather than thread state transition.
-      coroutine->_is_yielding || !thread->wisp_preempted() ||
-      thread->has_pending_exception() || thread->has_async_condition() ||
-      coroutine->in_critical(thread)) {
+  if (coroutine->_is_yielding || !thread->wisp_preempted()) {
     return;
   }
-
+  // filter unsupported state
+  JavaThreadState origin_thread_state = thread->thread_state();
+  if ((origin_thread_state != _thread_in_Java  && origin_thread_state != _thread_in_vm)
+    || thread->has_pending_exception() || thread->has_async_condition()) {
+    // clear preempted for next preempt
+    thread->set_wisp_preempted(false);
+    return;
+  }
+  // prevent preempting wisp internal
+  if (coroutine->in_critical(thread)) {
+    return;
+  }
+  // preempt only triggered by SafepointSynchronize::handle_polling_page_exception and
+  // InterpreterRuntime::at_safepoint.
   oop wisp_task = thread->current_coroutine()->_wisp_task;
   if (wisp_task != NULL) { // expose to perfCount and jstack
     int cnt = com_alibaba_wisp_engine_WispTask::get_preemptCount(wisp_task);
@@ -981,27 +1010,26 @@ void Coroutine::after_safepoint(JavaThread* thread) {
   // - The preempt mechanism should be enabled during "other" coroutines are executing
 
   thread->set_wisp_preempted(false);
-  ThreadInVMfromJava tiv(thread);
+  if (origin_thread_state == _thread_in_Java) {
+    ThreadStateTransition::transition_from_java(thread, _thread_in_vm);
+  }
+  assert(thread->thread_state() == _thread_in_vm, "illegal thread state");
   JavaValue result(T_VOID);
   JavaCallArguments args;
-  JavaCalls::call_static(&result,
-        KlassHandle(thread, SystemDictionary::Thread_klass()),
-        vmSymbols::yield_name(),
-        vmSymbols::void_method_signature(),
-        &args,
-        thread);
-  coroutine->_is_yielding = false;
-
-  if (thread->has_pending_exception()
-  && (thread->pending_exception()->klass() == SystemDictionary::OutOfMemoryError_klass()
-   || thread->pending_exception()->klass() == SystemDictionary::StackOverflowError_klass())) {
-      // throw expected vm error
-      return;
+  JavaCalls::call(&result, methodHandle(WispThread::yieldMethod), &args, thread);
+  if (origin_thread_state == _thread_in_Java) {
+    ThreadStateTransition::transition(thread, _thread_in_vm, _thread_in_Java);
   }
+  assert(thread->thread_state() == origin_thread_state, "illegal thread state");
+  coroutine->_is_yielding = false;
 
-  if (thread->has_pending_exception() || thread->has_async_condition()) {
-    guarantee(thread->pending_exception()->klass() == SystemDictionary::ThreadDeath_klass(),
-        "thread_death expected");
+  if (thread->has_pending_exception()) {
+    guarantee(thread->pending_exception()->klass() == SystemDictionary::OutOfMemoryError_klass() ||
+      thread->pending_exception()->klass() == SystemDictionary::StackOverflowError_klass() ||
+      thread->pending_exception()->klass() == SystemDictionary::ThreadDeath_klass(),
+      "Only SOF/OOM/ThreadDeath/TenantDeath happens here");
+    // If it's a SOF / OOM / ThreadDeath / TenantDeath exception, we'd clear it
+    // because polling page stub shouldn't have a pending exception.
     thread->clear_pending_exception();
   }
 }

src/share/vm/runtime/coroutine.hpp

@@ -346,6 +346,7 @@ class WispThread: public JavaThread {
   friend class Coroutine;
 private:
   static bool _wisp_booted;
+  static Method* yieldMethod;
   static Method* parkMethod;
   static Method* unparkMethod;
   static Method* runOutsideWispMethod;

src/share/vm/runtime/safepoint.cpp

@@ -1166,6 +1166,10 @@ void ThreadSafepointState::handle_polling_page_exception() {
     SafepointSynchronize::block(thread());
     set_at_poll_safepoint(false);
 
+    if (EnableCoroutine) {
+        Coroutine::after_safepoint(thread());
+    }
+
     // If we have a pending async exception deoptimize the frame
     // as otherwise we may never deliver it.
     if (thread()->has_async_condition()) {
@@ -1190,9 +1194,6 @@ void ThreadSafepointState::handle_polling_page_exception() {
         fatal("Exception installed and deoptimization is pending");
       }
     }
-    if (EnableCoroutine) {
-      Coroutine::after_safepoint(thread());
-    }
   }
 }
 

src/share/vm/runtime/safepoint.hpp

@@ -231,6 +231,8 @@ class ThreadSafepointState: public CHeapObj<mtThread> {
   JavaThread*  thread() const         { return _thread; }
   suspend_type type() const           { return _type; }
   bool         is_running() const     { return (_type==_running); }
+  bool         is_at_safepoint() const{ return (_type == _at_safepoint);}
+  bool         is_at_call_back() const{ return (_type == _call_back);}
   JavaThreadState orig_thread_state() const { return _orig_thread_state; }
 
   // Support for safepoint timeout (debugging)