SoundPool: Fix WeakRef lifetime

Ensure WeakRef lifetime is extended to duration of callback.

Test: atest SoundPoolAacTest
Test: atest SoundPoolHapticTest
Test: atest SoundPoolMidiTest
Test: atest SoundPoolOggTest
Test: atest AudioManagerTest#testSoundEffects
Bug: 223815163
Change-Id: I235ef32532d919823f1bbc4d7c90ec47311cc7ae

media/jni/soundpool/android_media_SoundPool.cpp

@@ -181,6 +181,116 @@ inline auto setSoundPool(
     return getSoundPoolManager().set(env, thiz, soundPool);
 }
 
+/**
+ * ConcurrentHashMap is a locked hash map
+ *
+ * As from the name, this class is thread_safe.
+ *
+ * The type V must have 3 properties in the current implementation.
+ *    1) A V{} default constructor which represents a nullValue.
+ *    2) V::operator bool() const efficient detection of such a nullValue.
+ *    3) V must be copyable.
+ *
+ * Note: The Key cannot be a Java LocalRef, as those change between JNI calls.
+ * The Key could be the raw native object pointer if one wanted to associate
+ * extra data with a native object.
+ *
+ * Using set() with a nullValue V results in erasing the key entry.
+ * A nullValue V is returned by get() if there is no underlying entry.
+ *
+ * Design notes:
+ * 1) For objects of type V that do not naturally have a "nullValue",
+ *    wrapping VOpt = std::optional<V> or VShared = std::shared<V> is recommended.
+ *
+ * 2) An overload for an explicit equality comparable nullValue such as
+ *    get(..., const V& nullValue) or set(..., const V& nullValue)
+ *    is omitted.  An alternative is to pass a fixed nullValue into a special
+ *    constructor (omitted) for equality comparisons and return value.
+ *
+ * 3) This ConcurrentHashMap currently allows only one thread at a time.
+ *    It is not optimized for heavy multi-threaded use.
+ */
+template <typename K, typename V>
+class ConcurrentHashMap
+{
+public:
+
+    // Sets the value and returns the old one.
+    //
+    // If the old value doesn't exist, then nullValue V is returned.
+    // If the new value is false by operator bool(), the internal value is destroyed.
+    // Note: The old value is returned so if V is a smart pointer, it can be held
+    // by the caller to be deleted outside of any external lock.
+
+    V set(const K& key, const V& value) {
+        std::lock_guard lg(mLock);
+        auto it = mMap.find(key);
+        if (it == mMap.end()) {
+            if (value) {
+                mMap[key] = value;
+            }
+            return {};
+        }
+        V oldValue = std::move(it->second);
+        if (value) {
+            it->second = value;
+        } else {
+            mMap.erase(it);
+        }
+        return oldValue;
+    }
+
+    // Retrieves the associated object, returns nullValue V if not available.
+    V get(const K& key) const {
+        std::lock_guard lg(mLock);
+        auto it = mMap.find(key);
+        return it != mMap.end() ? it->second : V{};
+    }
+private:
+    mutable std::mutex mLock;
+    std::unordered_map<K, V> mMap GUARDED_BY(mLock);
+};
+
+// *jobject is needed to fit the jobject into a std::shared_ptr.
+// This is the Android type _jobject, but we derive this type as JObjectValue.
+using JObjectValue = std::remove_pointer_t<jobject>; // _jobject
+
+// Check that jobject is really a pointer to JObjectValue.
+// The JNI contract is that jobject is NULL comparable,
+// so jobject is pointer equivalent; we check here to be sure.
+// Note std::remove_ptr_t<NonPointerType> == NonPointerType.
+static_assert(std::is_same_v<JObjectValue*, jobject>);
+
+// We store the ancillary data associated with a SoundPool object in a concurrent
+// hash map indexed on the SoundPool native object pointer.
+auto& getSoundPoolJavaRefManager() {
+    static ConcurrentHashMap<SoundPool *, std::shared_ptr<JObjectValue>> concurrentHashMap;
+    return concurrentHashMap;
+}
+
+// make_shared_globalref_from_localref() creates a sharable Java global
+// reference from a Java local reference. The equivalent type is
+// std::shared_ptr<_jobject> (where _jobject is JObjectValue,
+// and _jobject * is jobject),
+// and the jobject may be retrieved by .get() or pointer dereference.
+// This encapsulation gives the benefit of std::shared_ptr
+// ref counting, weak_ptr, etc.
+//
+// The Java global reference should be stable between JNI calls.  It is a limited
+// quantity so sparingly use global references.
+//
+// The Android JNI implementation is described here:
+// https://developer.android.com/training/articles/perf-jni
+// https://android-developers.googleblog.com/2011/11/jni-local-reference-changes-in-ics.html
+//
+inline auto make_shared_globalref_from_localref(JNIEnv *env, jobject localRef) {
+    return std::shared_ptr<JObjectValue>(
+            env->NewGlobalRef(localRef),
+            [](JObjectValue* object) { // cannot cache env as don't know which thread we're on.
+                if (object != nullptr) AndroidRuntime::getJNIEnv()->DeleteGlobalRef(object);
+            });
+}
+
 } // namespace
 
 static const char* const kAudioAttributesClassPathName = "android/media/AudioAttributes";
