Only check for deadlocks in deadlock busting thread

Signed-off-by: Alessandro Bellina <[email protected]>

Author: abellina

src/main/cpp/src/SparkResourceAdaptorJni.cpp

@@ -44,8 +44,6 @@ constexpr char const* CPU_RETRY_OOM_CLASS = "com/nvidia/spark/rapids/jni/CpuRetr
 constexpr char const* CPU_SPLIT_AND_RETRY_OOM_CLASS =
   "com/nvidia/spark/rapids/jni/CpuSplitAndRetryOOM";
 constexpr char const* THREAD_REG_CLASS      = "com/nvidia/spark/rapids/jni/ThreadStateRegistry";
-constexpr char const* IS_THREAD_BLOCKED     = "isThreadBlocked";
-constexpr char const* IS_THREAD_BLOCKED_SIG = "(J)Z";
 constexpr char const* REMOVE_THREAD         = "removeThread";
 constexpr char const* REMOVE_THREAD_SIG     = "(J)V";
 
@@ -55,7 +53,6 @@ std::mutex jni_mutex;
 bool is_jni_loaded = false;
 jclass ThreadStateRegistry_jclass;
 jmethodID removeThread_method;
-jmethodID isThreadBlocked_method;
 
 void cache_thread_reg_jni(JNIEnv* env)
 {
@@ -67,9 +64,6 @@ void cache_thread_reg_jni(JNIEnv* env)
   removeThread_method = env->GetStaticMethodID(cls, REMOVE_THREAD, REMOVE_THREAD_SIG);
   if (removeThread_method == nullptr) { return; }
 
-  isThreadBlocked_method = env->GetStaticMethodID(cls, IS_THREAD_BLOCKED, IS_THREAD_BLOCKED_SIG);
-  if (isThreadBlocked_method == nullptr) { return; }
-
   // Convert local reference to global so it cannot be garbage collected.
   ThreadStateRegistry_jclass = static_cast<jclass>(env->NewGlobalRef(cls));
   if (ThreadStateRegistry_jclass == nullptr) { return; }
@@ -1120,10 +1114,10 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
     return ret;
   }
 
-  void check_and_break_deadlocks()
+  void check_and_break_deadlocks(std::unordered_set<long> const& java_blocked_thread_ids)
   {
     std::unique_lock<std::mutex> lock(state_mutex);
-    check_and_update_for_bufn(lock);
+    check_and_update_for_bufn(lock, java_blocked_thread_ids);
   }
 
   bool cpu_prealloc(size_t const amount, bool const blocking)
@@ -1373,7 +1367,7 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
             // Before we can wait it is possible that the throw didn't release anything
             // and the other threads didn't get unblocked by this, so we need to
             // check again to see if this was fixed or not.
-            check_and_update_for_bufn(lock);
+            check_and_update_for_bufn_state_machine_only(lock);
             // If that caused us to transition to a new state, then we need to adjust to it
             // appropriately...
             if (is_blocked(thread->second.state)) {
@@ -1696,127 +1690,55 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
           break;
         default: break;
       }
