BaseBundle.java:: Clear underlying parcel when no longer referernced.

Follow up to ag/18795008.
Lazy Bundles, (aosp/1787847), introduced a change in behavior where a Parcel
created as part of initializing a Bundle is dependent on the next ART GCrun to be
recycled, causing a short term memory-leak. We previously force recycled
the parcel when bundle.clear() was called.

This commit takes it a step further and adds  reference counting for the number of
lazy values in mMap. When this reaches 0, we can safely eager recycle the
underlying parcel.

This still does not work when parcel references have been shared outside
the owning bundle (i.e, when mOwnsLazyValues is false.) We have two
options here:

1. Allow the GC to eventually claim the parcel (current behaviour).
2. Track lazy values refs globally across bundles. High complexity cost,
   and likely high performance impact.

We can not unmwrap the lazy values before copying, in case the
originating classloader can not instantiate the given class type. See
go/lazy-bundle

In initializeFromParcelLocked(), renamed boolean parameter recycleParcel to ownsParcel since the method no longer recycles in all cases. Also, only create lazy values if ownsParcel = true so that we don't have lazy values referencing the passed in parcel in case the caller doesn't own the parcel (which means the parcel could be recycled out of our control).

Bug: 233216232
Test: Reproduced linked bug on-device
Test: atest android.os.cts.ParcelTest android.os.cts.BundleTest
android.os.BundleTest android.os.ParcelTest android.os.BundleRecylingTest

Change-Id: I8a878a6d0055b04b2611909b4b17977ba5677a4f

core/java/android/os/BaseBundle.java

@@ -29,6 +29,7 @@ import android.util.SparseArray;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.IndentingPrintWriter;
+import com.android.internal.util.Preconditions;
 
 import java.io.Serializable;
 import java.lang.ref.WeakReference;
@@ -113,17 +114,21 @@ public class BaseBundle {
      */
     private boolean mParcelledByNative;
 
-    /*
+    /**
      * Flag indicating if mParcelledData is only referenced in this bundle.
-     * mParcelledData could be referenced by other bundles if mMap contains lazy values,
+     * mParcelledData could be referenced elsewhere if mMap contains lazy values,
      * and bundle data is copied to another bundle using putAll or the copy constructors.
      */
     boolean mOwnsLazyValues = true;
 
-    /*
+    /** Tracks how many lazy values are referenced in mMap */
+    private int mLazyValues = 0;
+
+    /**
      * As mParcelledData is set to null when it is unparcelled, we keep a weak reference to
      * it to aid in recycling it. Do not use this reference otherwise.
-     */
+     * Is non-null iff mMap contains lazy values.
+    */
     private WeakReference<Parcel> mWeakParcelledData = null;
 
     /**
@@ -310,7 +315,8 @@ public class BaseBundle {
         synchronized (this) {
             final Parcel source = mParcelledData;
             if (source != null) {
-                initializeFromParcelLocked(source, /*recycleParcel=*/ true, mParcelledByNative);
+                Preconditions.checkState(mOwnsLazyValues);
+                initializeFromParcelLocked(source, /*ownsParcel*/ true, mParcelledByNative);
             } else {
                 if (DEBUG) {
                     Log.d(TAG, "unparcel "
@@ -401,11 +407,23 @@ public class BaseBundle {
                 }
             }
             mMap.setValueAt(i, object);
+            mLazyValues--;
+            if (mOwnsLazyValues) {
+                Preconditions.checkState(mLazyValues >= 0, "Lazy values ref count below 0");
+                // No more lazy values in mMap, so we can recycle the parcel early rather than
+                // waiting for the next GC run
+                if (mLazyValues == 0) {
+                    Preconditions.checkState(mWeakParcelledData.get() != null,
+                            "Parcel recycled earlier than expected");
+                    recycleParcel(mWeakParcelledData.get());
+                    mWeakParcelledData = null;
+                }
+            }
         }
         return (clazz != null) ? clazz.cast(object) : (T) object;
     }
 
