Keep classes alive for stack traces with copied methods.

Test: Additional test in 141-class-unload.
Test: testrunner.py --host --optimizing
Bug: 263254495
Change-Id: Iea925ab96fb7c8d1b825c0b100e689ab722d3159
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index 8082d46..eaad1f0 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -2056,7 +2056,7 @@
     } else if (special_root->IsIntArray()) {
       size_t count = special_root->AsIntArray()->GetLength();
       if (oat_file->GetVdexFile()->GetNumberOfDexFiles() != count) {
-        *error_msg = "Cheksums count does not match";
+        *error_msg = "Checksums count does not match";
         return false;
       }
       static_assert(sizeof(VdexFile::VdexChecksum) == sizeof(int32_t));
@@ -2065,7 +2065,7 @@
       const VdexFile::VdexChecksum* vdex_checksums =
           oat_file->GetVdexFile()->GetDexChecksumsArray();
       if (memcmp(art_checksums, vdex_checksums, sizeof(VdexFile::VdexChecksum) * count) != 0) {
-        *error_msg = "Image and vdex cheksums did not match";
+        *error_msg = "Image and vdex checksums did not match";
         return false;
       }
     } else if (IsBootClassLoader(special_root.Get())) {
@@ -10399,9 +10399,73 @@
   CHECK(method->IsCopied());
   FindVirtualMethodHolderVisitor visitor(method, image_pointer_size_);
   VisitClasses(&visitor);
+  DCHECK(visitor.holder_ != nullptr);
   return visitor.holder_;
 }
 