-      wake_next_highest_priority_blocked(lock, false, is_for_cpu);
+      wake_next_highest_priority_blocked(lock, is_for_cpu);
     }
   }
 
   /**
-   * Wake the highest priority blocked (not BUFN) thread so it can make progress,
-   * or the highest priority BUFN thread if all of the tasks are in some form of BUFN
-   * and this was triggered by a free.
+   * Wake the highest priority blocked (not BUFN) thread so it can make progress
    *
    * This is typically called when a free happens, or an alloc succeeds.
-   * @param is_from_free true if a free happen.
    * @param is_for_cpu true if it was a CPU operation (free or alloc)
    */
   void wake_next_highest_priority_blocked(std::unique_lock<std::mutex> const& lock,
-                                          bool const is_from_free,
                                           bool const is_for_cpu)
   {
     // 1. Find the highest priority blocked thread, for the alloc that matches
     thread_priority to_wake(-1, -1);
-    bool is_to_wake_set = false;
-    for (auto const& [thread_d, t_state] : threads) {
+    full_thread_state* thread_to_wake = nullptr;
+    for (auto& [thread_d, t_state] : threads) {
       thread_state const& state = t_state.state;
       if (state == thread_state::THREAD_BLOCKED && is_for_cpu == t_state.is_cpu_alloc) {
         thread_priority current = t_state.priority();
-        if (!is_to_wake_set || to_wake < current) {
+        if (thread_to_wake == nullptr || to_wake < current) {
           to_wake        = current;
-          is_to_wake_set = true;
+          thread_to_wake = &t_state;
         }
       }
     }
     // 2. wake up that thread
-    long const thread_id_to_wake = to_wake.get_thread_id();
-    if (thread_id_to_wake > 0) {
-      auto const thread = threads.find(thread_id_to_wake);
-      if (thread != threads.end()) {
-        switch (thread->second.state) {
-          case thread_state::THREAD_BLOCKED:
-            transition(thread->second, thread_state::THREAD_RUNNING);
-            thread->second.wake_condition->notify_all();
-            break;
-          default: {
-            std::stringstream ss;
-            ss << "internal error expected to only wake up blocked threads " << thread_id_to_wake
-               << " " << as_str(thread->second.state);
-            throw std::runtime_error(ss.str());
-          }
-        }
-      }
-    } else if (is_from_free) {
-      // 3. Otherwise look to see if we are in a BUFN deadlock state.
-      //
-      // Memory was freed and if all of the tasks are in a BUFN state,
-      // then we want to wake up the highest priority one so it can make progress
-      // instead of trying to split its input. But we only do this if it
-      // is a different thread that is freeing memory from the one we want to wake up.
-      // This is because if the threads are the same no new memory is being added
-      // to what that task has access to and the task may never throw a retry and split.
-      // Instead it would just keep retrying and freeing the same memory each time.
-      std::map<long, long> pool_bufn_task_thread_count;
-      std::map<long, long> pool_task_thread_count;
-      std::unordered_set<long> bufn_task_ids;
-      std::unordered_set<long> all_task_ids;
-      is_in_deadlock(
-        pool_bufn_task_thread_count, pool_task_thread_count, bufn_task_ids, all_task_ids, lock);
-      bool const all_bufn = all_task_ids.size() == bufn_task_ids.size();
-      if (all_bufn) {
-        thread_priority to_wake(-1, -1);
-        bool is_to_wake_set = false;
-        for (auto const& [thread_id, t_state] : threads) {
-          switch (t_state.state) {
-            case thread_state::THREAD_BUFN: {
-              if (is_for_cpu == t_state.is_cpu_alloc) {
-                thread_priority current = t_state.priority();
-                if (!is_to_wake_set || to_wake < current) {
-                  to_wake        = current;
-                  is_to_wake_set = true;
-                }
-              }
-            } break;
-            default: break;
-          }
-        }
-        // 4. Wake up the BUFN thread if we should
-        if (is_to_wake_set) {
-          long const thread_id_to_wake = to_wake.get_thread_id();
-          if (thread_id_to_wake > 0) {
-            // Don't wake up yourself on a free. It is not adding more memory for this thread
-            // to use on a retry and we might need a split instead to break a deadlock
-            auto const this_id = static_cast<long>(pthread_self());
-            auto const thread  = threads.find(thread_id_to_wake);
-            if (thread != threads.end() && thread->first != this_id) {
-              switch (thread->second.state) {
-                case thread_state::THREAD_BUFN:
-                  transition(thread->second, thread_state::THREAD_RUNNING);
-                  thread->second.wake_condition->notify_all();
-                  break;
-                case thread_state::THREAD_BUFN_WAIT:
-                  transition(thread->second, thread_state::THREAD_RUNNING);
-                  // no need to notify anyone, we will just retry without blocking...
-                  break;
-                case thread_state::THREAD_BUFN_THROW:
-                  // This should really never happen, this is a temporary state that is here only
-                  // while the lock is held, but just in case we don't want to mess it up, or throw
-                  // an exception.
-                  break;
-                default: {
-                  std::stringstream ss;
-                  ss << "internal error expected to only wake up blocked threads "
-                     << thread_id_to_wake << " " << as_str(thread->second.state);
-                  throw std::runtime_error(ss.str());
-                }
-              }
-            }
-          }
-        }
-      }
+    if (thread_to_wake != nullptr) {
+      transition(*thread_to_wake, thread_state::THREAD_RUNNING);
+      thread_to_wake->wake_condition->notify_all();
     }
   }
 
-  bool is_thread_bufn_or_above(JNIEnv* env, full_thread_state const& state)
+  /**
+   * Returns a boolean indicating if the thread is in BUFN (Blocked Until
+   * Further Notice) state or above. A thread is considered BUFN or above if:
+   * - It has pool_blocked set to true, OR
+   * - Its state is THREAD_BUFN, OR
+   * - Its thread ID is in the java_blocked_thread_ids set if it is provided.
+   *
+   * Threads in THREAD_BLOCKED state are NOT considered BUFN or above.
+   */
+  bool is_thread_bufn_or_above(
+    full_thread_state const& state,
+    std::optional<std::unordered_set<long>> const& java_blocked_thread_ids)
   {
+    LOG_INFO("is_thread_bufn_or_above: state: {}, java_blocked_thread_ids: {}",
+             state.thread_id,
+             to_string(java_blocked_thread_ids));
     bool ret = false;
     if (state.pool_blocked) {
       ret = true;
@@ -1828,8 +1750,9 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
           ret = true;
           break;
         default:
-          ret = env->CallStaticBooleanMethod(
-            ThreadStateRegistry_jclass, isThreadBlocked_method, state.thread_id);
+          if (java_blocked_thread_ids.has_value()) {
+            ret = java_blocked_thread_ids.value().contains(state.thread_id);
+          }
           break;
       }
     }
@@ -1853,18 +1776,23 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
     return oss.str();
   }
 
-  bool is_in_deadlock(std::map<long, long>& pool_bufn_task_thread_count,
+  template <typename SetType>
+  std::string to_string(std::optional<SetType> const& set, std::string const& separator = ",")
+  {
+    if (set.has_value()) {
+      return to_string(set.value(), separator);
+    } else {
+      return "{}";
+    }
+  }
+
+  bool is_in_deadlock(std::unique_lock<std::mutex> const& lock,
+                      std::map<long, long>& pool_bufn_task_thread_count,
                       std::map<long, long>& pool_task_thread_count,
                       std::unordered_set<long>& bufn_task_ids,
                       std::unordered_set<long>& all_task_ids,
-                      std::unique_lock<std::mutex> const& lock)
+                      std::optional<std::unordered_set<long>> const& java_blocked_thread_ids)
   {
-    JNIEnv* env = nullptr;
-    if (jvm->GetEnv(reinterpret_cast<void**>(&env), cudf::jni::MINIMUM_JNI_VERSION) != JNI_OK) {
-      throw std::runtime_error("Cloud not init JNI callbacks");
-    }
-    cache_thread_reg_jni(env);
-
     // If all of the tasks are blocked, then we are in a deadlock situation
     // and we need to wake something up. In theory if any one thread is still
     // doing something, then we are not deadlocked. But the problem is detecting
@@ -1897,7 +1825,7 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
       long const task_id = t_state.task_id;
       if (task_id >= 0) {
         all_task_ids.insert(task_id);
-        bool const is_bufn_plus = is_thread_bufn_or_above(env, t_state);
+        bool const is_bufn_plus = is_thread_bufn_or_above(t_state, java_blocked_thread_ids);
         if (is_bufn_plus) { bufn_task_ids.insert(task_id); }
         if (is_bufn_plus || t_state.state == thread_state::THREAD_BLOCKED) {
           blocked_task_ids.insert(task_id);
@@ -1918,7 +1846,7 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
           }
         }
 
-        bool const is_bufn_plus = is_thread_bufn_or_above(env, t_state);
+        bool const is_bufn_plus = is_thread_bufn_or_above(t_state, java_blocked_thread_ids);
         if (is_bufn_plus) {
           for (auto const& task_id : t_state.pool_task_ids) {
             auto const it = pool_bufn_task_thread_count.find(task_id);
@@ -1974,19 +1902,44 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
       "DETAIL", -1, -1, thread_state::UNKNOWN, "States of all threads: {}", get_threads_string());
   }
 
+  /**
+   * This method is only called from code that only cares about the state machine state,
+   * and is happening very often. Our deadlock busting thread will invoke
+   * `check_and_update_for_bufn` directly, and pass along a set of java native thread ids
+   * that are blocked in the java state (Thread.getState).
+   *
+   * This split in invocation is done to make the critical sections faster, and leave
+   * deadlock busting to the deadlock thread.
+   */
+  void check_and_update_for_bufn_state_machine_only(const std::unique_lock<std::mutex>& lock)
+  {
+    // we pass nullopt because we are calling this method from a place in the code that has
+    // no java knowledge of blocked threads.
+    check_and_update_for_bufn(lock, /*java_blocked_thread_ids*/ std::nullopt);
+  }
+
   /**
    * Check to see if any threads need to move to BUFN. This should be
    * called when a task or shuffle thread becomes blocked so that we can
    * check to see if one of them needs to become BUFN or do a split and rollback.
+   *
+   * If this method is being called from the deadlock busting thread, we will pass
+   * along a set of java native thread ids that are blocked in the java state (Thread.getState).
    */
-  void check_and_update_for_bufn(const std::unique_lock<std::mutex>& lock)
+  void check_and_update_for_bufn(
+    const std::unique_lock<std::mutex>& lock,
+    std::optional<std::unordered_set<long>> const& java_blocked_thread_ids)
   {
     std::map<long, long> pool_bufn_task_thread_count;
     std::map<long, long> pool_task_thread_count;
     std::unordered_set<long> bufn_task_ids;
     std::unordered_set<long> all_task_ids;
-    bool const need_to_break_deadlock = is_in_deadlock(
-      pool_bufn_task_thread_count, pool_task_thread_count, bufn_task_ids, all_task_ids, lock);
+    bool const need_to_break_deadlock = is_in_deadlock(lock,
+                                                       pool_bufn_task_thread_count,
+                                                       pool_task_thread_count,
+                                                       bufn_task_ids,
+                                                       all_task_ids,
+                                                       java_blocked_thread_ids);
     if (need_to_break_deadlock) {
       // Find the task thread with the lowest priority that is not already BUFN
       thread_priority to_bufn(-1, -1);
@@ -2031,6 +1984,7 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
           thread->second.wake_condition->notify_all();
         }
       }
+
       // We now need a way to detect if we need to split the input and retry.
       // This happens when all of the tasks are also blocked until
       // further notice. So we are going to treat a task as blocked until
@@ -2154,7 +2108,7 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
       // do not retry if the thread is not registered...
       ret = false;
     }
-    check_and_update_for_bufn(lock);
+    check_and_update_for_bufn_state_machine_only(lock);
     return ret;
   }
 
@@ -2222,7 +2176,7 @@ class spark_resource_adaptor final : public rmm::mr::device_memory_resource {
         }
       }
     }
