Separate CCT dynamic module loading and incrementation of its usage counter.

Remove the WeakReference in the LoadModuleCallback.
Add check to not create the second loading task if dynamic module if it is currently loading.

Bug: 853793
Change-Id: Ie02a7dd63e3f9e8a21e48d18affb3c25d90ab0ed
Reviewed-on: https://chromium-review.googlesource.com/c/1286661Reviewed-by: Peter Conn <peconn@chromium.org>
Reviewed-by: Michael van Ouwerkerk <mvanouwerkerk@chromium.org>
Commit-Queue: Anna Malova <amalova@chromium.org>
Cr-Commit-Position: refs/heads/master@{#601584}

chrome/android/java/src/org/chromium/chrome/browser/customtabs/CustomTabActivity.java

@@ -72,6 +72,7 @@ import org.chromium.chrome.browser.contextual_suggestions.ContextualSuggestionsM
 import org.chromium.chrome.browser.customtabs.dynamicmodule.ActivityDelegate;
 import org.chromium.chrome.browser.customtabs.dynamicmodule.ActivityHostImpl;
 import org.chromium.chrome.browser.customtabs.dynamicmodule.ModuleEntryPoint;
+import org.chromium.chrome.browser.customtabs.dynamicmodule.ModuleLoader;
 import org.chromium.chrome.browser.customtabs.dynamicmodule.ModuleMetrics;
 import org.chromium.chrome.browser.dependency_injection.ChromeActivityCommonsModule;
 import org.chromium.chrome.browser.dependency_injection.CustomTabActivityComponent;
@@ -117,7 +118,6 @@ import org.chromium.ui.base.PageTransition;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
-import java.lang.ref.WeakReference;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
@@ -198,11 +198,12 @@ public class CustomTabActivity extends ChromeActivity<CustomTabActivityComponent
     private ModuleEntryPoint mModuleEntryPoint;
     @Nullable
     private ActivityDelegate mModuleActivityDelegate;
-    @Nullable
-    private Runnable mLoadModuleCancelRunnable;
+
     private boolean mModuleOnStartPending;
     private boolean mModuleOnResumePending;
     private boolean mHasSetOverlayView;
+    @Nullable
+    private LoadModuleCallback mModuleCallback;
 
     private ActivityTabTaskDescriptionHelper mTaskDescriptionHelper;
 
@@ -353,49 +354,40 @@ public class CustomTabActivity extends ChromeActivity<CustomTabActivityComponent
             return;
         }
 
-        mLoadModuleCancelRunnable =
-                mConnection.getModuleLoader(componentName).loadModule(new LoadModuleCallback(this));
+        ModuleLoader moduleLoader = mConnection.getModuleLoader(componentName);
+        moduleLoader.loadModule();
+        mModuleCallback = new LoadModuleCallback();
+        moduleLoader.addCallbackAndIncrementUseCount(mModuleCallback);
     }
 
     private boolean isModuleLoading() {
-        return mLoadModuleCancelRunnable != null;
-    }
-
-    private static class LoadModuleCallback implements Callback<ModuleEntryPoint> {
-        private final WeakReference<CustomTabActivity> mActivity;
-
-        LoadModuleCallback(CustomTabActivity activity) {
-            mActivity = new WeakReference<>(activity);
-        }
-
-        @Override
-        public void onResult(@Nullable ModuleEntryPoint entryPoint) {
-            CustomTabActivity activity = mActivity.get();
-            if (activity == null || activity.isActivityDestroyed()) return;
-            activity.onModuleLoaded(entryPoint);
-        }
+        return mModuleCallback != null;
     }
 
     /**
-     * Receives the entry point if it was loaded successfully, or null if there was a problem. This
-     * is always called on the UI thread.
+     * Callback to receive the entry point if it was loaded successfully,
+     * or null if there was a problem. This is always called on the UI thread.
      */