+ObjPtr<mirror::ClassLoader> ClassLinker::GetHoldingClassLoaderOfCopiedMethod(Thread* self,
+                                                                             ArtMethod* method) {
+  // Note: `GetHoldingClassOfCopiedMethod(method)` is a lot more expensive than finding
+  // the class loader, so we're using it only to verify the result in debug mode.
+  CHECK(method->IsCopied());
+  gc::Heap* heap = Runtime::Current()->GetHeap();
+  // Check if the copied method is in the boot class path.
+  if (heap->IsBootImageAddress(method) || GetAllocatorForClassLoader(nullptr)->Contains(method)) {
+    DCHECK(GetHoldingClassOfCopiedMethod(method)->GetClassLoader() == nullptr);
+    return nullptr;
+  }
+  // Check if the copied method is in an app image.
+  // Note: Continuous spaces contain boot image spaces and app image spaces.
+  // However, they are sorted by address, so boot images are not trivial to skip.
+  ArrayRef<gc::space::ContinuousSpace* const> spaces(heap->GetContinuousSpaces());
+  DCHECK_GE(spaces.size(), heap->GetBootImageSpaces().size());
+  for (gc::space::ContinuousSpace* space : spaces) {
+    if (space->IsImageSpace()) {
+      gc::space::ImageSpace* image_space = space->AsImageSpace();
+      size_t offset = reinterpret_cast<const uint8_t*>(method) - image_space->Begin();
+      const ImageSection& methods_section = image_space->GetImageHeader().GetMethodsSection();
+      if (offset - methods_section.Offset() < methods_section.Size()) {
+        // Grab the class loader from the first non-BCP class in the app image class table.
+        // Note: If we allow classes from arbitrary parent or library class loaders in app
+        // images, this shall need to be updated to actually search for the exact class.
+        const ImageSection& class_table_section =
+            image_space->GetImageHeader().GetClassTableSection();
+        CHECK_NE(class_table_section.Size(), 0u);
+        const uint8_t* ptr = image_space->Begin() + class_table_section.Offset();
+        size_t read_count = 0;
+        ClassTable::ClassSet class_set(ptr, /*make_copy_of_data=*/ false, &read_count);
+        CHECK(!class_set.empty());
+        auto it = class_set.begin();
+        // No read barrier needed for references to non-movable image classes.
+        while ((*it).Read<kWithoutReadBarrier>()->IsBootStrapClassLoaded()) {
+          ++it;
+          CHECK(it != class_set.end());
+        }
+        ObjPtr<mirror::ClassLoader> class_loader =
+            (*it).Read<kWithoutReadBarrier>()->GetClassLoader();
+        DCHECK(GetHoldingClassOfCopiedMethod(method)->GetClassLoader() == class_loader);
+        return class_loader;
+      }
+    }
+  }
+  // Otherwise, the method must be in one of the `LinearAlloc` memory areas.
+  jweak result = nullptr;
+  {
+    ReaderMutexLock mu(self, *Locks::classlinker_classes_lock_);
+    for (const ClassLoaderData& data : class_loaders_) {
+      if (data.allocator->Contains(method)) {
+        result = data.weak_root;
+        break;
+      }
+    }
+  }
+  CHECK(result != nullptr) << "Did not find allocator holding the copied method: " << method
+      << " " << method->PrettyMethod();
+  // The `method` is alive, so the class loader must also be alive.
+  return ObjPtr<mirror::ClassLoader>::DownCast(
+      Runtime::Current()->GetJavaVM()->DecodeWeakGlobalAsStrong(result));
+}
+
 bool ClassLinker::DenyAccessBasedOnPublicSdk(ArtMethod* art_method ATTRIBUTE_UNUSED) const
     REQUIRES_SHARED(Locks::mutator_lock_) {
   // Should not be called on ClassLinker, only on AotClassLinker that overrides this.
diff --git a/runtime/class_linker.h b/runtime/class_linker.h
index 9a57bd4..86e2881 100644
--- a/runtime/class_linker.h
+++ b/runtime/class_linker.h
@@ -764,6 +764,11 @@
   ObjPtr<mirror::Class> GetHoldingClassOfCopiedMethod(ArtMethod* method)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
+  // Get the class loader holding class for a copied method.
+  ObjPtr<mirror::ClassLoader> GetHoldingClassLoaderOfCopiedMethod(Thread* self, ArtMethod* method)
+      REQUIRES_SHARED(Locks::mutator_lock_)
+      REQUIRES(!Locks::classlinker_classes_lock_);
+
   // Returns null if not found.
   // This returns a pointer to the class-table, without requiring any locking - including the
   // boot class-table. It is the caller's responsibility to access this under lock, if required.
diff --git a/runtime/jni/java_vm_ext.cc b/runtime/jni/java_vm_ext.cc
index b4e7938..7ae6c99 100644
--- a/runtime/jni/java_vm_ext.cc
+++ b/runtime/jni/java_vm_ext.cc
@@ -866,6 +866,11 @@
   return weak_globals_.Get(ref);
 }
 