-    wake_next_highest_priority_blocked(lock, true, is_for_cpu);
+    wake_next_highest_priority_blocked(lock, is_for_cpu);
   }
 
   void do_deallocate(void* p, std::size_t size, rmm::cuda_stream_view stream) noexcept override
@@ -2585,13 +2539,16 @@ JNIEXPORT void JNICALL Java_com_nvidia_spark_rapids_jni_SparkResourceAdaptor_end
 }
 
 JNIEXPORT void JNICALL Java_com_nvidia_spark_rapids_jni_SparkResourceAdaptor_checkAndBreakDeadlocks(
-  JNIEnv* env, jclass, jlong ptr)
+  JNIEnv* env, jclass, jlong ptr, jlongArray jblocked_thread_ids)
 {
   JNI_NULL_CHECK(env, ptr, "resource_adaptor is null", );
   JNI_TRY
   {
     auto mr = reinterpret_cast<spark_resource_adaptor*>(ptr);
-    mr->check_and_break_deadlocks();
+    cudf::jni::native_jlongArray blocked_thread_ids(env, jblocked_thread_ids);
+    std::unordered_set<long> blocked_thread_ids_set(blocked_thread_ids.begin(),
+                                                    blocked_thread_ids.end());
+    mr->check_and_break_deadlocks(blocked_thread_ids_set);
   }
   JNI_CATCH(env, );
 }