@@ -321,10 +431,22 @@ android_media_SoundPool_setRate(JNIEnv *env, jobject thiz, jint channelID,
 static void android_media_callback(SoundPoolEvent event, SoundPool* soundPool, void* user)
 {
     ALOGV("callback: (%d, %d, %d, %p, %p)", event.mMsg, event.mArg1, event.mArg2, soundPool, user);
+    auto weakRef = getSoundPoolJavaRefManager().get(soundPool); // shared_ptr to WeakRef
+    if (weakRef == nullptr) {
+        ALOGD("%s: no weak ref, object released, ignoring callback", __func__);
+        return;
+    }
     JNIEnv *env = AndroidRuntime::getJNIEnv();
     env->CallStaticVoidMethod(
-            fields.mSoundPoolClass, fields.mPostEvent, user, event.mMsg, event.mArg1, event.mArg2,
+            fields.mSoundPoolClass, fields.mPostEvent,
+            weakRef.get(), event.mMsg, event.mArg1, event.mArg2,
             nullptr /* object */);
+
+    if (env->ExceptionCheck() != JNI_FALSE) {
+        ALOGE("%s: Uncaught exception returned from Java callback", __func__);
+        env->ExceptionDescribe();
+        env->ExceptionClear(); // Just clear it, hopefully all is ok.
+    }
 }
 
 static jint
@@ -353,13 +475,15 @@ android_media_SoundPool_native_setup(JNIEnv *env, jobject thiz, jobject weakRef,
     ALOGV("android_media_SoundPool_native_setup");
     ScopedUtfChars opPackageNameStr(env, opPackageName);
     auto soundPool = std::make_shared<SoundPool>(maxChannels, paa, opPackageNameStr.c_str());
-
-    // set callback with weak reference
-    jobject globalWeakRef = env->NewGlobalRef(weakRef);
-    soundPool->setCallback(android_media_callback, globalWeakRef);
+    soundPool->setCallback(android_media_callback, nullptr /* user */);
 
     // register with SoundPoolManager.
+
     auto oldSoundPool = setSoundPool(env, thiz, soundPool);
+    // register Java SoundPool WeakRef using native SoundPool * as the key, for the callback.
+    auto oldSoundPoolJavaRef = getSoundPoolJavaRefManager().set(
+            soundPool.get(), make_shared_globalref_from_localref(env, weakRef));
+
     ALOGW_IF(oldSoundPool != nullptr, "%s: Aliased SoundPool object %p",
             __func__, oldSoundPool.get());
 
@@ -375,19 +499,12 @@ android_media_SoundPool_release(JNIEnv *env, jobject thiz)
     ALOGV("android_media_SoundPool_release");
 
     // Remove us from SoundPoolManager.
-    auto oldSoundPool = setSoundPool(env, thiz, nullptr);
 
-    // Caution: Deleting the weakRef is not race free from invoking
-    // the Java callback because we may not have the last remaining
-    // reference to the SoundPool object - another method could still
-    // be in progress.
+    auto oldSoundPool = setSoundPool(env, thiz, nullptr);
     if (oldSoundPool != nullptr) {
-        // release weak reference and clear callback
-        auto weakRef = (jobject) oldSoundPool->getUserData();
-        oldSoundPool->setCallback(nullptr /* callback */, nullptr /* user */);
-        if (weakRef != nullptr) {
-            env->DeleteGlobalRef(weakRef);
-        }
+        // Note: setting the weak ref is thread safe in case there is a callback
+        // simultaneously occurring.
+        auto oldSoundPoolJavaRef = getSoundPoolJavaRefManager().set(oldSoundPool.get(), nullptr);
     }
     // destructor to oldSoundPool should occur at exit.
 }