+ObjPtr<mirror::Object> JavaVMExt::DecodeWeakGlobalAsStrong(IndirectRef ref) {
+  // The target is known to be alive. Simple `Get()` with read barrier is enough.
+  return weak_globals_.Get(ref);
+}
+
 ObjPtr<mirror::Object> JavaVMExt::DecodeWeakGlobalDuringShutdown(Thread* self, IndirectRef ref) {
   DCHECK_EQ(IndirectReferenceTable::GetIndirectRefKind(ref), kWeakGlobal);
   DCHECK(Runtime::Current()->IsShuttingDown(self));
diff --git a/runtime/jni/java_vm_ext.h b/runtime/jni/java_vm_ext.h
index 8c8b1b3..7f4f548 100644
--- a/runtime/jni/java_vm_ext.h
+++ b/runtime/jni/java_vm_ext.h
@@ -184,6 +184,11 @@
       REQUIRES_SHARED(Locks::mutator_lock_)
       REQUIRES(Locks::jni_weak_globals_lock_);
 
+  // Decode weak global as strong. Use only if the target object is known to be alive.
+  ObjPtr<mirror::Object> DecodeWeakGlobalAsStrong(IndirectRef ref)
+      REQUIRES_SHARED(Locks::mutator_lock_)
+      REQUIRES(!Locks::jni_weak_globals_lock_);
+
   // Like DecodeWeakGlobal() but to be used only during a runtime shutdown where self may be
   // null.
   ObjPtr<mirror::Object> DecodeWeakGlobalDuringShutdown(Thread* self, IndirectRef ref)
diff --git a/runtime/thread.cc b/runtime/thread.cc
index 78f9174..ff0a020 100644
--- a/runtime/thread.cc
+++ b/runtime/thread.cc
@@ -2980,9 +2980,19 @@
     methods_and_pcs->SetElementPtrSize</*kTransactionActive=*/ false, /*kCheckTransaction=*/ false>(
         static_cast<uint32_t>(methods_and_pcs->GetLength()) / 2 + count_, dex_pc, pointer_size_);
     // Save the declaring class of the method to ensure that the declaring classes of the methods
-    // do not get unloaded while the stack trace is live.
+    // do not get unloaded while the stack trace is live. However, this does not work for copied
+    // methods because the declaring class of a copied method points to an interface class which
+    // may be in a different class loader. Instead, retrieve the class loader associated with the
+    // allocator that holds the copied method. This is much cheaper than finding the actual class.
+    ObjPtr<mirror::Object> keep_alive;
+    if (UNLIKELY(method->IsCopied())) {
+      ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+      keep_alive = class_linker->GetHoldingClassLoaderOfCopiedMethod(self_, method);
+    } else {
+      keep_alive = method->GetDeclaringClass();
+    }
     trace_->Set</*kTransactionActive=*/ false, /*kCheckTransaction=*/ false>(
-        static_cast<int32_t>(count_) + 1, method->GetDeclaringClass());
+        static_cast<int32_t>(count_) + 1, keep_alive);
     ++count_;
   }
 
@@ -3000,11 +3010,14 @@
   uint32_t skip_depth_;
   // Current position down stack trace.
   uint32_t count_ = 0;
-  // An object array where the first element is a pointer array that contains the ArtMethod
-  // pointers on the stack and dex PCs. The rest of the elements are the declaring class of
-  // the ArtMethod pointers. trace_[i+1] contains the declaring class of the ArtMethod of the
-  // i'th frame. We're initializing a newly allocated trace, so we do not need to record that
-  // under a transaction. If the transaction is aborted, the whole trace shall be unreachable.
+  // An object array where the first element is a pointer array that contains the `ArtMethod`
+  // pointers on the stack and dex PCs. The rest of the elements are referencing objects
+  // that shall keep the methods alive, namely the declaring class of the `ArtMethod` for
+  // declared methods and the class loader for copied methods (because it's faster to find
+  // the class loader than the actual class that holds the copied method). The `trace_[i+1]`
+  // contains the declaring class or class loader of the `ArtMethod` of the i'th frame.
+  // We're initializing a newly allocated trace, so we do not need to record that under
+  // a transaction. If the transaction is aborted, the whole trace shall be unreachable.
   mirror::ObjectArray<mirror::Object>* trace_ = nullptr;
   // For cross compilation.
   const PointerSize pointer_size_;
diff --git a/test/141-class-unload/profile b/test/141-class-unload/profile
new file mode 100644
index 0000000..709f6ea
--- /dev/null
+++ b/test/141-class-unload/profile
@@ -0,0 +1,3 @@
+LMain;
+LIface;
+LImpl2;
diff --git a/test/141-class-unload/run.py b/test/141-class-unload/run.py
index c3fda2f..691d209 100644
--- a/test/141-class-unload/run.py
+++ b/test/141-class-unload/run.py
@@ -17,4 +17,4 @@
 
 def run(ctx, args):
   # Currently app images aren't unloaded when dex files are unloaded.
