Cache return values from BundleUtils.createIsolatedSplitContext()

This avoids getting a newly created ContextImpl from every call to
Context.createContextForSplit().

Bug: 1126301
Change-Id: I080841b697e2f28000d2a4d11ff197a015a1dca8
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2527852
Commit-Queue: Clark DuVall <cduvall@chromium.org>
Reviewed-by: Andrew Grieve <agrieve@chromium.org>
Cr-Commit-Position: refs/heads/master@{#825990}

base/BUILD.gn

@@ -3968,6 +3968,7 @@ if (is_android) {
   junit_binary("base_junit_tests") {
     sources = [
       "android/junit/src/org/chromium/base/ApplicationStatusTest.java",
+      "android/junit/src/org/chromium/base/BundleUtilsTest.java",
       "android/junit/src/org/chromium/base/CallbackControllerTest.java",
       "android/junit/src/org/chromium/base/CollectionUtilTest.java",
       "android/junit/src/org/chromium/base/DiscardableReferencePoolTest.java",

base/android/java/src/org/chromium/base/BundleUtils.java

@@ -8,6 +8,8 @@ import android.content.Context;
 import android.content.pm.PackageManager;
 import android.os.Build;
 
+import androidx.collection.SimpleArrayMap;
+
 import dalvik.system.BaseDexClassLoader;
 
 import org.chromium.base.annotations.CalledByNative;
@@ -33,6 +35,10 @@ import java.util.Arrays;
  * library APKs.
  */
 public final class BundleUtils {
+    // These Contexts will be needed throughout the lifetime of the application, so it is fine to
+    // have strong references in the map.
+    private static final SimpleArrayMap<String, Context> sIsolatedSplitContextCache =
+            new SimpleArrayMap<>();
     private static Boolean sIsBundle;
 
     /**
@@ -89,11 +95,18 @@ public final class BundleUtils {
             return base;
         }
 
-        try {
-            return ApiHelperForO.createContextForSplit(base, splitName);
-        } catch (PackageManager.NameNotFoundException e) {
-            throw new RuntimeException(e);
+        // Context.createContextForSplit() creates a new ContextImpl for each call, so we cache the
+        // returned context for future calls.
+        Context splitContext = sIsolatedSplitContextCache.get(splitName);
+        if (splitContext == null) {
+            try {
+                splitContext = ApiHelperForO.createContextForSplit(base, splitName);
+            } catch (PackageManager.NameNotFoundException e) {
+                throw new RuntimeException(e);
+            }
+            sIsolatedSplitContextCache.put(splitName, splitContext);
         }
+        return splitContext;
     }
 
     /* Returns absolute path to a native library in a feature module. */

base/android/junit/src/org/chromium/base/BundleUtilsTest.java

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package org.chromium.base;
+
+import android.content.Context;
+import android.content.ContextWrapper;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.annotation.Config;
+
+import org.chromium.base.test.BaseRobolectricTestRunner;
+
+/**
+ * Unit tests for the BundleUtils class.
+ */
+@RunWith(BaseRobolectricTestRunner.class)
+@Config(manifest = Config.NONE)
+public class BundleUtilsTest {
+    private static final String SPLIT_A = "a";
+    private static final String SPLIT_B = "b";
+
+    @Test
+    public void testIsolatedSplitContextCached() {
+        Context baseContext = ContextUtils.getApplicationContext();
+        Context context = new ContextWrapper(baseContext) {
+            @Override
+            public Context createContextForSplit(String splitName) {
+                return new ContextWrapper(baseContext);
+            }
+        };
+        Context contextA = BundleUtils.createIsolatedSplitContext(context, SPLIT_A);
+        Assert.assertEquals(contextA, BundleUtils.createIsolatedSplitContext(context, SPLIT_A));
+
+        Context contextB = BundleUtils.createIsolatedSplitContext(context, SPLIT_B);
+        Assert.assertNotEquals(contextA, contextB);
+        Assert.assertEquals(contextB, BundleUtils.createIsolatedSplitContext(context, SPLIT_B));
+    }
+}