-    private void onModuleLoaded(@Nullable ModuleEntryPoint entryPoint) {
-        mLoadModuleCancelRunnable = null;
-        if (entryPoint == null) return;
+    private class LoadModuleCallback implements Callback<ModuleEntryPoint> {
+        @Override
+        public void onResult(@Nullable ModuleEntryPoint entryPoint) {
+            if (entryPoint == null) return;
 
-        mModuleEntryPoint = entryPoint;
+            mModuleEntryPoint = entryPoint;
 
-        long createActivityDelegateStartTime = ModuleMetrics.now();
-        mModuleActivityDelegate = entryPoint.createActivityDelegate(new ActivityHostImpl(this));
-        ModuleMetrics.recordCreateActivityDelegateTime(createActivityDelegateStartTime);
-        mModuleActivityDelegate.onCreate(getSavedInstanceState());
+            long createActivityDelegateStartTime = ModuleMetrics.now();
+            mModuleActivityDelegate = entryPoint.createActivityDelegate(
+                    new ActivityHostImpl(CustomTabActivity.this));
+            ModuleMetrics.recordCreateActivityDelegateTime(createActivityDelegateStartTime);
+            mModuleActivityDelegate.onCreate(getSavedInstanceState());
 
-        if (mModuleOnStartPending) startModule();
-        if (mModuleOnResumePending) resumeModule();
+            if (mModuleOnStartPending) startModule();
+            if (mModuleOnResumePending) resumeModule();
 
-        mConnection.setActivityDelegateForSession(mSession, mModuleActivityDelegate,
-                mModuleEntryPoint.getModuleVersion());
+            mConnection.setActivityDelegateForSession(mSession, mModuleActivityDelegate,
+                    mModuleEntryPoint.getModuleVersion());
+            mModuleCallback = null;
+        }
     }
 
     private void startModule() {
@@ -997,10 +989,6 @@ public class CustomTabActivity extends ChromeActivity<CustomTabActivityComponent
     @Override
     protected void onDestroyInternal() {
         super.onDestroyInternal();
-        if (mLoadModuleCancelRunnable != null) {
-            mLoadModuleCancelRunnable.run();
-            mLoadModuleCancelRunnable = null;
-        }
         if (mModuleActivityDelegate != null) {
             mModuleActivityDelegate.onDestroy();
             mModuleActivityDelegate = null;
@@ -1011,7 +999,8 @@ public class CustomTabActivity extends ChromeActivity<CustomTabActivityComponent
             moduleComponentName = mIntentDataProvider.getModuleComponentName();
         }
         if (moduleComponentName != null) {
-            mConnection.getModuleLoader(moduleComponentName).decrementModuleUseCount();
+            mConnection.getModuleLoader(moduleComponentName)
+                    .removeCallbackAndDecrementUseCount(mModuleCallback);
         }
         if (mIncognitoTabHost != null) {
             IncognitoTabHostRegistry.getInstance().unregister(mIncognitoTabHost);

chrome/android/java/src/org/chromium/chrome/browser/customtabs/dynamicmodule/ModuleLoader.java

@@ -16,6 +16,8 @@ import android.support.annotation.Nullable;
 import org.chromium.base.Callback;
 import org.chromium.base.ContextUtils;
 import org.chromium.base.Log;
+import org.chromium.base.ObserverList;
+import org.chromium.base.VisibleForTesting;
 import org.chromium.base.task.AsyncTask;
 import org.chromium.chrome.browser.ChromeApplication;
 import org.chromium.chrome.browser.ChromeFeatureList;
@@ -39,6 +41,10 @@ public class ModuleLoader {
      */
     private int mModuleUseCount;
 
+    private boolean mIsModuleLoading;
+
+    private final ObserverList<Callback<ModuleEntryPoint>> mCallbacks = new ObserverList<>();
+
     /**
      * The timestamp of the moment the module became unused. This is used to determine whether or
      * not to continue caching it. A value of -1 indicates there is no usable value.
@@ -86,46 +92,53 @@ public class ModuleLoader {
     }
 
     /**
-     * If the module is not loaded yet, dynamically loads the module entry point class. If
-     * successful, the callback will receive a {@link ModuleEntryPoint} asynchronously. If the
-     * module fails to load, the callback will receive null. If the module was already loaded and a
-     * reference to it is still held, the callback will synchronously receive a
-     * {@link ModuleEntryPoint}.
+     * If the module is not loaded yet, dynamically loads the module entry point class.
+     */
+    public void loadModule() {
+        if (mIsModuleLoading) return;
+
+        // If module has been already loaded all callbacks must be notified synchronously.
+        // {@see #addCallbackAndIncrementUseCount}
+        if (mModuleEntryPoint != null) {
+            assert mCallbacks.isEmpty();
+            return;
+        }
+
+        Context moduleContext = getModuleContext(mComponentName.getPackageName());
+        if (moduleContext == null) {
+            runAndClearCallbacks();
+            return;
+        }
+
+        mIsModuleLoading = true;
+        new LoadClassTask(moduleContext).executeOnExecutor(AsyncTask.THREAD_POOL_EXECUTOR);
+    }
+
+    /**
+     * Register a callback to receive a {@link ModuleEntryPoint} asynchronously.
+     * If the module fails to load, the callback will receive null.
+     * If the module was already loaded and a reference to it is still held,
+     * the callback will synchronously receive a {@link ModuleEntryPoint}.
+     *
+     * Module use count is incremented when a callback notified.
      *
      * @param callback The callback to receive the result.
-     * @return If the module is not loaded yet, an {@link AsyncTask} will be used to load it and
-     *     a {@link Runnable} returned to the caller for cancelling the task if necessary. If the
-     *     module was already loaded, null is returned.
      */
-    @Nullable
-    public Runnable loadModule(Callback<ModuleEntryPoint> callback) {
+    public void addCallbackAndIncrementUseCount(Callback<ModuleEntryPoint> callback) {
         if (mModuleEntryPoint != null) {
             mModuleUseCount++;
             mModuleUnusedTimeMs = -1;
             ModuleMetrics.recordLoadResult(ModuleMetrics.LoadResult.SUCCESS_CACHED);
             callback.onResult(mModuleEntryPoint);
-            return null;
-        }
-
-        Context moduleContext = getModuleContext(mComponentName.getPackageName());
-        if (moduleContext == null) {
-            callback.onResult(null);
-            return null;
+            return;
         }
-
-        // TODO(crbug.com/864520): Define and return a CancellablePromise instead.
-        final LoadClassTask task = new LoadClassTask(moduleContext, callback);
-        task.executeOnExecutor(AsyncTask.THREAD_POOL_EXECUTOR);
-        return new Runnable() {
-            @Override
-            public void run() {
-                task.cancel(false);
-            }
-        };
+        mCallbacks.addObserver(callback);
     }
 
-    public void decrementModuleUseCount() {
-        if (mModuleEntryPoint == null) return;
+    public void removeCallbackAndDecrementUseCount(Callback<ModuleEntryPoint> callback) {
+        boolean isPendingCallback = mCallbacks.removeObserver(callback);
+        if (mModuleEntryPoint == null || isPendingCallback) return;
+
         mModuleUseCount--;
         if (mModuleUseCount == 0) {
             mModuleUnusedTimeMs = ModuleMetrics.now();
@@ -172,22 +185,35 @@ public class ModuleLoader {
         mModuleUnusedTimeMs = -1;
     }
 
+    /**
+     * Notify all callbacks which are waiting for module loading. Each callback is needed to notify
+     * only once therefore all callbacks are cleared after call.
+     */
+    private void runAndClearCallbacks() {
+        assert !mIsModuleLoading;
+        if (mModuleEntryPoint != null && mCallbacks.size() > 0) {
+            mModuleUseCount += mCallbacks.size();
+            mModuleUnusedTimeMs = -1;
+        }
+
+        for (Callback<ModuleEntryPoint> callback: mCallbacks) {
+            callback.onResult(mModuleEntryPoint);
+        }
+        mCallbacks.clear();
+    }
+
     /**
      * A task for loading the module entry point class on a background thread.
      */
     private class LoadClassTask extends AsyncTask<Class<?>> {
         private final Context mModuleContext;
-        private final Callback<ModuleEntryPoint> mCallback;
 
         /**
          * Constructs the task.
          * @param moduleContext The context for the package to load the class from.
-         * @param callback The callback to receive the result of the task. If there was a problem
-         *     the result will be null.
          */
-        LoadClassTask(Context moduleContext, Callback<ModuleEntryPoint> callback) {
+        LoadClassTask(Context moduleContext) {
             mModuleContext = moduleContext;
-            mCallback = callback;
         }
 
         @Override
@@ -219,8 +245,9 @@ public class ModuleLoader {
 
         @Override
         protected void onPostExecute(@Nullable Class<?> clazz) {
+            mIsModuleLoading = false;
             if (clazz == null) {
-                mCallback.onResult(null);
+                runAndClearCallbacks();
                 return;
             }
 
@@ -242,7 +269,7 @@ public class ModuleLoader {
                             moduleHost.getHostVersion(), entryPoint.getMinimumHostVersion(),
                             entryPoint.getModuleVersion(), moduleHost.getMinimumModuleVersion());
                     ModuleMetrics.recordLoadResult(ModuleMetrics.LoadResult.INCOMPATIBLE_VERSION);
-                    mCallback.onResult(null);
+                    runAndClearCallbacks();
                     return;
                 }
 
@@ -253,9 +280,8 @@ public class ModuleLoader {
                 entryPoint.init(moduleHost);
                 ModuleMetrics.recordLoadResult(ModuleMetrics.LoadResult.SUCCESS_NEW);
                 mModuleEntryPoint = entryPoint;
-                mModuleUseCount = 1;
-                mModuleUnusedTimeMs = -1;
-                mCallback.onResult(entryPoint);
+                mModuleUnusedTimeMs = ModuleMetrics.now();
+                runAndClearCallbacks();
                 return;
             } catch (Exception e) {
                 // No multi-catch below API level 19 for reflection exceptions.
@@ -263,7 +289,7 @@ public class ModuleLoader {
                 Log.e(TAG, "Could not instantiate class %s", mComponentName.getClassName(), e);
                 ModuleMetrics.recordLoadResult(ModuleMetrics.LoadResult.INSTANTIATION_EXCEPTION);
             }
-            mCallback.onResult(null);
+            runAndClearCallbacks();
         }
     }
 
@@ -289,4 +315,9 @@ public class ModuleLoader {
         return entryPoint.getModuleVersion() >= moduleHost.getMinimumModuleVersion()
                 && moduleHost.getHostVersion() >= entryPoint.getMinimumHostVersion();
     }
+
+    @VisibleForTesting
+    public int getModuleUseCount() {
+        return mModuleUseCount;
+    }
 }

chrome/android/javatests/src/org/chromium/chrome/browser/customtabs/CustomTabsDynamicModuleTest.java

@@ -4,6 +4,7 @@
 
 package org.chromium.chrome.browser.customtabs;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 
 import android.content.ComponentName;
@@ -11,18 +12,26 @@ import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import org.chromium.base.Callback;
 import org.chromium.base.library_loader.LibraryLoader;
 import org.chromium.base.library_loader.LibraryProcessType;
 import org.chromium.base.test.util.CallbackHelper;
+import org.chromium.chrome.browser.ChromeFeatureList;
 import org.chromium.chrome.browser.customtabs.dynamicmodule.BaseModuleEntryPoint;
 import org.chromium.chrome.browser.customtabs.dynamicmodule.IActivityDelegate;
 import org.chromium.chrome.browser.customtabs.dynamicmodule.IActivityHost;
 import org.chromium.chrome.browser.customtabs.dynamicmodule.IModuleHost;
+import org.chromium.chrome.browser.customtabs.dynamicmodule.ModuleEntryPoint;
+import org.chromium.chrome.browser.customtabs.dynamicmodule.ModuleLoader;
 import org.chromium.chrome.test.ChromeJUnit4ClassRunner;
+import org.chromium.chrome.test.util.browser.Features;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.TimeoutException;
 
 /**
@@ -35,20 +44,78 @@ public class CustomTabsDynamicModuleTest {
             InstrumentationRegistry.getInstrumentation().getContext().getPackageName(),
             FakeCCTDynamicModule.class.getName());
 
+    @Rule
+    public Features.JUnitProcessor mProcessor = new Features.JUnitProcessor();
+
     @Before
     public void setUp() throws Exception {
         LibraryLoader.getInstance().ensureInitialized(LibraryProcessType.PROCESS_BROWSER);
     }
 
+    /**
+     * Test fake dynamic module is correctly loaded.
+     */
     @Test
     @SmallTest
     public void testModuleLoading() throws TimeoutException, InterruptedException {
-        CallbackHelper loadingWaiter = new CallbackHelper();
-        mConnection.getModuleLoader(mModuleComponentName).loadModule(result -> {
-            assertNotNull(result);
-            loadingWaiter.notifyCalled();
-        });
-        loadingWaiter.waitForCallback();
+        ModuleLoader moduleLoader = new ModuleLoader(mModuleComponentName);
+        moduleLoader.loadModule();
+
+        CallbackHelper onLoaded = new CallbackHelper();
+        moduleLoader.addCallbackAndIncrementUseCount(
+           result -> {
+              assertNotNull(result);
+              onLoaded.notifyCalled();
+           });
+
+        onLoaded.waitForCallback();
+    }
+
+    /**
+     * Test the ModuleLoader correctly update module usage counter.
+     */
+    @Test
+    @SmallTest
+    @Features.EnableFeatures(ChromeFeatureList.CCT_MODULE_CACHE)
+    public void testModuleUseCounter() throws TimeoutException, InterruptedException {
+        final int callbacksNumber = 3;
+        ModuleLoader moduleLoader = new ModuleLoader(mModuleComponentName);
+
+        CallbackHelper onLoaded = new CallbackHelper();
+
+        // Test we correctly unregister callbacks which were never notified.
+        List<Callback<ModuleEntryPoint>> unusedCallbacks = new ArrayList<>();
+        for (int i = 0; i < callbacksNumber; i++) {
+            unusedCallbacks.add(result -> {});
+            moduleLoader.addCallbackAndIncrementUseCount(unusedCallbacks.get(i));
+        }
+        // module has not been loaded, therefore there is no usage
+        assertEquals(0, moduleLoader.getModuleUseCount());
+
+        // unregister all callbacks so they should not increment module usage
+        for (int i = 0; i < callbacksNumber; i++) {
+            moduleLoader.removeCallbackAndDecrementUseCount(unusedCallbacks.get(i));
+        }
+
+        assertEquals(0, moduleLoader.getModuleUseCount());
+
+        moduleLoader.loadModule();
+
+        List<Callback<ModuleEntryPoint>> callbacks = new ArrayList<>();
+        // register callbacks and wait for the notification when module is loaded
+        for (int i = 0; i < callbacksNumber; i++) {
+            callbacks.add(result -> onLoaded.notifyCalled());
+            moduleLoader.addCallbackAndIncrementUseCount(callbacks.get(i));
+        }
+        onLoaded.waitForCallback(0, callbacksNumber);
+
+        assertEquals(callbacksNumber, moduleLoader.getModuleUseCount());
+
+        // unregister already notified callbacks
+        for (int i = 0; i < callbacksNumber; i++) {
+            moduleLoader.removeCallbackAndDecrementUseCount(callbacks.get(i));
+        }
+        assertEquals(0, moduleLoader.getModuleUseCount());
     }
 
     /**