[Ravenwood] Support NativeAllocationRegistry

We don't track native allocation size, but we still need to
release native allocations.

Bug: 337110712
Test: atest RavenwoodBivalentTest
Test: atest RavenwoodBivalentTest_device
Test: ./ravenwood/scripts/run-ravenwood-tests.sh
Change-Id: Ia50c963731a26fc951a8040cbf353a5c56505b6a

13 files changed, 400 insertions, 3 deletions

diff --git a/Ravenwood.bp b/Ravenwood.bp
index 74382a6e8d44..3ab0934e9acc 100644
--- a/Ravenwood.bp
+++ b/Ravenwood.bp
@@ -270,6 +270,7 @@ android_ravenwood_libgroup {
     ],
     jni_libs: [
         "libandroid_runtime",
+        "libravenwood_runtime",
     ],
 }
 
diff --git a/ravenwood/Android.bp b/ravenwood/Android.bp
index 33374195e674..e06f40070a0a 100644
--- a/ravenwood/Android.bp
+++ b/ravenwood/Android.bp
@@ -155,6 +155,31 @@ filegroup {
     visibility: ["//frameworks/base"],
 }
 
+cc_library_shared {
+    name: "libravenwood_runtime",
+    host_supported: true,
+
+    cflags: [
+        "-Wall",
+        "-Werror",
+        "-Wno-unused-parameter",
+        "-Wthread-safety",
+    ],
+
+    srcs: [
+        "runtime-helper-src/jni/*.cpp",
+    ],
+
+    shared_libs: [
+        "libbase",
+        "liblog",
+        "libnativehelper",
+        "libutils",
+        "libcutils",
+    ],
+    visibility: ["//frameworks/base"],
+}
+
 // For collecting the *stats.csv files in a known directory under out/host/linux-x86/testcases/.
 // The "test" just shows the available stats filenames.
 sh_test_host {
diff --git a/ravenwood/bivalenttest/Android.bp b/ravenwood/bivalenttest/Android.bp
index a6b6ed934ce7..2d94894c54be 100644
--- a/ravenwood/bivalenttest/Android.bp
+++ b/ravenwood/bivalenttest/Android.bp
@@ -45,7 +45,6 @@ android_ravenwood_test {
     jni_libs: [
         "libravenwoodbivalenttest_jni",
     ],
-    sdk_version: "test_current",
     auto_gen_config: true,
 }
 
diff --git a/ravenwood/bivalenttest/AndroidTest.xml b/ravenwood/bivalenttest/AndroidTest.xml
index ac4695bb9a50..9e5dd11b60cb 100644
--- a/ravenwood/bivalenttest/AndroidTest.xml
+++ b/ravenwood/bivalenttest/AndroidTest.xml
@@ -25,5 +25,8 @@
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.ravenwood.bivalenttest" />
         <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+
+        <!-- Need to use MODULE_LIBRARIES APIs, so no hidden API check. -->
+        <option name="hidden-api-checks" value="false" />
     </test>
 </configuration>
diff --git a/ravenwood/bivalenttest/jni/ravenwood_core_test_jni.cpp b/ravenwood/bivalenttest/jni/ravenwood_core_test_jni.cpp
index 83f756e9a0bf..956d79c0d25f 100644
--- a/ravenwood/bivalenttest/jni/ravenwood_core_test_jni.cpp
+++ b/ravenwood/bivalenttest/jni/ravenwood_core_test_jni.cpp
@@ -15,19 +15,70 @@
  */
 
 #include <nativehelper/JNIHelp.h>
+#include <atomic>
 #include "jni.h"
 #include "utils/Log.h"
 #include "utils/misc.h"
 
+// JNI methods for RavenwoodJniTest
+
 static jint add(JNIEnv* env, jclass clazz, jint a, jint b) {
     return a + b;
 }
 
-static const JNINativeMethod sMethods[] =
+static const JNINativeMethod sMethods_JniTest[] =
 {
     { "add", "(II)I", (void*)add },
 };
 
+// JNI methods for RavenwoodNativeAllocationRegistryTest
+std::atomic<int> numTotalAlloc = 0;
+
+class NarTestData {
+public:
+    NarTestData(jint v): value(v) {
+        numTotalAlloc++;
+    }
+
+    ~NarTestData() {
+        value = -1;
+        numTotalAlloc--;
+    }
+
+    volatile jint value;
+};
+
+static jlong NarTestData_nMalloc(JNIEnv* env, jclass clazz, jint value) {
+    NarTestData* p = new NarTestData(value);
+    return reinterpret_cast<jlong>(p);
+}
+
+static jint NarTestData_nGet(JNIEnv* env, jclass clazz, jlong ptr) {
+    NarTestData* p = reinterpret_cast<NarTestData*>(ptr);
+    return p->value;
+}
+
+static void NarTestData_free(jlong ptr) {
+    NarTestData* p = reinterpret_cast<NarTestData*>(ptr);
+    delete p;
+}
+
+static jlong NarTestData_nGetNativeFinalizer(JNIEnv* env, jclass clazz) {
+    return reinterpret_cast<jlong>(NarTestData_free);
+}
+
+static jint NarTestData_nGetTotalAlloc(JNIEnv* env, jclass clazz) {
+    return numTotalAlloc;
+}
+
+static const JNINativeMethod sMethods_NarTestData[] =
+{
+    { "nMalloc", "(I)J", (void*)NarTestData_nMalloc },
+    { "nGet", "(J)I", (void*)NarTestData_nGet },
+    { "nGetNativeFinalizer", "()J", (void*)NarTestData_nGetNativeFinalizer },
+    { "nGetTotalAlloc", "()I", (void*)NarTestData_nGetTotalAlloc },
+};
+
 extern "C" jint JNI_OnLoad(JavaVM* vm, void* /* reserved */)
 {
     JNIEnv* env = NULL;
@@ -43,7 +94,13 @@ extern "C" jint JNI_OnLoad(JavaVM* vm, void* /* reserved */)
 
     int res = jniRegisterNativeMethods(env,
             "com/android/ravenwoodtest/bivalenttest/RavenwoodJniTest",
-            sMethods, NELEM(sMethods));
+            sMethods_JniTest, NELEM(sMethods_JniTest));
+    if (res < 0) {
+        return res;
+    }
+    res = jniRegisterNativeMethods(env,
+            "com/android/ravenwoodtest/bivalenttest/RavenwoodNativeAllocationRegistryTest$Data",
+            sMethods_NarTestData, NELEM(sMethods_NarTestData));
     if (res < 0) {
         return res;
     }
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodJniTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodJniTest.java
index 59467e9d7d14..0cc2adc6b26b 100644
--- a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodJniTest.java
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodJniTest.java
@@ -29,6 +29,10 @@ import org.junit.runner.RunWith;
 @RunWith(AndroidJUnit4.class)
 public final class RavenwoodJniTest {
     static {
+        initializeJni();
+    }
+
+    public static void initializeJni() {
         RavenwoodUtils.loadJniLibrary("ravenwoodbivalenttest_jni");
     }
 
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodNativeAllocationRegistryTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodNativeAllocationRegistryTest.java
new file mode 100644
index 000000000000..415b4676f46b
--- /dev/null
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodNativeAllocationRegistryTest.java
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+package com.android.ravenwoodtest.bivalenttest;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.Log;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import libcore.util.NativeAllocationRegistry;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+public class RavenwoodNativeAllocationRegistryTest {
+    private static final String TAG = RavenwoodNativeAllocationRegistryTest.class.getSimpleName();
+    static {
+        RavenwoodJniTest.initializeJni();
+    }
+
+    @Rule
+    public final RavenwoodRule mRavenwoodRule = new RavenwoodRule();
+
+    private static class Data {
+        private final long mNativePtr;
+
+        private static native long nMalloc(int value);
+        private static native int nGet(long ptr);
+        private static native long nGetNativeFinalizer();
+
+        public static native int nGetTotalAlloc();
+
+        public int get() {
+            return nGet(mNativePtr);
+        }
+
+        private static class NarHolder {
+            public static final NativeAllocationRegistry sRegistry =
+                    NativeAllocationRegistry.createMalloced(
+                            Data.class.getClassLoader(), nGetNativeFinalizer());
+        }
+
+        public Data(int value) {
+            mNativePtr = nMalloc(value);
+            NarHolder.sRegistry.registerNativeAllocation(this, mNativePtr);
+        }
+    }
+
+    @Test
+    public void testNativeAllocationRegistry() {
+
+        final long timeoutTime = mRavenwoodRule.realCurrentTimeMillis() + 10_000;
+
+        final int startAlloc = Data.nGetTotalAlloc();
+
+        int totalAlloc = 0;
+
+        // Keep allocation new objects, until some get released.
+
+        while (true) {
+            for (int i = 0; i < 1000; i++) {
+                totalAlloc++;
+                Data d = new Data(i);
+                assertEquals(i, d.get());
+            }
+            System.gc();
+
+            final int currentAlloc = Data.nGetTotalAlloc() - startAlloc;
+            Log.i(TAG, "# of currently allocated objects=" + currentAlloc);
+
+            if (currentAlloc < totalAlloc) {
+                break; // Good, some objects have been released;
+            }
+            if (mRavenwoodRule.realCurrentTimeMillis() > timeoutTime) {
+                fail("No objects have been released before timeout");
+            }
+        }
+    }
+}
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
index 21d8019fcd87..9d12f855c0ea 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
@@ -426,4 +426,15 @@ public class RavenwoodRule implements TestRule {
             return ENABLE_OPTIONAL_VALIDATION;
         }
     }
+
+    /**
+     * Returns the "real" result from {@link System#currentTimeMillis()}.
+     *
+     * Currently, it's the same thing as calling {@link System#currentTimeMillis()},
+     * but this one is guaranteeed to return the real value, even when Ravenwood supports
+     * injecting a time to{@link System#currentTimeMillis()}.
+     */
+    public long realCurrentTimeMillis() {
+        return System.currentTimeMillis();
+    }
 }
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodUtils.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodUtils.java
index 37ceac601e37..99ab32788235 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodUtils.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodUtils.java
@@ -26,6 +26,15 @@ public class RavenwoodUtils {
     private RavenwoodUtils() {
     }
 
+    private static final String RAVENWOOD_NATIVE_RUNTIME_NAME = "ravenwood_runtime";
+
+    // LibcoreRavenwoodUtils calls it with reflections.
+    public static void loadRavenwoodNativeRuntime() {
+        if (RavenwoodRule.isOnRavenwood()) {
+            RavenwoodUtils.loadJniLibrary(RAVENWOOD_NATIVE_RUNTIME_NAME);
+        }
+    }
+
     /**
      * Load a JNI library respecting {@code java.library.path}
      * (which reflects {@code LD_LIBRARY_PATH}).
diff --git a/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
index e951351bd829..773a89a3df4d 100644
--- a/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
+++ b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
@@ -39,4 +39,8 @@ public class RavenwoodRuleImpl {
     public static void validate(Statement base, Description description,
             boolean enableOptionalValidation) {
     }
+
+    public static long realCurrentTimeMillis() {
+        return System.currentTimeMillis();
+    }
 }
diff --git a/ravenwood/runtime-helper-src/jni/ravenwood_runtime.cpp b/ravenwood/runtime-helper-src/jni/ravenwood_runtime.cpp
new file mode 100644
index 000000000000..8e3a21dd6d87
--- /dev/null
+++ b/ravenwood/runtime-helper-src/jni/ravenwood_runtime.cpp
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#include <nativehelper/JNIHelp.h>
+#include "jni.h"
+#include "utils/Log.h"
+#include "utils/misc.h"
+
+
+typedef void (*FreeFunction)(void*);
+
+static void NativeAllocationRegistry_applyFreeFunction(JNIEnv*,
+                                                       jclass,
+                                                       jlong freeFunction,
+                                                       jlong ptr) {
+    void* nativePtr = reinterpret_cast<void*>(static_cast<uintptr_t>(ptr));
+    FreeFunction nativeFreeFunction
+        = reinterpret_cast<FreeFunction>(static_cast<uintptr_t>(freeFunction));
+    nativeFreeFunction(nativePtr);
+}
+
+static const JNINativeMethod sMethods_NAR[] =
+{
+    { "applyFreeFunction", "(JJ)V", (void*)NativeAllocationRegistry_applyFreeFunction },
+};
+
+extern "C" jint JNI_OnLoad(JavaVM* vm, void* /* reserved */)
+{
+    JNIEnv* env = NULL;
+    jint result = -1;
+
+    if (vm->GetEnv((void**) &env, JNI_VERSION_1_4) != JNI_OK) {
+        ALOGE("GetEnv failed!");
+        return result;
+    }
+    ALOG_ASSERT(env, "Could not retrieve the env!");
+
+    ALOGI("%s: JNI_OnLoad", __FILE__);
+
+    // Initialize the Ravenwood version of NativeAllocationRegistry.
+    // We don't use this JNI on the device side, but if we ever have to do, skip this part.
+#ifndef __ANDROID__
+    int res = jniRegisterNativeMethods(env, "libcore/util/NativeAllocationRegistry",
+            sMethods_NAR, NELEM(sMethods_NAR));
+    if (res < 0) {
+        return res;
+    }
+#endif
+
+    return JNI_VERSION_1_4;
+}
diff --git a/ravenwood/runtime-helper-src/libcore-fake/libcore/ravenwood/LibcoreRavenwoodUtils.java b/ravenwood/runtime-helper-src/libcore-fake/libcore/ravenwood/LibcoreRavenwoodUtils.java
new file mode 100644
index 000000000000..839b62a39471
--- /dev/null
+++ b/ravenwood/runtime-helper-src/libcore-fake/libcore/ravenwood/LibcoreRavenwoodUtils.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+package libcore.ravenwood;
+
+public class LibcoreRavenwoodUtils {
+    private LibcoreRavenwoodUtils() {
+    }
+
+    public static void loadRavenwoodNativeRuntime() {
+        // TODO Stop using reflections.
+        // We need to call RavenwoodUtils.loadRavenwoodNativeRuntime(), but due to the build
+        // structure complexity, we can't refer to to this method directly from here,
+        // so let's use reflections for now...
+        try {
+            final var clazz = Class.forName("android.platform.test.ravenwood.RavenwoodUtils");
+            final var method = clazz.getMethod("loadRavenwoodNativeRuntime");
+            method.invoke(null);
+        } catch (Throwable th) {
+            throw new IllegalStateException("Failed to load Ravenwood native runtime", th);
+        }
+    }
+}
diff --git a/ravenwood/runtime-helper-src/libcore-fake/libcore/util/NativeAllocationRegistry.java b/ravenwood/runtime-helper-src/libcore-fake/libcore/util/NativeAllocationRegistry.java
new file mode 100644
index 000000000000..93861e87318f
--- /dev/null
+++ b/ravenwood/runtime-helper-src/libcore-fake/libcore/util/NativeAllocationRegistry.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+package libcore.util;
+
+import libcore.ravenwood.LibcoreRavenwoodUtils;
+
+import java.lang.ref.Cleaner;
+import java.lang.ref.Reference;
+
+/**
+ * Re-implementation of ART's NativeAllocationRegistry for Ravenwood.
+ * - We don't track the native allocation size on Ravenwood.
+ * - sun.misc.Cleaner isn't available on the desktop JVM, so we use java.lang.ref.Cleaner.
+ *   (Should ART switch to java.lang.ref.Cleaner?)
+ */
+public class NativeAllocationRegistry {
+    static {
+        // Initialize the JNI method.
+        LibcoreRavenwoodUtils.loadRavenwoodNativeRuntime();
+    }
+
+    private final long mFreeFunction;
+    private static final Cleaner sCleaner = Cleaner.create();
+
+    public static NativeAllocationRegistry createNonmalloced(
+            ClassLoader classLoader, long freeFunction, long size) {
+        return new NativeAllocationRegistry(classLoader, freeFunction, size, false);
+    }
+
+    public static NativeAllocationRegistry createMalloced(
+            ClassLoader classLoader, long freeFunction, long size) {
+        return new NativeAllocationRegistry(classLoader, freeFunction, size, true);
+    }
+
+    public static NativeAllocationRegistry createMalloced(
+            ClassLoader classLoader, long freeFunction) {
+        return new NativeAllocationRegistry(classLoader, freeFunction, 0, true);
+    }
+
+    public NativeAllocationRegistry(ClassLoader classLoader, long freeFunction, long size) {
+        this(classLoader, freeFunction, size, size == 0);
+    }
+
+    private NativeAllocationRegistry(ClassLoader classLoader, long freeFunction, long size,
+            boolean mallocAllocation) {
+        if (size < 0) {
+            throw new IllegalArgumentException("Invalid native allocation size: " + size);
+        }
+        mFreeFunction = freeFunction;
+    }
+
+    public Runnable registerNativeAllocation(Object referent, long nativePtr) {
+        if (referent == null) {
+            throw new IllegalArgumentException("referent is null");
+        }
+        if (nativePtr == 0) {
+            throw new IllegalArgumentException("nativePtr is null");
+        }
+
+        final Runnable releaser = () -> {
+            applyFreeFunction(mFreeFunction, nativePtr);
+        };
+        sCleaner.register(referent, releaser);
+
+        // Ensure that cleaner doesn't get invoked before we enable it.
+        Reference.reachabilityFence(referent);
+        return releaser;
+    }
+
+    /**
+     * Calls {@code freeFunction}({@code nativePtr}).
+     */
+    public static native void applyFreeFunction(long freeFunction, long nativePtr);
+}
+