Add tab persistence support for CustomTabs.

BUG=606513

Review-Url: https://codereview.chromium.org/2296833002
Cr-Commit-Position: refs/heads/master@{#415786}

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

@@ -43,6 +43,17 @@ public class ApiCompatibilityUtils {
     private ApiCompatibilityUtils() {
     }
 
+    /**
+     * @see Long#compare(long, long)
+     */
+    public static int compareLong(long lhs, long rhs) {
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.KITKAT) {
+            return Long.compare(lhs, rhs);
+        } else {
+            return lhs < rhs ? -1 : (lhs == rhs ? 0 : 1);
+        }
+    }
+
     /**
      * Returns true if view's layout direction is right-to-left.
      *

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

@@ -58,7 +58,6 @@ import org.chromium.chrome.browser.tabmodel.TabModel.TabLaunchType;
 import org.chromium.chrome.browser.tabmodel.TabModelObserver;
 import org.chromium.chrome.browser.tabmodel.TabModelSelectorImpl;
 import org.chromium.chrome.browser.tabmodel.TabPersistencePolicy;
-import org.chromium.chrome.browser.tabmodel.TabbedModeTabPersistencePolicy;
 import org.chromium.chrome.browser.toolbar.ToolbarControlContainer;
 import org.chromium.chrome.browser.util.ColorUtils;
 import org.chromium.chrome.browser.util.FeatureUtilities;
@@ -264,11 +263,8 @@ public class CustomTabActivity extends ChromeActivity {
     @Override
     public void postInflationStartup() {
         super.postInflationStartup();
-        // TODO(tedchoc): Replace with an actual CCT implementation for the persistence policy.
-        //                Currently, this relies on checks for CUSTOM_TABS_SELECTOR_INDEX in the
-        //                main TabPersitentStore that ignores certain actions.
-        TabPersistencePolicy persistencePolicy = new TabbedModeTabPersistencePolicy(
-                this, TabModelSelectorImpl.CUSTOM_TABS_SELECTOR_INDEX);
+        TabPersistencePolicy persistencePolicy = new CustomTabTabPersistencePolicy(
+                getTaskId(), getSavedInstanceState() != null);
         setTabModelSelector(new TabModelSelectorImpl(
                 this, persistencePolicy, getWindowAndroid(), false));
 

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

+// Copyright 2016 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.customtabs;
+
+import android.app.Activity;
+import android.os.AsyncTask;
+import android.os.StrictMode;
+import android.util.Pair;
+import android.util.SparseBooleanArray;
+
+import org.chromium.base.ApiCompatibilityUtils;
+import org.chromium.base.ApplicationStatus;
+import org.chromium.base.Callback;
+import org.chromium.base.Log;
+import org.chromium.base.StreamUtil;
+import org.chromium.base.ThreadUtils;
+import org.chromium.chrome.browser.TabState;
+import org.chromium.chrome.browser.compositor.layouts.content.TabContentManager;
+import org.chromium.chrome.browser.tabmodel.TabModel;
+import org.chromium.chrome.browser.tabmodel.TabModelSelector;
+import org.chromium.chrome.browser.tabmodel.TabPersistencePolicy;
+import org.chromium.chrome.browser.tabmodel.TabPersistentStore;
+
+import java.io.BufferedInputStream;
+import java.io.DataInputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+
+import javax.annotation.Nullable;
+
+/**
+ * Handles the Custom Tab specific behaviors of tab persistence.
+ */
+public class CustomTabTabPersistencePolicy implements TabPersistencePolicy {
+
+    static final String SAVED_STATE_DIRECTORY = "custom_tabs";
+
+    /** Threshold where old state files should be deleted (30 days). */
+    protected static final long STATE_EXPIRY_THRESHOLD = 30L * 24 * 60 * 60 * 1000;
+
+    /** Maximum number of state files before we should start deleting old ones. */
+    protected static final int MAXIMUM_STATE_FILES = 30;
+
+    private static final String TAG = "tabmodel";
+
+    /** Prevents two state directories from getting created simultaneously. */
+    private static final Object DIR_CREATION_LOCK = new Object();
+
+    /**
+     * Prevents two clean up tasks from getting created simultaneously. Also protects against
+     * incorrectly interleaving create/run/cancel on the task.
+     */
+    private static final Object CLEAN_UP_TASK_LOCK = new Object();
+
+    private static File sStateDirectory;
+    private static AsyncTask<Void, Void, Void> sCleanupTask;
+
+    /**
+     * The folder where the state should be saved to.
+     * @return A file representing the directory that contains TabModelSelector states.
+     */
+    public static File getOrCreateCustomTabModeStateDirectory() {
+        synchronized (DIR_CREATION_LOCK) {
+            if (sStateDirectory == null) {
+                sStateDirectory = new File(
+                        TabPersistentStore.getOrCreateBaseStateDirectory(), SAVED_STATE_DIRECTORY);
+                StrictMode.ThreadPolicy oldPolicy = StrictMode.allowThreadDiskReads();
+                StrictMode.allowThreadDiskWrites();
+                try {
+                    if (!sStateDirectory.exists() && !sStateDirectory.mkdirs()) {
+                        Log.e(TAG, "Failed to create state folder: " + sStateDirectory);
+                    }
+                } finally {
+                    StrictMode.setThreadPolicy(oldPolicy);
+                }
+            }
+        }
+        return sStateDirectory;
+    }
+
+    private final int mTaskId;
+    private final boolean mShouldRestore;
+
+    private AsyncTask<Void, Void, Void> mInitializationTask;
+    private boolean mDestroyed;
+
+    /**
+     * Constructs a persistence policy for a given Custom Tab.
+     *
+     * @param taskId The task ID that the owning Custom Tab is in.
+     * @param shouldRestore Whether an attempt to restore tab state information should be done on
+     *                      startup.
+     */
+    public CustomTabTabPersistencePolicy(int taskId, boolean shouldRestore) {
+        mTaskId = taskId;
+        mShouldRestore = shouldRestore;
+    }
+
+    @Override
+    public File getOrCreateStateDirectory() {
+        return getOrCreateCustomTabModeStateDirectory();
+    }
+
+    @Override
+    public String getStateFileName() {
+        return TabPersistentStore.getStateFileName(Integer.toString(mTaskId));
+    }
+
+    @Override
+    @Nullable
+    public String getStateToBeMergedFileName() {
+        return null;
+    }
+
+    @Override
+    public boolean performInitialization(Executor executor) {
+        mInitializationTask = new AsyncTask<Void, Void, Void>() {
+            @Override
+            protected Void doInBackground(Void... params) {
+                File stateDir = getOrCreateStateDirectory();
+                File metadataFile = new File(stateDir, getStateFileName());
+                if (metadataFile.exists()) {
+                    if (mShouldRestore) {
+                        if (!metadataFile.setLastModified(System.currentTimeMillis())) {
+                            Log.e(TAG, "Unable to update last modified time: " + metadataFile);
+                        }
+                    } else {
+                        if (!metadataFile.delete()) {
+                            Log.e(TAG, "Failed to delete file: " + metadataFile);
+                        }
+                    }
+                }
+                return null;
+            }
+        }.executeOnExecutor(executor);
+
+        return true;
+    }
+
+    @Override
+    public void waitForInitializationToFinish() {
+        if (mInitializationTask == null) return;
+        try {
+            mInitializationTask.get();
+        } catch (InterruptedException | ExecutionException e) {
+            // Ignore and proceed.
+        }
+    }
+
+    @Override
+    public boolean isMergeInProgress() {
+        return false;
+    }
+
+    @Override
+    public void setMergeInProgress(boolean isStarted) {
+        assert false : "Merge not supported in Custom Tabs";
+    }
+
+    @Override
+    public void cancelCleanupInProgress() {
+        synchronized (CLEAN_UP_TASK_LOCK) {
+            if (sCleanupTask != null) sCleanupTask.cancel(true);
+        }
+    }
+
+    @Override
+    public void cleanupUnusedFiles(Callback<List<String>> filesToDelete) {
+        synchronized (CLEAN_UP_TASK_LOCK) {
+            if (sCleanupTask != null) sCleanupTask.cancel(true);
+            sCleanupTask = new CleanUpTabStateDataTask(filesToDelete);
+            sCleanupTask.executeOnExecutor(AsyncTask.THREAD_POOL_EXECUTOR);
+        }
+    }
+
+    @Override
+    public void setTabContentManager(TabContentManager cache) {
+    }
+
+    @Override
+    public void destroy() {
+        mDestroyed = true;
+    }
+
+    /**
+     * Given a list of metadata files, determine which are applicable for deletion based on the
+     * deletion strategy of Custom Tabs.
+     *
+     * @param currentTimeMillis The current time in milliseconds
+     *                          ({@link System#currentTimeMillis()}.
+     * @param allMetadataFiles The complete list of all metadata files to check.
+     * @return The list of metadata files that are applicable for deletion.
+     */
+    protected static List<File> getMetadataFilesForDeletion(
+            long currentTimeMillis, List<File> allMetadataFiles) {
+        Collections.sort(allMetadataFiles, new Comparator<File>() {
+            @Override
+            public int compare(File lhs, File rhs) {
+                long lhsModifiedTime = lhs.lastModified();
+                long rhsModifiedTime = rhs.lastModified();
+
+                // Sort such that older files (those with an lower timestamp number) are at the
+                // end of the sorted listed.
+                return ApiCompatibilityUtils.compareLong(rhsModifiedTime, lhsModifiedTime);
+            }
+        });
+
+        List<File> stateFilesApplicableForDeletion = new ArrayList<File>();
+        for (int i = 0; i < allMetadataFiles.size(); i++) {
+            File file = allMetadataFiles.get(i);
+            long fileAge = currentTimeMillis - file.lastModified();
+            if (i >= MAXIMUM_STATE_FILES || fileAge >= STATE_EXPIRY_THRESHOLD) {
+                stateFilesApplicableForDeletion.add(file);
+            }
+        }
+        return stateFilesApplicableForDeletion;
+    }
+
+    /**
+     * Get all current Tab IDs used by the specified activity.
+     *
+     * @param activity The activity whose tab IDs are to be collected from.
+     * @param tabIds Where the tab IDs should be added to.
+     */
+    private static void getAllTabIdsForActivity(CustomTabActivity activity, Set<Integer> tabIds) {
+        if (activity == null) return;
+        TabModelSelector selector = activity.getTabModelSelector();
+        if (selector == null) return;
+        List<TabModel> models = selector.getModels();
+        for (int i = 0; i < models.size(); i++) {
+            TabModel model = models.get(i);
+            for (int j = 0; j < model.getCount(); j++) {
+                tabIds.add(model.getTabAt(j).getId());
+            }
+        }
+    }
+
+    /**
+     * Gathers all of the tab IDs and task IDs for all currently live Custom Tabs.
+     *
+     * @param liveTabIds Where tab IDs will be added.
+     * @param liveTaskIds Where task IDs will be added.
+     */
+    protected static void getAllLiveTabAndTaskIds(
+            Set<Integer> liveTabIds, Set<Integer> liveTaskIds) {
+        ThreadUtils.assertOnUiThread();
+
+        List<WeakReference<Activity>> activities = ApplicationStatus.getRunningActivities();
+        for (int i = 0; i < activities.size(); i++) {
+            Activity activity = activities.get(i).get();
+            if (activity == null) continue;
+            if (!(activity instanceof CustomTabActivity)) continue;
+            getAllTabIdsForActivity((CustomTabActivity) activity, liveTabIds);
+            liveTaskIds.add(activity.getTaskId());
+        }
+    }
+
+    private class CleanUpTabStateDataTask extends AsyncTask<Void, Void, Void> {
+        private final Callback<List<String>> mFilesToDeleteCallback;
+
+        private Set<Integer> mUnreferencedTabIds;
+        private List<File> mDeletableMetadataFiles;
+        private Map<File, SparseBooleanArray> mTabIdsByMetadataFile;
+
+        CleanUpTabStateDataTask(Callback<List<String>> filesToDelete) {
+            mFilesToDeleteCallback = filesToDelete;
+        }
+
+        @Override
+        protected Void doInBackground(Void... voids) {
+            if (mDestroyed) return null;
+
+            mTabIdsByMetadataFile = new HashMap<>();
+            mUnreferencedTabIds = new HashSet<>();
+
+            File[] stateFiles = getOrCreateStateDirectory().listFiles();
+            if (stateFiles == null) return null;
+
+            Set<Integer> allTabIds = new HashSet<>();
+            Set<Integer> allReferencedTabIds = new HashSet<>();
+            List<File> metadataFiles = new ArrayList<>();
+            for (File file : stateFiles) {
+                if (TabPersistentStore.isStateFile(file.getName())) {
+                    metadataFiles.add(file);
+
+                    SparseBooleanArray tabIds = new SparseBooleanArray();
+                    mTabIdsByMetadataFile.put(file, tabIds);
+                    getTabsFromStateFile(tabIds, file);
+                    for (int i = 0; i < tabIds.size(); i++) {
+                        allReferencedTabIds.add(tabIds.keyAt(i));
+                    }
+                    continue;
+                }
+
+                Pair<Integer, Boolean> tabInfo = TabState.parseInfoFromFilename(file.getName());
+                if (tabInfo == null) continue;
+                allTabIds.add(tabInfo.first);
+            }
+
+            mUnreferencedTabIds.addAll(allTabIds);
+            mUnreferencedTabIds.removeAll(allReferencedTabIds);
+
+            mDeletableMetadataFiles = getMetadataFilesForDeletion(
+                    System.currentTimeMillis(), metadataFiles);
+            return null;
+        }
+
+        @Override
+        protected void onPostExecute(Void unused) {
+            List<String> filesToDelete = new ArrayList<>();
+            if (mDestroyed) {
+                mFilesToDeleteCallback.onResult(filesToDelete);
+                return;
+            }
+
+            if (mUnreferencedTabIds.isEmpty() && mDeletableMetadataFiles.isEmpty()) {
+                mFilesToDeleteCallback.onResult(filesToDelete);
+                return;
+            }
+
+            Set<Integer> liveTabIds = new HashSet<>();
+            Set<Integer> liveTaskIds = new HashSet<>();
+            getAllLiveTabAndTaskIds(liveTabIds, liveTaskIds);
+
+            for (Integer unreferencedTabId : mUnreferencedTabIds) {
+                // Ignore tabs that are referenced by live activities as they might not have been
+                // able to write out their state yet.
+                if (liveTabIds.contains(unreferencedTabId)) continue;
+
+                // The tab state is not referenced by any current activities or any metadata files,
+                // so mark it for deletion.
+                filesToDelete.add(TabState.getTabStateFilename(unreferencedTabId, false));
+            }
+
+            for (int i = 0; i < mDeletableMetadataFiles.size(); i++) {
+                File metadataFile = mDeletableMetadataFiles.get(i);
+                String id = TabPersistentStore.getStateFileUniqueId(metadataFile.getName());
+                try {
+                    int taskId = Integer.parseInt(id);
+
+                    // Ignore the metadata file if it belongs to a currently live CustomTabActivity.
+                    if (liveTaskIds.contains(taskId)) continue;
+
+                    filesToDelete.add(metadataFile.getName());
+
+                    SparseBooleanArray unusedTabIds = mTabIdsByMetadataFile.get(metadataFile);
+                    if (unusedTabIds == null) continue;
+                    for (int j = 0; j < unusedTabIds.size(); j++) {
+                        filesToDelete.add(TabState.getTabStateFilename(
+                                unusedTabIds.keyAt(j), false));
+                    }
+                } catch (NumberFormatException ex) {
+                    assert false : "Unexpected tab metadata file found: " + metadataFile.getName();
+                    continue;
+                }
+            }
+
+            mFilesToDeleteCallback.onResult(filesToDelete);
+        }
+
+        private void getTabsFromStateFile(SparseBooleanArray tabIds, File metadataFile) {
+            DataInputStream stream = null;
+            try {
+                stream = new DataInputStream(
+                        new BufferedInputStream(new FileInputStream(metadataFile)));
+                TabPersistentStore.readSavedStateFile(stream, null, tabIds, false);
+            } catch (Exception e) {
+                Log.e(TAG, "Unable to read state for " + metadataFile.getName() + ": " + e);
+            } finally {
+                StreamUtil.closeQuietly(stream);
+            }
+        }
+    }
+}

chrome/android/java/src/org/chromium/chrome/browser/tab/Tab.java

@@ -1627,6 +1627,7 @@ public class Tab implements ViewGroup.OnHierarchyChangeListener,
 
         reparentingParams.finalizeTabReparenting();
         mIsDetachedForReparenting = false;
+        mIsTabStateDirty = true;
 
         for (TabObserver observer : mObservers) {
             observer.onReparentingFinished(this);

chrome/android/java/src/org/chromium/chrome/browser/tabmodel/TabModelSelectorImpl.java

@@ -227,6 +227,17 @@ public class TabModelSelectorImpl extends TabModelSelectorBase implements TabMod
         };
     }
 
