Cache Callbacks in PersistedTabData

PersistedTabData provides a mechanism for acquiring serialized UserData
objects in storage and attaching them to a Tab. The caching mechanism
provides the capability to cache callbacks acquiring the same data
which yields a performance improvement. It also means we avoid a bug
of the following:
1 FooTabData.from(...)
2 FooTabData.from(...)
3 setUserData changing what is stored on FooTabData
4 2) returning after 3) and changing what is stored on FooTabData

Bug: 1059602
Change-Id: Ib3534922ace1a653068caa17a78a7eb46d307dd7
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2236045Reviewed-by: David Trainor <dtrainor@chromium.org>
Reviewed-by: Tommy Nyquist <nyquist@chromium.org>
Commit-Queue: David Maunder <davidjm@chromium.org>
Cr-Commit-Position: refs/heads/master@{#779097}

chrome/android/chrome_java_sources.gni

@@ -1558,6 +1558,7 @@ chrome_java_sources = [
   "java/src/org/chromium/chrome/browser/tab/state/CriticalPersistedTabData.java",
   "java/src/org/chromium/chrome/browser/tab/state/EncryptedFilePersistedTabDataStorage.java",
   "java/src/org/chromium/chrome/browser/tab/state/FilePersistedTabDataStorage.java",
+  "java/src/org/chromium/chrome/browser/tab/state/MockPersistedTabData.java",
   "java/src/org/chromium/chrome/browser/tab/state/PersistedTabData.java",
   "java/src/org/chromium/chrome/browser/tab/state/PersistedTabDataConfiguration.java",
   "java/src/org/chromium/chrome/browser/tab/state/PersistedTabDataFactory.java",

chrome/android/chrome_test_java_sources.gni

@@ -514,6 +514,7 @@ chrome_test_java_sources = [
   "javatests/src/org/chromium/chrome/browser/tab/UndoIntegrationTest.java",
   "javatests/src/org/chromium/chrome/browser/tab/state/CriticalPersistedTabDataTest.java",
   "javatests/src/org/chromium/chrome/browser/tab/state/FilePersistedTabDataStorageTest.java",
+  "javatests/src/org/chromium/chrome/browser/tab/state/PersistedTabDataTest.java",
   "javatests/src/org/chromium/chrome/browser/tabbed_mode/TabbedNavigationBarColorControllerTest.java",
   "javatests/src/org/chromium/chrome/browser/tabmodel/AsyncTabCreationParamsManagerTest.java",
   "javatests/src/org/chromium/chrome/browser/tabmodel/ChromeTabCreatorTest.java",

chrome/android/java/src/org/chromium/chrome/browser/tab/DEPS

@@ -13,6 +13,7 @@ include_rules = [
   "+chrome/android/java/src/org/chromium/chrome/browser/ssl/ChromeSecurityStateModelDelegate.java",
   "+chrome/android/java/src/org/chromium/chrome/browser/ssl/SecurityStateModel.java",
   "+chrome/android/java/src/org/chromium/chrome/browser/tab",
+  "+chrome/browser/tab/java/src/org/chromium/chrome/browser/tab/Tab.java",
   "+chrome/android/java/src/org/chromium/chrome/browser/ui/TabObscuringHandler.java",
   "+chrome/android/java/src/org/chromium/chrome/browser/webapps/WebDisplayMode.java",
   "+chrome/browser/flags",

chrome/android/java/src/org/chromium/chrome/browser/tab/state/MockPersistedTabData.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.chrome.browser.tab.state;
+
+import org.chromium.base.Callback;
+import org.chromium.chrome.browser.tab.Tab;
+
+import java.nio.ByteBuffer;
+
+/**
+ * MockPersistedTabData object used for testing
+ */
+public class MockPersistedTabData extends PersistedTabData {
+    private int mField;
+
+    /**
+     * @param tab   tab associated with {@link MockPersistedTabData}
+     * @param field field stored in {@link MockPersistedTabData}
+     */
+    public MockPersistedTabData(Tab tab, int field) {
+        super(tab,
+                PersistedTabDataConfiguration.get(MockPersistedTabData.class, tab.isIncognito())
+                        .storage,
+                PersistedTabDataConfiguration.get(MockPersistedTabData.class, tab.isIncognito())
+                        .id);
+        mField = field;
+    }
+
+    private MockPersistedTabData(Tab tab, byte[] data, PersistedTabDataStorage storage, String id) {
+        super(tab, data, storage, id);
+    }
+
+    /**
+     * Acquire {@link MockPersistedTabData} from storage or create it and
+     * associate with {@link Tab}
+     * @param tab      {@link Tab} {@link MockPersistedTabData} will be associated with
+     * @param callback callback {@link MockPersistedTabData} will be passed back in
+     */
+    public static void from(Tab tab, Callback<MockPersistedTabData> callback) {
+        PersistedTabData.from(tab,
+                (data, storage, id)
+                        -> { return new MockPersistedTabData(tab, data, storage, id); },
+                ()
+                        -> {
+                    return null; /** Currently unused */
+                },
+                MockPersistedTabData.class, callback);
+    }
+
+    /**
+     * @return field stored in {@link MockPersistedTabData}
+     */
+    public int getField() {
+        return mField;
+    }
+
+    /**
+     * Sets field
+     * @param field new value of field
+     */
+    public void setField(int field) {
+        mField = field;
+        save();
+    }
+
+    @Override
+    public byte[] serialize() {
+        return ByteBuffer.allocate(4).putInt(mField).array();
+    }
+
+    @Override
+    public void deserialize(byte[] data) {
+        mField = ByteBuffer.wrap(data).getInt();
+    }
+
+    @Override
+    public void destroy() {}
+}

chrome/android/java/src/org/chromium/chrome/browser/tab/state/PersistedTabData.java

@@ -8,10 +8,17 @@ import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import org.chromium.base.Callback;
+import org.chromium.base.ThreadUtils;
 import org.chromium.base.UserData;
 import org.chromium.base.supplier.Supplier;
 import org.chromium.chrome.browser.tab.Tab;
 
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+
 /**
  * PersistedTabData is Tab data persisted across restarts
  * A constructor of taking a Tab and a byte[] (serialized
@@ -22,6 +29,7 @@ import org.chromium.chrome.browser.tab.Tab;
  */
 public abstract class PersistedTabData implements UserData {
     private static final String TAG = "PTD";
+    private static final Map<String, List<Callback>> sCachedCallbacks = new HashMap<>();
     protected final Tab mTab;
     private final PersistedTabDataStorage mPersistedTabDataStorage;
     private final String mPersistedTabDataId;
@@ -67,6 +75,7 @@ public abstract class PersistedTabData implements UserData {
     protected static <T extends PersistedTabData> void from(Tab tab,
             PersistedTabDataFactory<T> factory, Supplier<T> supplier, Class<T> clazz,
             Callback<T> callback) {
+        ThreadUtils.assertOnUiThread();
         // TODO(crbug.com/1059602) cache callbacks
         T persistedTabDataFromTab = getUserData(tab, clazz);
         if (persistedTabDataFromTab != null) {
@@ -74,6 +83,10 @@ public abstract class PersistedTabData implements UserData {
             callback.onResult(persistedTabDataFromTab);
             return;
         }
+        String key = String.format(Locale.ENGLISH, "%d-%s", tab.getId(), clazz.toString());
+        addCallback(key, callback);
+        // Only load data for the same key once
+        if (sCachedCallbacks.get(key).size() > 1) return;
         PersistedTabDataConfiguration config =
                 PersistedTabDataConfiguration.get(clazz, tab.isIncognito());
         config.storage.restore(tab.getId(), config.id, (data) -> {
@@ -86,10 +99,20 @@ public abstract class PersistedTabData implements UserData {
             if (persistedTabData != null) {
                 setUserData(tab, clazz, persistedTabData);
             }
-            callback.onResult(persistedTabData);
+            for (Callback cachedCallback : sCachedCallbacks.get(key)) {
+                cachedCallback.onResult(persistedTabData);
+            }
+            sCachedCallbacks.remove(key);
         });
     }
 
+    private static <T extends PersistedTabData> void addCallback(String key, Callback<T> callback) {
+        if (!sCachedCallbacks.containsKey(key)) {
+            sCachedCallbacks.put(key, new LinkedList<>());
+        }
+        sCachedCallbacks.get(key).add(callback);
+    }
+
     /**
      * Acquire {@link PersistedTabData} stored in {@link UserData} on a {@link Tab}
      * @param tab the {@link Tab}

chrome/android/java/src/org/chromium/chrome/browser/tab/state/PersistedTabDataConfiguration.java

@@ -15,7 +15,9 @@ public enum PersistedTabDataConfiguration {
     // TODO(crbug.com/1059650) investigate should this go in the app code?
     // Also investigate if the storage instance should be shared.
     CRITICAL_PERSISTED_TAB_DATA("CPTD", new FilePersistedTabDataStorage()),
-    ENCRYPTED_CRITICAL_PERSISTED_TAB_DATA("ECPTD", new EncryptedFilePersistedTabDataStorage());
+    ENCRYPTED_CRITICAL_PERSISTED_TAB_DATA("ECPTD", new EncryptedFilePersistedTabDataStorage()),
+    MOCK_PERSISTED_TAB_DATA("MPTD", new FilePersistedTabDataStorage()),
+    ENCRYPTED_MOCK_PERSISTED_TAB_DATA("EMPTD", new EncryptedFilePersistedTabDataStorage());
 
     private static final Map<Class<? extends PersistedTabData>, PersistedTabDataConfiguration>
             sLookup = new HashMap<>();
@@ -26,6 +28,8 @@ public enum PersistedTabDataConfiguration {
         // TODO(crbug.com/1060187) remove static initializer and initialization lazy
         sLookup.put(CriticalPersistedTabData.class, CRITICAL_PERSISTED_TAB_DATA);
         sEncryptedLookup.put(CriticalPersistedTabData.class, ENCRYPTED_CRITICAL_PERSISTED_TAB_DATA);
+        sLookup.put(MockPersistedTabData.class, MOCK_PERSISTED_TAB_DATA);
+        sEncryptedLookup.put(MockPersistedTabData.class, ENCRYPTED_MOCK_PERSISTED_TAB_DATA);
     }
 
     public final String id;

chrome/android/javatests/src/org/chromium/chrome/browser/tab/state/PersistedTabDataTest.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.chrome.browser.tab.state;
+
+import android.support.test.annotation.UiThreadTest;
+import android.support.test.filters.SmallTest;
+import android.support.test.rule.UiThreadTestRule;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+
+import org.chromium.base.test.BaseJUnit4ClassRunner;
+import org.chromium.chrome.browser.tab.MockTab;
+import org.chromium.chrome.browser.tab.Tab;
+
+/**
+ * Test relating to {@link PersistedTabData}
+ */
+@RunWith(BaseJUnit4ClassRunner.class)
+public class PersistedTabDataTest {
+    private static final int INITIAL_VALUE = 42;
+    private static final int CHANGED_VALUE = 51;
+
+    @Rule
+    public UiThreadTestRule mRule = new UiThreadTestRule();
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    @SmallTest
+    @UiThreadTest
+    @Test
+    public void testCacheCallbacks() throws InterruptedException {
+        Tab tab = new MockTab(1, false);
+        MockPersistedTabData mockPersistedTabData = new MockPersistedTabData(tab, INITIAL_VALUE);
+        mockPersistedTabData.save();
+        // 1
+        MockPersistedTabData.from(tab, (res) -> {
+            Assert.assertEquals(INITIAL_VALUE, res.getField());
+            tab.getUserDataHost().getUserData(MockPersistedTabData.class).setField(CHANGED_VALUE);
+            // Caching callbacks means 2) shouldn't overwrite CHANGED_VALUE
+            // back to INITIAL_VALUE in the callback.
+            MockPersistedTabData.from(
+                    tab, (ares) -> { Assert.assertEquals(CHANGED_VALUE, ares.getField()); });
+        });
+        // 2
+        MockPersistedTabData.from(tab, (res) -> {
+            Assert.assertEquals(CHANGED_VALUE, res.getField());
+            mockPersistedTabData.delete();
+        });
+    }
+}