From 1585f19982686fba9e5ae835a844ed302b975e87 Mon Sep 17 00:00:00 2001
From: ramindani <ramindani@google.com>
Date: Thu, 26 Jan 2023 13:42:19 -0800
Subject: [PATCH] Use NativeAllocationRegistry for DisplayEventReceiver

Test: ChoreographerTests, libsurfaceflinger_unittest, CtsGraphicsTestCases,
SurfaceFlinger_test, libcompositionengine_test
BUG: 266121235

Change-Id: I7167e29ddf48de707f5362942aa679f0b2d001ac
---
 .../android/view/DisplayEventReceiver.java    | 27 +++++++++----------
 core/java/android/view/SurfaceControl.java    |  1 -
 .../jni/android_view_DisplayEventReceiver.cpp | 11 +++++---
 3 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/core/java/android/view/DisplayEventReceiver.java b/core/java/android/view/DisplayEventReceiver.java
index 26fda3444c4ec..edce001813350 100644
--- a/core/java/android/view/DisplayEventReceiver.java
+++ b/core/java/android/view/DisplayEventReceiver.java
@@ -27,6 +27,8 @@ import com.android.internal.annotations.VisibleForTesting;
 
 import dalvik.annotation.optimization.FastNative;
 
+import libcore.util.NativeAllocationRegistry;
+
 import java.lang.ref.WeakReference;
 
 /**
@@ -81,11 +83,17 @@ public abstract class DisplayEventReceiver {
 
     private static native long nativeInit(WeakReference<DisplayEventReceiver> receiver,
             MessageQueue messageQueue, int vsyncSource, int eventRegistration, long layerHandle);
-    private static native void nativeDispose(long receiverPtr);
+    private static native long nativeGetDisplayEventReceiverFinalizer();
     @FastNative
     private static native void nativeScheduleVsync(long receiverPtr);
     private static native VsyncEventData nativeGetLatestVsyncEventData(long receiverPtr);
 
+    private static final NativeAllocationRegistry sNativeAllocationRegistry =
+            NativeAllocationRegistry.createMalloced(
+                    DisplayEventReceiver.class.getClassLoader(),
+                    nativeGetDisplayEventReceiverFinalizer());
+    private Runnable mFreeNativeResources;
+
     /**
      * Creates a display event receiver.
      *
@@ -118,27 +126,16 @@ public abstract class DisplayEventReceiver {
         mMessageQueue = looper.getQueue();
         mReceiverPtr = nativeInit(new WeakReference<DisplayEventReceiver>(this), mMessageQueue,
                 vsyncSource, eventRegistration, layerHandle);
-    }
-
-    @Override
-    protected void finalize() throws Throwable {
-        try {
-            dispose(true);
-        } finally {
-            super.finalize();
-        }
+        mFreeNativeResources = sNativeAllocationRegistry.registerNativeAllocation(this,
+                mReceiverPtr);
     }
 
     /**
      * Disposes the receiver.
      */
     public void dispose() {
-        dispose(false);
-    }
-
-    private void dispose(boolean finalized) {
         if (mReceiverPtr != 0) {
-            nativeDispose(mReceiverPtr);
+            mFreeNativeResources.run();
             mReceiverPtr = 0;
         }
         mMessageQueue = null;
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index b003659affd32..4ff86d33ea491 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -1386,7 +1386,6 @@ public final class SurfaceControl implements Parcelable {
             synchronized (mChoreographerLock) {
                 if (mChoreographer != null) {
                     mChoreographer.invalidate();
-                    // TODO(b/266121235): Use NativeAllocationRegistry to clean up Choreographer.
                     mChoreographer = null;
                 }
             }
diff --git a/core/jni/android_view_DisplayEventReceiver.cpp b/core/jni/android_view_DisplayEventReceiver.cpp
index 8855b7806cc18..b09a9c3a505de 100644
--- a/core/jni/android_view_DisplayEventReceiver.cpp
+++ b/core/jni/android_view_DisplayEventReceiver.cpp
@@ -240,13 +240,15 @@ static jlong nativeInit(JNIEnv* env, jclass clazz, jobject receiverWeak, jobject
     return reinterpret_cast<jlong>(receiver.get());
 }
 
-static void nativeDispose(JNIEnv* env, jclass clazz, jlong receiverPtr) {
-    NativeDisplayEventReceiver* receiver =
-            reinterpret_cast<NativeDisplayEventReceiver*>(receiverPtr);
+static void release(NativeDisplayEventReceiver* receiver) {
     receiver->dispose();
     receiver->decStrong(gDisplayEventReceiverClassInfo.clazz); // drop reference held by the object
 }
 
+static jlong nativeGetDisplayEventReceiverFinalizer(JNIEnv*, jclass) {
+    return static_cast<jlong>(reinterpret_cast<uintptr_t>(&release));
+}
+
 static void nativeScheduleVsync(JNIEnv* env, jclass clazz, jlong receiverPtr) {
     sp<NativeDisplayEventReceiver> receiver =
             reinterpret_cast<NativeDisplayEventReceiver*>(receiverPtr);
@@ -274,7 +276,8 @@ static const JNINativeMethod gMethods[] = {
         /* name, signature, funcPtr */
         {"nativeInit", "(Ljava/lang/ref/WeakReference;Landroid/os/MessageQueue;IIJ)J",
          (void*)nativeInit},
-        {"nativeDispose", "(J)V", (void*)nativeDispose},
+        {"nativeGetDisplayEventReceiverFinalizer", "()J",
+         (void*)nativeGetDisplayEventReceiverFinalizer},
         // @FastNative
         {"nativeScheduleVsync", "(J)V", (void*)nativeScheduleVsync},
         {"nativeGetLatestVsyncEventData", "(J)Landroid/view/DisplayEventReceiver$VsyncEventData;",