-    private void initializeFromParcelLocked(@NonNull Parcel parcelledData, boolean recycleParcel,
+    private void initializeFromParcelLocked(@NonNull Parcel parcelledData, boolean ownsParcel,
             boolean parcelledByNative) {
         if (isEmptyParcel(parcelledData)) {
             if (DEBUG) {
@@ -437,9 +455,10 @@ public class BaseBundle {
             map.erase();
             map.ensureCapacity(count);
         }
+        int numLazyValues = 0;
         try {
-            recycleParcel &= parcelledData.readArrayMap(map, count, !parcelledByNative,
+            numLazyValues = parcelledData.readArrayMap(map, count, !parcelledByNative,
-                    /* lazy */ true, mClassLoader);
+                    /* lazy */ ownsParcel, mClassLoader);
         } catch (BadParcelableException e) {
             if (sShouldDefuse) {
                 Log.w(TAG, "Failed to parse Bundle, but defusing quietly", e);
@@ -448,14 +467,19 @@ public class BaseBundle {
                 throw e;
             }
         } finally {
-            mMap = map;
+            mWeakParcelledData = null;
-            if (recycleParcel) {
+            if (ownsParcel) {
-                recycleParcel(parcelledData);
+                if (numLazyValues == 0) {
-                mWeakParcelledData = null;
+                    recycleParcel(parcelledData);
-            } else {
+                } else {
-                mWeakParcelledData = new WeakReference<>(parcelledData);
+                    mWeakParcelledData = new WeakReference<>(parcelledData);
+                }
             }
+
+            mLazyValues = numLazyValues;
             mParcelledByNative = false;
+            mMap = map;
+            // Set field last as it is volatile
             mParcelledData = null;
         }
         if (DEBUG) {
@@ -592,13 +616,17 @@ public class BaseBundle {
 
     /**
      * Removes all elements from the mapping of this Bundle.
+     * Recycles the underlying parcel if it is still present.
      */
     public void clear() {
         unparcel();
         if (mOwnsLazyValues && mWeakParcelledData != null) {
             recycleParcel(mWeakParcelledData.get());
-            mWeakParcelledData = null;
         }
+
+        mWeakParcelledData = null;
+        mLazyValues = 0;
+        mOwnsLazyValues = true;
         mMap.clear();
     }
 
@@ -1844,8 +1872,8 @@ public class BaseBundle {
             // had been constructed with parcel or to make sure they trigger deserialization of the
             // bundle immediately; neither of which is obvious.
             synchronized (this) {
-                initializeFromParcelLocked(parcel, /*recycleParcel=*/ false, isNativeBundle);
+                mOwnsLazyValues = false;
-                unparcel(/* itemwise */ true);
+                initializeFromParcelLocked(parcel, /*ownsParcel*/ false, isNativeBundle);
             }
             return;
         }
@@ -1862,6 +1890,7 @@ public class BaseBundle {
                 + ": " + length + " bundle bytes starting at " + offset);
         p.setDataPosition(0);
 
+        mOwnsLazyValues = true;
         mParcelledByNative = isNativeBundle;
         mParcelledData = p;
     }

core/java/android/os/Parcel.java

@@ -71,7 +71,6 @@ import java.util.Set;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.IntFunction;
-import java.util.function.Supplier;
 
 /**
  * Container for a message (data and object references) that can
@@ -5242,20 +5241,20 @@ public final class Parcel {
      * Reads a map into {@code map}.
      *
      * @param sorted Whether the keys are sorted by their hashes, if so we use an optimized path.
-     * @param lazy   Whether to populate the map with lazy {@link Supplier} objects for
+     * @param lazy   Whether to populate the map with lazy {@link Function} objects for
      *               length-prefixed values. See {@link Parcel#readLazyValue(ClassLoader)} for more
      *               details.
-     * @return whether the parcel can be recycled or not.
+     * @return a count of the lazy values in the map
      * @hide
      */
-    boolean readArrayMap(ArrayMap<? super String, Object> map, int size, boolean sorted,
+    int readArrayMap(ArrayMap<? super String, Object> map, int size, boolean sorted,
             boolean lazy, @Nullable ClassLoader loader) {
-        boolean recycle = true;
+        int lazyValues = 0;
         while (size > 0) {
             String key = readString();
             Object value = (lazy) ? readLazyValue(loader) : readValue(loader);
             if (value instanceof LazyValue) {
-                recycle = false;
+                lazyValues++;
             }
             if (sorted) {
                 map.append(key, value);
@@ -5267,7 +5266,7 @@ public final class Parcel {
         if (sorted) {
             map.validate();
         }
-        return recycle;
+        return lazyValues;
     }
 
     /**

core/tests/mockingcoretests/src/android/os/BundleRecyclingTest.java

@@ -23,6 +23,7 @@ import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -34,7 +35,6 @@ import androidx.test.filters.SmallTest;
 
 import com.android.dx.mockito.inline.extended.StaticMockitoSession;
 
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.quality.Strictness;
@@ -56,16 +56,11 @@ public class BundleRecyclingTest {
     private Parcel mParcelSpy;
     private Bundle mBundle;
 
-    @Before
-    public void setUp() throws Exception {
-        setUpBundle(/* hasLazy */ true);
-    }
-
     @Test
     public void bundleClear_whenUnparcelledWithoutLazy_recyclesParcelOnce() {
-        setUpBundle(/* hasLazy */ false);
+        setUpBundle(/* lazy */ 0, /* nonLazy */ 1);
         // Will unparcel and immediately recycle parcel
-        assertNotNull(mBundle.getString("key"));
+        assertNotNull(mBundle.getString("nonLazy0"));
         verify(mParcelSpy, times(1)).recycle();
         assertFalse(mBundle.isDefinitelyEmpty());
 
@@ -77,6 +72,7 @@ public class BundleRecyclingTest {
 
     @Test
     public void bundleClear_whenParcelled_recyclesParcel() {
+        setUpBundle(/* lazy */ 1);
         assertTrue(mBundle.isParcelled());
         verify(mParcelSpy, times(0)).recycle();
 
@@ -89,24 +85,10 @@ public class BundleRecyclingTest {
         verify(mParcelSpy, times(1)).recycle();
     }
 
-    @Test
-    public void bundleClear_whenUnparcelledWithLazyValueUnwrapped_recyclesParcel() {
-        // Will unparcel with a lazy value, and immediately unwrap the lazy value,
-        // with no lazy values left at the end of getParcelable
-        assertNotNull(mBundle.getParcelable("key", CustomParcelable.class));
-        verify(mParcelSpy, times(0)).recycle();
-
-        mBundle.clear();
-        verify(mParcelSpy, times(1)).recycle();
-        assertTrue(mBundle.isDefinitelyEmpty());
-
-        // Should not recycle again
-        mBundle.clear();
-        verify(mParcelSpy, times(1)).recycle();
-    }
-
     @Test
     public void bundleClear_whenUnparcelledWithLazy_recyclesParcel() {
+        setUpBundle(/* lazy */ 1);
+
         // Will unparcel but keep the CustomParcelable lazy
         assertFalse(mBundle.isEmpty());
         verify(mParcelSpy, times(0)).recycle();
@@ -122,6 +104,8 @@ public class BundleRecyclingTest {
 
     @Test
     public void bundleClear_whenClearedWithSharedParcel_doesNotRecycleParcel() {
+        setUpBundle(/* lazy */ 1);
+
         Bundle copy = new Bundle();
         copy.putAll(mBundle);
 
@@ -136,6 +120,8 @@ public class BundleRecyclingTest {
 
     @Test
     public void bundleClear_whenClearedWithCopiedParcel_doesNotRecycleParcel() {
+        setUpBundle(/* lazy */ 1);
+
         // Will unparcel but keep the CustomParcelable lazy
         assertFalse(mBundle.isEmpty());
 
@@ -151,7 +137,101 @@ public class BundleRecyclingTest {
         verify(mParcelSpy, never()).recycle();
     }
 
-    private void setUpBundle(boolean hasLazy) {
+    @Test
+    public void bundleGet_whenUnparcelledWithLazyValueUnwrapped_recyclesParcel() {
+        setUpBundle(/* lazy */ 1);
+
+        // Will unparcel with a lazy value, and immediately unwrap the lazy value,
+        // with no lazy values left at the end of getParcelable
+        // Ref counting should immediately recycle
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        verify(mParcelSpy, times(1)).recycle();
+
+        // Should not recycle again
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        mBundle.clear();
+        verify(mParcelSpy, times(1)).recycle();
+    }
+
+    @Test
+    public void bundleGet_whenMultipleLazy_recyclesParcelWhenAllUnwrapped() {
+        setUpBundle(/* lazy */ 2);
+
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        verify(mParcelSpy, times(0)).recycle();
+
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        verify(mParcelSpy, times(0)).recycle();
+
+        assertNotNull(mBundle.getParcelable("lazy1", CustomParcelable.class));
+        verify(mParcelSpy, times(1)).recycle();
+
+        // Should not recycle again
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        mBundle.clear();
+        verify(mParcelSpy, times(1)).recycle();
+        assertTrue(mBundle.isDefinitelyEmpty());
+    }
+
+    @Test
+    public void bundleGet_whenLazyAndNonLazy_recyclesParcelWhenAllUnwrapped() {
+        setUpBundle(/* lazy */ 1, /* nonLazy */ 1);
+
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        verify(mParcelSpy, times(1)).recycle();
+
+        // Should not recycle again
+        assertNotNull(mBundle.getString("nonLazy0"));
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        mBundle.clear();
+        verify(mParcelSpy, times(1)).recycle();
+    }
+
+    @Test
+    public void bundleGet_whenLazyAndNonLazy_doesNotRecycleWhenOnlyNonLazyRetrieved() {
+        setUpBundle(/* lazy */ 1, /* nonLazy */ 1);
+
+        assertNotNull(mBundle.getString("nonLazy0"));
+        verify(mParcelSpy, times(0)).recycle();
+
+        assertNotNull(mBundle.getString("nonLazy0"));
+        verify(mParcelSpy, times(0)).recycle();
+
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        verify(mParcelSpy, times(1)).recycle();
+    }
+
+    @Test
+    public void bundleGet_withWithSharedParcel_doesNotRecycleParcel() {
+        setUpBundle(/* lazy */ 1);
+
+        Bundle copy = new Bundle();
+        copy.putAll(mBundle);
+
+        assertNotNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        mBundle.clear();
+
+        assertNotNull(copy.getParcelable("lazy0", CustomParcelable.class));
+        copy.clear();
+
+        verify(mParcelSpy, never()).recycle();
+    }
+
+    @Test
+    public void bundleGet_getAfterLazyCleared_doesNotRecycleAgain() {
+        setUpBundle(/* lazy */ 1);
+        mBundle.clear();
+        verify(mParcelSpy, times(1)).recycle();
+
+        assertNull(mBundle.getParcelable("lazy0", CustomParcelable.class));
+        verify(mParcelSpy, times(1)).recycle();
+    }
+
+    private void setUpBundle(int lazy) {
+        setUpBundle(lazy, /* nonLazy */ 0);
+    }
+
+    private void setUpBundle(int lazy, int nonLazy) {
         AtomicReference<Parcel> parcel = new AtomicReference<>();
         StaticMockitoSession session = mockitoSession()
                 .strictness(Strictness.STRICT_STUBS)
@@ -166,7 +246,7 @@ public class BundleRecyclingTest {
 
         Bundle bundle = new Bundle();
         bundle.setClassLoader(getClass().getClassLoader());
-        Parcel p = createBundle(hasLazy);
+        Parcel p = createBundle(lazy, nonLazy);
         bundle.readFromParcel(p);
         p.recycle();
 
@@ -179,13 +259,17 @@ public class BundleRecyclingTest {
     /**
      * Create a test bundle, parcel it and return the parcel.
      */
-    private Parcel createBundle(boolean hasLazy) {
+    private Parcel createBundle(int lazy, int nonLazy) {
         final Bundle source = new Bundle();
-        if (hasLazy) {
+
-            source.putParcelable("key", new CustomParcelable(13, "Tiramisu"));
+        for (int i = 0; i < nonLazy; i++) {
-        } else {
+            source.putString("nonLazy" + i, "Tiramisu");
-            source.putString("key", "tiramisu");
         }
+
+        for (int i = 0; i < lazy; i++) {
+            source.putParcelable("lazy" + i, new CustomParcelable(13, "Tiramisu"));
+        }
+
         return getParcelledBundle(source);
     }