-  ctx.default_run(args, secondary_app_image=False)
+  ctx.default_run(args, profile=True, secondary_app_image=False)
diff --git a/test/141-class-unload/src-ex/Impl.java b/test/141-class-unload/src-ex/Impl.java
new file mode 100644
index 0000000..978b3de
--- /dev/null
+++ b/test/141-class-unload/src-ex/Impl.java
@@ -0,0 +1,18 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public class Impl implements Iface {
+}
diff --git a/test/141-class-unload/src/Iface.java b/test/141-class-unload/src/Iface.java
new file mode 100644
index 0000000..37b513f
--- /dev/null
+++ b/test/141-class-unload/src/Iface.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public interface Iface {
+    default void invokeRun(Runnable r) {
+        r.run();
+    }
+}
diff --git a/test/141-class-unload/src/Impl2.java b/test/141-class-unload/src/Impl2.java
new file mode 100644
index 0000000..8fd77e5
--- /dev/null
+++ b/test/141-class-unload/src/Impl2.java
@@ -0,0 +1,18 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public class Impl2 implements Iface {
+}
diff --git a/test/141-class-unload/src/Main.java b/test/141-class-unload/src/Main.java
index 8031dba..716f055 100644
--- a/test/141-class-unload/src/Main.java
+++ b/test/141-class-unload/src/Main.java
@@ -20,6 +20,8 @@
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.function.Consumer;
 
 public class Main {
     static final String DEX_FILE = System.getenv("DEX_LOCATION") + "/141-class-unload-ex.jar";
@@ -49,6 +51,14 @@
             testOatFilesUnloaded(getPid());
             // Test that objects keep class loader live for sticky GC.
             testStickyUnload(constructor);
+            // Test that copied methods recorded in a stack trace prevents unloading.
+            testCopiedMethodInStackTrace(constructor);
+            // Test that code preventing unloading holder classes of copied methods recorded in
+            // a stack trace does not crash when processing a copied method in the boot class path.
+            testCopiedBcpMethodInStackTrace();
+            // Test that code preventing unloading holder classes of copied methods recorded in
+            // a stack trace does not crash when processing a copied method in an app image.
+            testCopiedAppImageMethodInStackTrace();
         } catch (Exception e) {
             e.printStackTrace(System.out);
         }
@@ -103,13 +113,12 @@
         System.out.println(klass2.get());
     }
 
-    private static void testUnloadLoader(Constructor<?> constructor)
-        throws Exception {
-      WeakReference<ClassLoader> loader = setUpUnloadLoader(constructor, true);
-      // No strong references to class loader, should get unloaded.
-      doUnloading();
-      // If the weak reference is cleared, then it was unloaded.
-      System.out.println(loader.get());
+    private static void testUnloadLoader(Constructor<?> constructor) throws Exception {
+        WeakReference<ClassLoader> loader = setUpUnloadLoader(constructor, true);
+        // No strong references to class loader, should get unloaded.
+        doUnloading();
+        // If the weak reference is cleared, then it was unloaded.
+        System.out.println(loader.get());
     }
 
     private static void testStackTrace(Constructor<?> constructor) throws Exception {
@@ -138,20 +147,20 @@
     }
 
     static class Pair {
-      public Pair(Object o, ClassLoader l) {
-        object = o;
-        classLoader = new WeakReference<ClassLoader>(l);
-      }
+        public Pair(Object o, ClassLoader l) {
+            object = o;
+            classLoader = new WeakReference<ClassLoader>(l);
+        }
 
-      public Object object;
-      public WeakReference<ClassLoader> classLoader;
+        public Object object;
+        public WeakReference<ClassLoader> classLoader;
     }
 
     // Make the method not inline-able to prevent the compiler optimizing away the allocation.
     private static Pair $noinline$testNoUnloadInstanceHelper(Constructor<?> constructor)
             throws Exception {
         ClassLoader loader = (ClassLoader) constructor.newInstance(
-            DEX_FILE, LIBRARY_SEARCH_PATH, ClassLoader.getSystemClassLoader());
+                DEX_FILE, LIBRARY_SEARCH_PATH, ClassLoader.getSystemClassLoader());
         Object o = testNoUnloadHelper(loader);
         return new Pair(o, loader);
     }