+    /**
+     * Exposed to allow tests to initialize the selector with different tab models.
+     * @param normalModel The normal tab model.
+     * @param incognitoModel The incognito tab model.
+     */
+    @VisibleForTesting
+    public void initializeForTesting(TabModel normalModel, TabModel incognitoModel) {
+        initialize(isIncognitoSelected(), normalModel, incognitoModel);
+        mActiveState = true;
+    }
+
     @Override
     public void setCloseAllTabsDelegate(CloseAllTabsDelegate delegate) {
         mCloseAllTabsDelegate = delegate;

chrome/android/java/src/org/chromium/chrome/browser/tabmodel/TabPersistencePolicy.java

@@ -43,18 +43,21 @@ public interface TabPersistencePolicy {
     String getStateToBeMergedFileName();
 
     /**
-     * Performs any necessary migration required before accessing the tab information.
+     * Performs any necessary initialization required before accessing the tab information.  This
+     * can include cleanups or migrations that must occur before the tab state information can be
+     * read reliably.
      *
      * @param executor The executor that any asynchronous tasks should be run on.
-     * @return Whether any migration is necessary.
+     * @return Whether any blocking initialization is necessary.
      */
-    boolean performMigration(Executor executor);
+    boolean performInitialization(Executor executor);
 
     /**
-     * Waits for the task that migrates all state files to their new location to
-     * finish.
+     * Waits for the any pending initialization to finish.
+     *
+     * @see #performInitialization(Executor)
      */
-    void waitForMigrationToFinish();
+    void waitForInitializationToFinish();
 
     /**
      * @return Whether a merge is currently in progress.

chrome/android/java/src/org/chromium/chrome/browser/tabmodel/TabPersistentStore.java

@@ -219,11 +219,12 @@ public class TabPersistentStore extends TabPersister {
         mPreferences = ContextUtils.getAppSharedPreferences();
 
         assert isStateFile(policy.getStateFileName()) : "State file name is not valid";
-        boolean needsMigration = mPersistencePolicy.performMigration(AsyncTask.SERIAL_EXECUTOR);
+        boolean needsInitialization = mPersistencePolicy.performInitialization(
+                AsyncTask.SERIAL_EXECUTOR);
 
         if (mPersistencePolicy.isMergeInProgress()) return;
 
-        Executor executor = needsMigration
+        Executor executor = needsInitialization
                 ? AsyncTask.SERIAL_EXECUTOR : AsyncTask.THREAD_POOL_EXECUTOR;
 
         mPrefetchTabListTask =
@@ -247,7 +248,7 @@ public class TabPersistentStore extends TabPersister {
      */
     @VisibleForTesting
     public void waitForMigrationToFinish() {
-        mPersistencePolicy.waitForMigrationToFinish();
+        mPersistencePolicy.waitForInitializationToFinish();
     }
 
     /**
@@ -998,7 +999,7 @@ public class TabPersistentStore extends TabPersister {
      * @param forMerge Whether this state file was read as part of a merge.
      * @return The next available tab ID based on the maximum ID referenced in this state file.
      */
-    protected static int readSavedStateFile(
+    public static int readSavedStateFile(
             DataInputStream stream, @Nullable OnTabStateReadCallback callback,
             @Nullable SparseBooleanArray tabIds, boolean forMerge) throws IOException {
         if (stream == null) return 0;
@@ -1366,10 +1367,20 @@ public class TabPersistentStore extends TabPersister {
         return TabPersistencePolicy.SAVED_STATE_FILE_PREFIX + uniqueId;
     }
 
+    /**
+     * Parses the state file name and returns the unique ID encoded into it.
+     * @param stateFileName The state file name to be parsed.
+     * @return The unique ID used when generating the file name.
+     */
+    public static String getStateFileUniqueId(String stateFileName) {
+        assert isStateFile(stateFileName);
+        return stateFileName.substring(TabPersistencePolicy.SAVED_STATE_FILE_PREFIX.length());
+    }
+
     /**
      * @return Whether the specified filename matches the expected pattern of the tab state files.
      */
-    private static boolean isStateFile(String fileName) {
+    public static boolean isStateFile(String fileName) {
         return fileName.startsWith(TabPersistencePolicy.SAVED_STATE_FILE_PREFIX);
     }
 

chrome/android/java/src/org/chromium/chrome/browser/tabmodel/TabbedModeTabPersistencePolicy.java

@@ -143,7 +143,7 @@ public class TabbedModeTabPersistencePolicy implements TabPersistencePolicy {
     }
 
     @Override
-    public boolean performMigration(Executor executor) {
+    public boolean performInitialization(Executor executor) {
         ThreadUtils.assertOnUiThread();
 
         final boolean hasRunLegacyMigration =
@@ -291,7 +291,7 @@ public class TabbedModeTabPersistencePolicy implements TabPersistencePolicy {
     }
 
     @Override
-    public void waitForMigrationToFinish() {
+    public void waitForInitializationToFinish() {
         if (sMigrationTask == null) return;
         try {
             sMigrationTask.get();

chrome/android/java_sources.gni

@@ -269,6 +269,7 @@ chrome_java_sources = [
   "java/src/org/chromium/chrome/browser/customtabs/CustomTabObserver.java",
   "java/src/org/chromium/chrome/browser/customtabs/CustomTabsConnection.java",
   "java/src/org/chromium/chrome/browser/customtabs/CustomTabsConnectionService.java",
+  "java/src/org/chromium/chrome/browser/customtabs/CustomTabTabPersistencePolicy.java",
   "java/src/org/chromium/chrome/browser/customtabs/RequestThrottler.java",
   "java/src/org/chromium/chrome/browser/customtabs/SeparateTaskCustomTabActivity.java",
   "java/src/org/chromium/chrome/browser/customtabs/SeparateTaskCustomTabActivity0.java",
@@ -1157,6 +1158,7 @@ chrome_test_java_sources = [
   "javatests/src/org/chromium/chrome/browser/customtabs/CustomTabExternalNavigationTest.java",
   "javatests/src/org/chromium/chrome/browser/customtabs/CustomTabsConnectionTest.java",
   "javatests/src/org/chromium/chrome/browser/customtabs/CustomTabsTestUtils.java",
+  "javatests/src/org/chromium/chrome/browser/customtabs/CustomTabTabPersistencePolicyTest.java",
   "javatests/src/org/chromium/chrome/browser/customtabs/RequestThrottlerTest.java",
   "javatests/src/org/chromium/chrome/browser/document/LauncherActivityTest.java",
   "javatests/src/org/chromium/chrome/browser/dom_distiller/DistillabilityServiceTest.java",

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

+// Copyright 2016 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.customtabs;
+
+import static android.test.MoreAsserts.assertContentsInAnyOrder;
+import static android.test.MoreAsserts.assertEmpty;
+
+import android.app.Activity;
+import android.os.AsyncTask;
+import android.test.InstrumentationTestCase;
+import android.test.UiThreadTest;
+import android.test.suitebuilder.annotation.MediumTest;
+import android.test.suitebuilder.annotation.SmallTest;
+
+import org.chromium.base.ActivityState;
+import org.chromium.base.ApplicationStatus;
+import org.chromium.base.Callback;
+import org.chromium.base.ContextUtils;
+import org.chromium.base.StreamUtil;
+import org.chromium.base.ThreadUtils;
+import org.chromium.base.test.util.AdvancedMockContext;
+import org.chromium.base.test.util.Feature;
+import org.chromium.chrome.browser.ChromeTabbedActivity;
+import org.chromium.chrome.browser.TabState;
+import org.chromium.chrome.browser.compositor.layouts.content.TabContentManager;
+import org.chromium.chrome.browser.tab.Tab;
+import org.chromium.chrome.browser.tabmodel.TabModelSelector;
+import org.chromium.chrome.browser.tabmodel.TabModelSelectorImpl;
+import org.chromium.chrome.browser.tabmodel.TabPersistencePolicy;
+import org.chromium.chrome.browser.tabmodel.TabPersistentStore;
+import org.chromium.chrome.browser.tabmodel.TestTabModelDirectory;
+import org.chromium.chrome.test.util.browser.tabmodel.MockTabModel;
+import org.chromium.content.browser.test.util.CallbackHelper;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
+
+import javax.annotation.Nullable;
+
+/**
+ * Tests for the Custom Tab persistence logic.
+ */
+public class CustomTabTabPersistencePolicyTest extends InstrumentationTestCase {
+
+    private TestTabModelDirectory mMockDirectory;
+    private AdvancedMockContext mAppContext;
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+
+        mAppContext = new AdvancedMockContext(
+                getInstrumentation().getTargetContext().getApplicationContext());
+        ContextUtils.initApplicationContextForTests(mAppContext);
+
+        mMockDirectory = new TestTabModelDirectory(
+                mAppContext, "CustomTabTabPersistencePolicyTest",
+                CustomTabTabPersistencePolicy.SAVED_STATE_DIRECTORY);
+        TabPersistentStore.setBaseStateDirectoryForTests(mMockDirectory.getBaseDirectory());
+    }
+
+    @Override
+    public void tearDown() throws Exception {
+        mMockDirectory.tearDown();
+
+        List<WeakReference<Activity>> activities = ApplicationStatus.getRunningActivities();
+        for (int i = 0; i < activities.size(); i++) {
+            Activity activity = activities.get(i).get();
+            if (activity == null) continue;
+            ApplicationStatus.onStateChangeForTesting(activity, ActivityState.DESTROYED);
+        }
+        super.tearDown();
+    }
+
+    @Feature("TabPersistentStore")
+    @SmallTest
+    public void testDeletableMetadataSelection_NoFiles() {
+        List<File> deletableFiles = CustomTabTabPersistencePolicy.getMetadataFilesForDeletion(
+                System.currentTimeMillis(), new ArrayList<File>());
+        assertEmpty(deletableFiles);
+    }
+
+    @Feature("TabPersistentStore")
+    @SmallTest
+    public void testDeletableMetadataSelection_MaximumValidFiles() {
+        long currentTime = System.currentTimeMillis();
+
+        // Test the maximum allowed number of state files where they are all valid in terms of age.
+        List<File> filesToTest = new ArrayList<>();
+        filesToTest.addAll(generateMaximumStateFiles(currentTime));
+        List<File> deletableFiles = CustomTabTabPersistencePolicy.getMetadataFilesForDeletion(
+                currentTime, filesToTest);
+        assertEmpty(deletableFiles);
+    }
+
+    @Feature("TabPersistentStore")
+    @SmallTest
+    public void testDeletableMetadataSelection_ExceedsMaximumValidFiles() {
+        long currentTime = System.currentTimeMillis();
+
+        // Test where we exceed the maximum number of allowed state files and ensure it chooses the
+        // older file to delete.
+        List<File> filesToTest = new ArrayList<>();
+        filesToTest.addAll(generateMaximumStateFiles(currentTime));
+        File slightlyOlderFile = buildTestFile("slightlyolderfile", currentTime - 1L);
+        // Insert it into the middle just to ensure it is not picking the last file.
+        filesToTest.add(filesToTest.size() / 2, slightlyOlderFile);
+        List<File> deletableFiles = CustomTabTabPersistencePolicy.getMetadataFilesForDeletion(
+                currentTime, filesToTest);
+        assertContentsInAnyOrder(deletableFiles, slightlyOlderFile);
+    }
+
+    @Feature("TabPersistentStore")
+    @SmallTest
+    public void testDeletableMetadataSelection_ExceedExpiryThreshold() {
+        long currentTime = System.currentTimeMillis();
+
+        // Ensure that files that exceed the allowed time threshold are removed regardless of the
+        // number of possible files.
+        List<File> filesToTest = new ArrayList<>();
+        File expiredFile = buildTestFile("expired_file",
+                currentTime - CustomTabTabPersistencePolicy.STATE_EXPIRY_THRESHOLD);
+        filesToTest.add(expiredFile);
+        List<File> deletableFiles = CustomTabTabPersistencePolicy.getMetadataFilesForDeletion(
+                currentTime, filesToTest);
+        assertContentsInAnyOrder(deletableFiles, expiredFile);
+    }
+
+    /**
+     * Test to ensure that an existing metadata files are deleted if no restore is requested.
+     */
+    @Feature("TabPersistentStore")
+    @MediumTest
+    public void testExistingMetadataFileDeletedIfNoRestore() throws Exception {
+        File baseStateDirectory = TabPersistentStore.getOrCreateBaseStateDirectory();
+        assertNotNull(baseStateDirectory);
+
+        CustomTabTabPersistencePolicy policy = new CustomTabTabPersistencePolicy(7, false);
+        File stateDirectory = policy.getOrCreateStateDirectory();
+        assertNotNull(stateDirectory);
+
+        String stateFileName = policy.getStateFileName();
+        File existingStateFile = new File(stateDirectory, stateFileName);
+        assertTrue(existingStateFile.createNewFile());
+
+        assertTrue(existingStateFile.exists());
+        policy.performInitialization(AsyncTask.SERIAL_EXECUTOR);
+        policy.waitForInitializationToFinish();
+        assertFalse(existingStateFile.exists());
+    }
+
+    /**
+     * Test the logic that gets all the live tab and task IDs.
+     */
+    @Feature("TabPersistentStore")
+    @SmallTest
+    @UiThreadTest
+    public void testGettingTabAndTaskIds() {
+        Set<Integer> tabIds = new HashSet<>();
+        Set<Integer> taskIds = new HashSet<>();
+        CustomTabTabPersistencePolicy.getAllLiveTabAndTaskIds(tabIds, taskIds);
+        assertEmpty(tabIds);
+        assertEmpty(taskIds);
+
+        tabIds.clear();
+        taskIds.clear();
+
+        CustomTabActivity cct1 = buildTestCustomTabActivity(1, new int[] {4, 8, 9}, null);
+        ApplicationStatus.onStateChangeForTesting(cct1, ActivityState.CREATED);
+
+        CustomTabActivity cct2 = buildTestCustomTabActivity(5, new int[] {458}, new int[] {9878});
+        ApplicationStatus.onStateChangeForTesting(cct2, ActivityState.CREATED);
+
+        // Add a tabbed mode activity to ensure that its IDs are not included in the returned CCT
+        // ID sets.
+        final TabModelSelectorImpl tabbedSelector =
+                buildTestTabModelSelector(new int[] {12121212}, new int[] {1515151515});
+        ChromeTabbedActivity tabbedActivity = new ChromeTabbedActivity() {
+            @Override
+            public int getTaskId() {
+                return 888;
+            }
+
+            @Override
+            public TabModelSelector getTabModelSelector() {
+                return tabbedSelector;
+            }
+        };
+        ApplicationStatus.onStateChangeForTesting(tabbedActivity, ActivityState.CREATED);
+
+        CustomTabTabPersistencePolicy.getAllLiveTabAndTaskIds(tabIds, taskIds);
+        assertContentsInAnyOrder(tabIds, 4, 8, 9, 458, 9878);
+        assertContentsInAnyOrder(taskIds, 1, 5);
+    }
+
+    /**
+     * Test the full cleanup task path that determines what files are eligible for deletion.
+     */
+    @Feature("TabPersistentStore")
+    @MediumTest
+    public void testCleanupTask() throws Exception {
+        File baseStateDirectory = TabPersistentStore.getOrCreateBaseStateDirectory();
+        assertNotNull(baseStateDirectory);
+
+        CustomTabTabPersistencePolicy policy = new CustomTabTabPersistencePolicy(2, false);
+        File stateDirectory = policy.getOrCreateStateDirectory();
+        assertNotNull(stateDirectory);
+
+        final AtomicReference<List<String>> filesToDelete = new AtomicReference<>();
+        final CallbackHelper callbackSignal = new CallbackHelper();
+        Callback<List<String>> filesToDeleteCallback = new Callback<List<String>>() {
+            @Override
+            public void onResult(List<String> fileNames) {
+                filesToDelete.set(fileNames);
+                callbackSignal.notifyCalled();
+            }
+        };
+
+        // Test when no files have been created.
+        policy.cleanupUnusedFiles(filesToDeleteCallback);
+        callbackSignal.waitForCallback(0);
+        assertEmpty(filesToDelete.get());
+
+        // Create an unreferenced tab state file and ensure it is marked for deletion.
+        File tab999File = TabState.getTabStateFile(stateDirectory, 999, false);
+        assertTrue(tab999File.createNewFile());
+        policy.cleanupUnusedFiles(filesToDeleteCallback);
+        callbackSignal.waitForCallback(1);
+        assertContentsInAnyOrder(filesToDelete.get(), tab999File.getName());
+
+        // Reference the tab state file and ensure it is no longer marked for deletion.
+        CustomTabActivity cct1 = buildTestCustomTabActivity(1, new int[] {999}, null);
+        ApplicationStatus.onStateChangeForTesting(cct1, ActivityState.CREATED);
+        policy.cleanupUnusedFiles(filesToDeleteCallback);
+        callbackSignal.waitForCallback(2);
+        assertEmpty(filesToDelete.get());
+
+        // Create a tab model and associated tabs. Ensure it is not marked for deletion as it is
+        // new enough.
+        final TabModelSelectorImpl selectorImpl = buildTestTabModelSelector(
+                new int[] {111, 222, 333 }, null);
+        byte[] data = ThreadUtils.runOnUiThreadBlockingNoException(new Callable<byte[]>() {
+            @Override
+            public byte[] call() throws Exception {
+                return TabPersistentStore.serializeTabModelSelector(selectorImpl, null);
+            }
+        });
+        FileOutputStream fos = null;
+        File metadataFile = new File(
+                stateDirectory, TabPersistentStore.getStateFileName("3"));
+        try {
+            fos = new FileOutputStream(metadataFile);
+            fos.write(data);
+        } finally {
+            StreamUtil.closeQuietly(fos);
+        }
+        File tab111File = TabState.getTabStateFile(stateDirectory, 111, false);
+        assertTrue(tab111File.createNewFile());
+        File tab222File = TabState.getTabStateFile(stateDirectory, 222, false);
+        assertTrue(tab222File.createNewFile());
+        File tab333File = TabState.getTabStateFile(stateDirectory, 333, false);
+        assertTrue(tab333File.createNewFile());
+        policy.cleanupUnusedFiles(filesToDeleteCallback);
+        callbackSignal.waitForCallback(3);
+        assertEmpty(filesToDelete.get());
+
+        // Set the age of the metadata file to be past the expiration threshold and ensure it along
+        // with the associated tab files are marked for deletion.
+        assertTrue(metadataFile.setLastModified(1234));
+        policy.cleanupUnusedFiles(filesToDeleteCallback);
+        callbackSignal.waitForCallback(4);
+        assertContentsInAnyOrder(filesToDelete.get(), tab111File.getName(), tab222File.getName(),
+                tab333File.getName(), metadataFile.getName());
+    }
+
+    /**
+     * Ensure that the metadata file's last modified timestamp is updated on initialization.
+     */
+    @Feature("TabPersistentStore")
+    @MediumTest
+    public void testMetadataTimestampRefreshed() throws Exception {
+        File baseStateDirectory = TabPersistentStore.getOrCreateBaseStateDirectory();
+        assertNotNull(baseStateDirectory);
+
+        CustomTabTabPersistencePolicy policy = new CustomTabTabPersistencePolicy(2, true);
+        File stateDirectory = policy.getOrCreateStateDirectory();
+        assertNotNull(stateDirectory);
+
+        File metadataFile = new File(stateDirectory, policy.getStateFileName());
+        assertTrue(metadataFile.createNewFile());
+
+        long previousTimestamp =
+                System.currentTimeMillis() - CustomTabTabPersistencePolicy.STATE_EXPIRY_THRESHOLD;
+        assertTrue(metadataFile.setLastModified(previousTimestamp));
+
+        policy.performInitialization(AsyncTask.SERIAL_EXECUTOR);
+        policy.waitForInitializationToFinish();
+
+        assertTrue(metadataFile.lastModified() > previousTimestamp);
+    }
+
+    private static List<File> generateMaximumStateFiles(long currentTime) {
+        List<File> validFiles = new ArrayList<>();
+        for (int i = 0; i < CustomTabTabPersistencePolicy.MAXIMUM_STATE_FILES; i++) {
+            validFiles.add(buildTestFile("testfile" + i, currentTime));
+        }
+        return validFiles;
+    }
+
+    private static File buildTestFile(String filename, final long lastModifiedTime) {
+        return new File(filename) {
+            @Override
+            public long lastModified() {
+                return lastModifiedTime;
+            }
+        };
+    }
+
+    private static CustomTabActivity buildTestCustomTabActivity(
+            final int taskId, int[] normalTabIds, int[] incognitoTabIds) {
+        final TabModelSelectorImpl selectorImpl =
+                buildTestTabModelSelector(normalTabIds, incognitoTabIds);
+        return new CustomTabActivity() {
+            @Override
+            public int getTaskId() {
+                return taskId;
+            }
+
+            @Override
+            public TabModelSelectorImpl getTabModelSelector() {
+                return selectorImpl;
+            }
+        };
+    }
+
+    private static TabPersistencePolicy buildTestPersistencePolicy() {
+        return new TabPersistencePolicy() {
+            @Override
+            public void waitForInitializationToFinish() {
+            }
+
+            @Override
+            public void setTabContentManager(TabContentManager cache) {
+            }
+
+            @Override
+            public void setMergeInProgress(boolean isStarted) {
+            }
+
+            @Override
+            public boolean performInitialization(Executor executor) {
+                return false;
+            }
+
+            @Override
+            public boolean isMergeInProgress() {
+                return false;
+            }
+
+            @Override
+            @Nullable
+            public String getStateToBeMergedFileName() {
+                return null;
+            }
+
+            @Override
+            public String getStateFileName() {
+                return "testing_is_niftytown";
+            }
+
+            @Override
+            public File getOrCreateStateDirectory() {
+                return new File(
+                        TabPersistentStore.getOrCreateBaseStateDirectory(), "cct_tests_zor");
+            }
+
+            @Override
+            public void destroy() {
+            }
+
+            @Override
+            public void cleanupUnusedFiles(Callback<List<String>> filesToDelete) {
+            }
+
+            @Override
+            public void cancelCleanupInProgress() {
+            }
+        };
+    }
+
+    private static TabModelSelectorImpl buildTestTabModelSelector(
+            int[] normalTabIds, int[] incognitoTabIds) {
+        MockTabModel.MockTabModelDelegate tabModelDelegate =
+                new MockTabModel.MockTabModelDelegate() {
+                    @Override
+                    public Tab createTab(int id, boolean incognito) {
+                        return new Tab(id, incognito, null) {
+                            @Override
+                            public String getUrl() {
+                                return "https://www.google.com";
+                            }
+                        };
+                    }
+                };
+        final MockTabModel normalTabModel = new MockTabModel(false, tabModelDelegate);
+        if (normalTabIds != null) {
+            for (int tabId : normalTabIds) normalTabModel.addTab(tabId);
+        }
+        final MockTabModel incognitoTabModel = new MockTabModel(true, tabModelDelegate);
+        if (incognitoTabIds != null) {
+            for (int tabId : incognitoTabIds) incognitoTabModel.addTab(tabId);
+        }
+
+        CustomTabActivity activity = new CustomTabActivity();
+        ApplicationStatus.onStateChangeForTesting(activity, ActivityState.CREATED);
+        TabModelSelectorImpl selector = new TabModelSelectorImpl(
+                activity, buildTestPersistencePolicy(), null, false);
+        selector.initializeForTesting(normalTabModel, incognitoTabModel);
+        ApplicationStatus.onStateChangeForTesting(activity, ActivityState.DESTROYED);
+        return selector;
+    }
+}