From 5da1dfbd64af26b46e2a265741b916eff45b7658 Mon Sep 17 00:00:00 2001
From: Zhanty <Ztiany3@gmail.com>
Date: Thu, 8 Aug 2019 17:45:06 +0800
Subject: [PATCH] fix SingleLiveData memory leak

---
 .../android/base/app/aac/SingleLiveData.java  | 60 +++++++++++++++----
 1 file changed, 49 insertions(+), 11 deletions(-)

diff --git a/lib_base/src/main/java/com/android/base/app/aac/SingleLiveData.java b/lib_base/src/main/java/com/android/base/app/aac/SingleLiveData.java
index 2b7bbb2..7d24096 100644
--- a/lib_base/src/main/java/com/android/base/app/aac/SingleLiveData.java
+++ b/lib_base/src/main/java/com/android/base/app/aac/SingleLiveData.java
@@ -4,9 +4,12 @@ import android.arch.lifecycle.LifecycleOwner;
 import android.arch.lifecycle.MediatorLiveData;
 import android.arch.lifecycle.Observer;
 import android.support.annotation.NonNull;
+import android.support.annotation.Nullable;
 
-import java.util.Map;
-import java.util.WeakHashMap;
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.ListIterator;
 
 import timber.log.Timber;
 
@@ -14,7 +17,8 @@ public class SingleLiveData<T> extends MediatorLiveData<T> {
 
     private int mVersion = 0;
 
-    private final Map<Observer<T>, Observer<T>> cache = new WeakHashMap<>();
+    private final List<WeakReference<ObserverWrapper<T>>> mWrapperObserverList = new ArrayList<>();
+
 
     @Override
     public void observe(@NonNull LifecycleOwner owner, @NonNull Observer<T> observer) {
@@ -34,23 +38,37 @@ public class SingleLiveData<T> extends MediatorLiveData<T> {
 
     @Override
     public void removeObserver(@NonNull Observer<T> observer) {
-        Observer<T> wrapper = cache.remove(observer);
+        Observer<T> wrapper = findWrapper(observer);
         Timber.d("removeObserver() called with: observer = [" + observer + "], wrapper = [" + wrapper + "]");
         if (wrapper != null) {
             super.removeObserver(wrapper);
         }
     }
 
+    private ObserverWrapper<T> findWrapper(Observer<T> observer) {
+        ListIterator<WeakReference<ObserverWrapper<T>>> iterator = mWrapperObserverList.listIterator();
+
+        ObserverWrapper<T> target = null;
+
+        while (iterator.hasNext()) {
+            WeakReference<ObserverWrapper<T>> next = iterator.next();
+            ObserverWrapper<T> item = next.get();
+            if (item == null) {
+                iterator.remove();
+            } else if (item.mOrigin == observer) {
+                target =  item;
+            }
+        }
+
+        return target;
+    }
+
     private Observer<T> getOrNewObserver(@NonNull Observer<T> observer, int observerVersion) {
-        Observer<T> wrapper = cache.get(observer);
+        ObserverWrapper<T> wrapper = findWrapper(observer);
 
         if (wrapper == null) {
-            wrapper = t -> {
-                if (observerVersion < mVersion) {
-                    observer.onChanged(t);
-                }
-            };
-            cache.put(observer, wrapper);
+            wrapper = new ObserverWrapper<>(observerVersion, observer);
+            mWrapperObserverList.add(new WeakReference<>(wrapper));
         }
 
         Timber.d("getOrNewObserver() called with: observer = [" + observer + "], observerVersion = [" + observerVersion + "], wrapper = [" + wrapper + "]");
@@ -58,4 +76,24 @@ public class SingleLiveData<T> extends MediatorLiveData<T> {
         return wrapper;
     }
 
+    private class ObserverWrapper<E> implements Observer<E> {
+
+        private final int mObserverVersion;
+
+        private final Observer<E> mOrigin;
+
+        private ObserverWrapper(int observerVersion, Observer<E> origin) {
+            mObserverVersion = observerVersion;
+            mOrigin = origin;
+        }
+
+        @Override
+        public void onChanged(@Nullable E t) {
+            if (mObserverVersion < mVersion && mOrigin != null) {
+                mOrigin.onChanged(t);
+            }
+        }
+
+    }
+
 }
\ No newline at end of file