@@ -165,7 +174,7 @@
 
     private static Class<?> setUpUnloadClass(Constructor<?> constructor) throws Exception {
         ClassLoader loader = (ClassLoader) constructor.newInstance(
-            DEX_FILE, LIBRARY_SEARCH_PATH, ClassLoader.getSystemClassLoader());
+                DEX_FILE, LIBRARY_SEARCH_PATH, ClassLoader.getSystemClassLoader());
         Class<?> intHolder = loader.loadClass("IntHolder");
         Method getValue = intHolder.getDeclaredMethod("getValue");
         Method setValue = intHolder.getDeclaredMethod("setValue", Integer.TYPE);
@@ -202,6 +211,75 @@
         System.out.println("Too small " + (s.length() < 1000));
     }
 
+    private static void assertStackTraceContains(Throwable t, String className, String methodName) {
+        boolean found = false;
+        for (StackTraceElement e : t.getStackTrace()) {
+            if (className.equals(e.getClassName()) && methodName.equals(e.getMethodName())) {
+                found = true;
+                break;
+            }
+        }
+        if (!found) {
+            throw new Error("Did not find " + className + "." + methodName);
+        }
+    }
+
+    private static void testCopiedMethodInStackTrace(Constructor<?> constructor) throws Exception {
+        Throwable t = $noinline$createStackTraceWithCopiedMethod(constructor);
+        doUnloading();
+        assertStackTraceContains(t, "Iface", "invokeRun");
+    }
+
+    private static Throwable $noinline$createStackTraceWithCopiedMethod(Constructor<?> constructor)
+            throws Exception {
+      ClassLoader loader = (ClassLoader) constructor.newInstance(
+              DEX_FILE, LIBRARY_SEARCH_PATH, Main.class.getClassLoader());
+      Iface impl = (Iface) loader.loadClass("Impl").newInstance();
+      Runnable throwingRunnable = new Runnable() {
+          public void run() {
+              throw new Error();
+          }
+      };
+      try {
+          impl.invokeRun(throwingRunnable);
+          System.out.println("UNREACHABLE");
+          return null;
+      } catch (Error expected) {
+          return expected;
+      }
+    }
+
+    private static void testCopiedBcpMethodInStackTrace() {
+        Consumer<Object> consumer = new Consumer<Object>() {
+            public void accept(Object o) {
+                throw new Error();
+            }
+        };
+        Error err = null;
+        try {
+            Arrays.asList(new Object[] { new Object() }).iterator().forEachRemaining(consumer);
+        } catch (Error expected) {
+            err = expected;
+        }
+        assertStackTraceContains(err, "Main", "testCopiedBcpMethodInStackTrace");
+    }
+
+    private static void testCopiedAppImageMethodInStackTrace() throws Exception {
+        Iface limpl = (Iface) Class.forName("Impl2").newInstance();
+        Runnable throwingRunnable = new Runnable() {
+            public void run() {
+                throw new Error();
+            }
+        };
+        Error err = null;
+        try {
+            limpl.invokeRun(throwingRunnable);
+        } catch (Error expected) {
+            err = expected;
+        }
+        assertStackTraceContains(err, "Main", "testCopiedAppImageMethodInStackTrace");
+    }
+
     private static WeakReference<Class> setUpUnloadClassWeak(Constructor<?> constructor)
             throws Exception {
         return new WeakReference<Class>(setUpUnloadClass(constructor));
@@ -242,7 +320,7 @@
     }
 
     private static int getPid() throws Exception {
-      return Integer.parseInt(new File("/proc/self").getCanonicalFile().getName());
+        return Integer.parseInt(new File("/proc/self").getCanonicalFile().getName());
     }
 
     public static native void stopJit();