Migrate ApkAssets to use NativeAllocationRegistry

This change ensures that mNativePtr can be deleted only when its
ApkAssets is really inaccessible and prevents finalization race
conditions.

Bug: 159041693
Test: boot and add logs to ensure destructor runs
Change-Id: Iff0612fe6862a17b8c889d9e7ba230094e3cb14a
diff --git a/core/java/android/content/res/ApkAssets.java b/core/java/android/content/res/ApkAssets.java
index f5e1719..2d381eb 100644
--- a/core/java/android/content/res/ApkAssets.java
+++ b/core/java/android/content/res/ApkAssets.java
@@ -25,6 +25,8 @@
 
 import com.android.internal.annotations.GuardedBy;
 
+import libcore.util.NativeAllocationRegistry;
+
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.lang.annotation.Retention;
@@ -101,11 +103,11 @@
     public @interface FormatType {}
 
     @GuardedBy("this")
-    private long mNativePtr;  // final, except cleared in finalizer.
+    private final long mNativePtr;
 
     @Nullable
     @GuardedBy("this")
-    private final StringBlock mStringBlock;  // null or closed if mNativePtr = 0.
+    private final StringBlock mStringBlock;
 
     @PropertyFlags
     private final int mFlags;
@@ -113,6 +115,19 @@
     @Nullable
     private final AssetsProvider mAssets;
 
+    @GuardedBy("this")
+    @Nullable
+    private final Runnable mRunNativeCleanup;
+
+    // Use a Holder to allow static initialization of ApkAssets in the boot image, and
+    // possibly to avoid some initialization ordering issues.
+    private static class NoImagePreloadHolder {
+        // TODO(175425996): Make size estimate more accurate
+        public static final NativeAllocationRegistry REGISTRY =
+                NativeAllocationRegistry.createMalloced(ApkAssets.class.getClassLoader(),
+                                                        nativeGetFinalizer());
+    }
+
     /**
      * Creates a new ApkAssets instance from the given path on disk.
      *
@@ -287,6 +302,8 @@
         mFlags = flags;
         mNativePtr = nativeLoad(format, path, flags, assets);
         mStringBlock = new StringBlock(nativeGetStringBlock(mNativePtr), true /*useSparse*/);
+        mRunNativeCleanup = NoImagePreloadHolder.REGISTRY.registerNativeAllocation(
+                this, mNativePtr);
         mAssets = assets;
     }
 
@@ -298,6 +315,8 @@
         mFlags = flags;
         mNativePtr = nativeLoadFd(format, fd, friendlyName, flags, assets);
         mStringBlock = new StringBlock(nativeGetStringBlock(mNativePtr), true /*useSparse*/);
+        mRunNativeCleanup = NoImagePreloadHolder.REGISTRY.registerNativeAllocation(
+                this, mNativePtr);
         mAssets = assets;
     }
 
@@ -309,6 +328,8 @@
         mFlags = flags;
         mNativePtr = nativeLoadFdOffsets(format, fd, friendlyName, offset, length, flags, assets);
         mStringBlock = new StringBlock(nativeGetStringBlock(mNativePtr), true /*useSparse*/);
+        mRunNativeCleanup = NoImagePreloadHolder.REGISTRY.registerNativeAllocation(
+                this, mNativePtr);
         mAssets = assets;
     }
 
@@ -316,6 +337,7 @@
         mFlags = flags;
         mNativePtr = nativeLoadEmpty(flags, assets);
         mStringBlock = null;
+        mRunNativeCleanup = null;
         mAssets = assets;
     }
 
@@ -403,22 +425,16 @@
         return "ApkAssets{path=" + getAssetPath() + "}";
     }
 
-    @Override
-    protected void finalize() throws Throwable {
-        close();
-    }
-
     /**
      * Closes this class and the contained {@link #mStringBlock}.
      */
     public void close() {
         synchronized (this) {
-            if (mNativePtr != 0) {
-                if (mStringBlock != null) {
-                    mStringBlock.close();
-                }
-                nativeDestroy(mNativePtr);
-                mNativePtr = 0;
+            if (mStringBlock != null) {
+                mStringBlock.close();
+            }
+            if (mRunNativeCleanup != null) {
+                mRunNativeCleanup.run();
             }
         }
     }
@@ -433,7 +449,6 @@
     private static native long nativeLoadFdOffsets(@FormatType int format,
             @NonNull FileDescriptor fd, @NonNull String friendlyName, long offset, long length,
             @PropertyFlags int flags, @Nullable AssetsProvider asset) throws IOException;
-    private static native void nativeDestroy(long ptr);
     private static native @NonNull String nativeGetAssetPath(long ptr);
     private static native long nativeGetStringBlock(long ptr);
     private static native boolean nativeIsUpToDate(long ptr);
@@ -441,4 +456,5 @@
     private static native @Nullable OverlayableInfo nativeGetOverlayableInfo(long ptr,
             String overlayableName) throws IOException;
     private static native boolean nativeDefinesOverlayable(long ptr) throws IOException;
+    private static native final long nativeGetFinalizer();
 }
diff --git a/core/jni/android_content_res_ApkAssets.cpp b/core/jni/android_content_res_ApkAssets.cpp
index 444bb66..66753e4 100644
--- a/core/jni/android_content_res_ApkAssets.cpp
+++ b/core/jni/android_content_res_ApkAssets.cpp
@@ -315,10 +315,14 @@
   return reinterpret_cast<jlong>(apk_assets.release());
 }
 
-static void NativeDestroy(JNIEnv* /*env*/, jclass /*clazz*/, jlong ptr) {
+static void NativeDestroy(void* ptr) {
   delete reinterpret_cast<ApkAssets*>(ptr);
 }
 
+static jlong NativeGetFinalizer(JNIEnv* /*env*/, jclass /*clazz*/) {
+  return reinterpret_cast<jlong>(&NativeDestroy);
+}
+
 static jstring NativeGetAssetPath(JNIEnv* env, jclass /*clazz*/, jlong ptr) {
   const ApkAssets* apk_assets = reinterpret_cast<const ApkAssets*>(ptr);
   return env->NewStringUTF(apk_assets->GetPath().c_str());
@@ -427,7 +431,7 @@
     {"nativeLoadFdOffsets",
      "(ILjava/io/FileDescriptor;Ljava/lang/String;JJILandroid/content/res/loader/AssetsProvider;)J",
      (void*)NativeLoadFromFdOffset},
-    {"nativeDestroy", "(J)V", (void*)NativeDestroy},
+    {"nativeGetFinalizer", "()J", (void*)NativeGetFinalizer},
     {"nativeGetAssetPath", "(J)Ljava/lang/String;", (void*)NativeGetAssetPath},
     {"nativeGetStringBlock", "(J)J", (void*)NativeGetStringBlock},
     {"nativeIsUpToDate", "(J)Z", (void*)NativeIsUpToDate},