Fix deoptmize static methods

lsplant/src/main/jni/art/runtime/class_linker.hpp

@@ -68,27 +68,41 @@ private:
         std::list<std::tuple<art::ArtMethod *, void *>> out;
         auto class_def = mirror_class->GetClassDef();
         if (!class_def) return out;
+        {
             std::shared_lock lk(hooked_classes_lock_);
             if (auto found = hooked_classes_.find(class_def); found != hooked_classes_.end())
                 [[unlikely]] {
-            LOGD("Before fixup %s, backup hooked methods' trampoline",
+                LOGV("Before fixup %s, backup %zu hooked methods' trampoline",
-                 mirror_class->GetDescriptor().c_str());
+                     mirror_class->GetDescriptor().c_str(), found->second.size());
                 for (auto method : found->second) {
                     out.emplace_back(method, method->GetEntryPoint());
                 }
             }
+        }
+        {
+            std::shared_lock lk(deoptimized_methods_lock_);
+            if (auto found = deoptimized_classes_.find(class_def);
+                found != deoptimized_classes_.end()) [[unlikely]] {
+                LOGV("Before fixup %s, backup %zu deoptimized methods' trampoline",
+                     mirror_class->GetDescriptor().c_str(), found->second.size());
+                for (auto method : found->second) {
+                    out.emplace_back(method, method->GetEntryPoint());
+                }
+            }
+        }
         return out;
     }
 
     static void FixTrampoline(const std::list<std::tuple<art::ArtMethod *, void *>> &methods) {
         std::shared_lock lk(hooked_methods_lock_);
         for (const auto &[art_method, old_trampoline] : methods) {
+            auto new_trampoline = art_method->GetEntryPoint();
+            art_method->SetEntryPoint(old_trampoline);
+            if (IsDeoptimized(art_method)) continue;
             if (auto backup_method = IsHooked(art_method); backup_method) [[likely]] {
-                if (auto new_trampoline = art_method->GetEntryPoint();
+                if (new_trampoline != old_trampoline) [[unlikely]] {
-                    new_trampoline != old_trampoline) [[unlikely]] {
                     LOGV("propagate entrypoint for %s", backup_method->PrettyMethod(true).data());
                     backup_method->SetEntryPoint(new_trampoline);
-                    art_method->SetEntryPoint(old_trampoline);
                 }
             }
         }
@@ -159,9 +173,14 @@ public:
         // Android 13
         if (art_quick_to_interpreter_bridgeSym && art_quick_generic_jni_trampolineSym) [[likely]] {
             if (art_method->GetAccessFlags() & ArtMethod::kAccNative) [[unlikely]] {
+                LOGV("deoptimize native method %s from %p to %p",
+                     art_method->PrettyMethod(true).data(), art_method->GetEntryPoint(),
+                     art_quick_generic_jni_trampolineSym);
                 art_method->SetEntryPoint(
                     reinterpret_cast<void *>(art_quick_generic_jni_trampolineSym));
             } else {
+                LOGV("deoptimize method %s from %p to %p", art_method->PrettyMethod(true).data(),
+                     art_method->GetEntryPoint(), art_quick_to_interpreter_bridgeSym);
                 art_method->SetEntryPoint(
                     reinterpret_cast<void *>(art_quick_to_interpreter_bridgeSym));
             }

lsplant/src/main/jni/common.hpp

@@ -1,12 +1,12 @@
 #pragma once
 
+#include <absl/container/flat_hash_map.h>
+#include <absl/container/flat_hash_set.h>
 #include <sys/system_properties.h>
 
 #include <list>
 #include <shared_mutex>
 #include <string_view>
-#include <absl/container/flat_hash_map.h>
-#include <absl/container/flat_hash_set.h>
 
 #include "logging.hpp"
 #include "lsplant.hpp"
@@ -64,14 +64,13 @@ inline auto IsJavaDebuggable(JNIEnv *env) {
             LOGE("Failed to find VMRuntime");
             return false;
         }
-        auto get_runtime_method = JNI_GetStaticMethodID(env, runtime_class, "getRuntime",
+        auto get_runtime_method =
-                                                        "()Ldalvik/system/VMRuntime;");
+            JNI_GetStaticMethodID(env, runtime_class, "getRuntime", "()Ldalvik/system/VMRuntime;");
         if (!get_runtime_method) {
             LOGE("Failed to find VMRuntime.getRuntime()");
             return false;
         }
-        auto is_debuggable_method =
+        auto is_debuggable_method = JNI_GetMethodID(env, runtime_class, "isJavaDebuggable", "()Z");
-                JNI_GetMethodID(env, runtime_class, "isJavaDebuggable", "()Z");
         if (!is_debuggable_method) {
             LOGE("Failed to find VMRuntime.isJavaDebuggable()");
             return false;
@@ -105,12 +104,19 @@ namespace {
 inline absl::flat_hash_map<art::ArtMethod *, std::pair<jobject, art::ArtMethod *>> hooked_methods_;
 inline std::shared_mutex hooked_methods_lock_;
 
-inline std::list<std::pair<art::ArtMethod *, art::ArtMethod *>> jit_movements_;
-inline std::shared_mutex jit_movements_lock_;
-
 inline absl::flat_hash_map<const art::dex::ClassDef *, absl::flat_hash_set<art::ArtMethod *>>
     hooked_classes_;
 inline std::shared_mutex hooked_classes_lock_;
+
+inline absl::flat_hash_set<art::ArtMethod *> deoptimized_methods_set_;
+inline std::shared_mutex deoptimized_methods_lock_;
+
+inline absl::flat_hash_map<const art::dex::ClassDef *, absl::flat_hash_set<art::ArtMethod *>>
+    deoptimized_classes_;
+inline std::shared_mutex deoptimized_classes_lock_;
+
+inline std::list<std::pair<art::ArtMethod *, art::ArtMethod *>> jit_movements_;
+inline std::shared_mutex jit_movements_lock_;
 }  // namespace
 
 inline art::ArtMethod *IsHooked(art::ArtMethod *art_method, bool including_backup = false) {
@@ -131,6 +137,11 @@ inline art::ArtMethod *IsBackup(art::ArtMethod *art_method) {
     return nullptr;
 }
 
+inline bool IsDeoptimized(art::ArtMethod *art_method) {
+    std::shared_lock lk(deoptimized_methods_lock_);
+    return deoptimized_methods_set_.contains(art_method);
+}
+
 inline std::list<std::pair<art::ArtMethod *, art::ArtMethod *>> GetJitMovements() {
     std::unique_lock lk(jit_movements_lock_);
     return std::move(jit_movements_);
@@ -138,14 +149,25 @@ inline std::list<std::pair<art::ArtMethod *, art::ArtMethod *>> GetJitMovements(
 
 inline void RecordHooked(art::ArtMethod *target, const art::dex::ClassDef *class_def,
                          jobject reflected_backup, art::ArtMethod *backup) {
+    {
+        std::unique_lock lk(hooked_classes_lock_);
+        hooked_classes_[class_def].emplace(target);
+    }
     {
         std::unique_lock lk(hooked_methods_lock_);
         hooked_methods_[target] = {reflected_backup, backup};
         hooked_methods_[backup] = {nullptr, target};
     }
+}
+
+inline void RecordDeoptimized(const art::dex::ClassDef *class_def, art::ArtMethod *art_method) {
     {
-        std::unique_lock lk(hooked_classes_lock_);
+        std::unique_lock lk(deoptimized_classes_lock_);
-        hooked_classes_[class_def].emplace(target);
+        deoptimized_classes_[class_def].emplace(art_method);
+    }
+    {
+        std::unique_lock lk(deoptimized_methods_lock_);
+        deoptimized_methods_set_.insert(art_method);
     }
 }
 

lsplant/src/main/jni/lsplant.cc

@@ -489,12 +489,12 @@ bool DoHook(ArtMethod *target, ArtMethod *hook, ArtMethod *backup) {
     LOGV("Hooking: target = %s(%p), hook = %s(%p), backup = %s(%p)", target->PrettyMethod().c_str(),
          target, hook->PrettyMethod().c_str(), hook, backup->PrettyMethod().c_str(), backup);
 
-    if (auto *trampoline = GenerateTrampolineFor(hook); !trampoline) {
+    if (auto *entrypoint = GenerateTrampolineFor(hook); !entrypoint) {
         LOGE("Failed to generate trampoline");
         return false;
         // NOLINTNEXTLINE
     } else {
-        LOGV("Generated trampoline %p", trampoline);
+        LOGV("Generated trampoline %p", entrypoint);
 
         target->SetNonCompilable();
         hook->SetNonCompilable();
@@ -504,7 +504,7 @@ bool DoHook(ArtMethod *target, ArtMethod *hook, ArtMethod *backup) {
 
         target->ClearFastInterpretFlag();
 
-        target->SetEntryPoint(trampoline);
+        target->SetEntryPoint(entrypoint);
 
         if (!backup->IsStatic()) backup->SetPrivate();
 
@@ -621,6 +621,10 @@ using ::lsplant::IsHooked;
         if (!is_proxy) [[likely]] {
             RecordJitMovement(target, backup);
         }
+        // Always record backup as deoptimized since we dont want its entrypoint to be updated
+        // by FixupStaticTrampolines on hooker class
+        // Used hook's declaring class here since backup's is no longer the same with hook's
+        RecordDeoptimized(hook->GetDeclaringClass()->GetClassDef(), backup);
         return global_backup;
     }
 
@@ -687,6 +691,8 @@ using ::lsplant::IsHooked;
         return false;
     }
     auto *art_method = ArtMethod::FromReflectedMethod(env, method);
+    // record the original but not the backup
+    RecordDeoptimized(art_method->GetDeclaringClass()->GetClassDef(), art_method);
     if (auto *backup = IsHooked(art_method); backup) {
         art_method = backup;
     }