Fixing Launcher preview leaking surface and memory

> Closing existing preview if a new request comes for same host token
> Closing in-memory icon db when closing preview
> Removing unnecessary wait blocks on UI thread and rendering
  view asynchronously
> Fixing preview loading failing on LauncherAppState access

Bug: 186712316
Bug: 187140897
Test: Manual
Change-Id: I045930b007e5dc015320224a197eee20a8354d17
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index b6cc6d6..dabbdd3 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -81,6 +81,8 @@
 
     public LauncherAppState(Context context) {
         this(context, LauncherFiles.APP_ICONS_DB);
+        Log.v(Launcher.TAG, "LauncherAppState initiated");
+        Preconditions.assertUIThread();
 
         mInvariantDeviceProfile.addOnChangeListener(idp -> refreshAndReloadLauncher());
 
@@ -132,8 +134,6 @@
     }
 
     public LauncherAppState(Context context, @Nullable String iconCacheFileName) {
-        Log.v(Launcher.TAG, "LauncherAppState initiated");
-        Preconditions.assertUIThread();
         mContext = context;
 
         mInvariantDeviceProfile = InvariantDeviceProfile.INSTANCE.get(context);
@@ -142,6 +142,7 @@
                 iconCacheFileName, mIconProvider);
         mWidgetCache = new WidgetPreviewLoader(mContext, mIconCache);
         mModel = new LauncherModel(context, this, mIconCache, new AppFilter(mContext));
+        mOnTerminateCallback.add(mIconCache::close);
     }
 
     private void onNotificationSettingsChanged(boolean areNotificationDotsEnabled) {
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
index a03e48d..60a1732 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
@@ -4,6 +4,7 @@
 import static com.android.launcher3.util.Themes.KEY_THEMED_ICONS;
 import static com.android.launcher3.util.Themes.isThemedIconEnabled;
 
+import android.annotation.TargetApi;
 import android.content.ContentProvider;
 import android.content.ContentValues;
 import android.content.pm.PackageManager;
@@ -12,14 +13,24 @@
 import android.database.MatrixCursor;
 import android.net.Uri;
 import android.os.Binder;
+import android.os.Build;
 import android.os.Bundle;
+import android.os.Handler;
+import android.os.IBinder;
+import android.os.IBinder.DeathRecipient;
+import android.os.Looper;
+import android.os.Message;
+import android.os.Messenger;
+import android.util.ArrayMap;
 import android.util.Log;
 import android.util.Xml;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile.GridOption;
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.util.Executors;
 
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
@@ -65,6 +76,11 @@
     private static final String ICON_THEMED = "/icon_themed";
     private static final String BOOLEAN_VALUE = "boolean_value";
 
+    private static final String KEY_SURFACE_PACKAGE = "surface_package";
+    private static final String KEY_CALLBACK = "callback";
+
+    private final ArrayMap<IBinder, PreviewLifecycleObserver> mActivePreviews = new ArrayMap<>();
+
     @Override
     public boolean onCreate() {
         return true;
@@ -177,10 +193,74 @@
             return null;
         }
 
-        if (!METHOD_GET_PREVIEW.equals(method)) {
+        if (!Utilities.ATLEAST_R || !METHOD_GET_PREVIEW.equals(method)) {
             return null;
         }
+        return getPreview(extras);
+    }
 
-        return new PreviewSurfaceRenderer(getContext(), extras).render();
+    @TargetApi(Build.VERSION_CODES.R)
+    private synchronized Bundle getPreview(Bundle request) {
+        PreviewLifecycleObserver observer = null;
+        try {
+            PreviewSurfaceRenderer renderer = new PreviewSurfaceRenderer(getContext(), request);
+
+            // Destroy previous
+            destroyObserver(mActivePreviews.get(renderer.getHostToken()));
+
+            observer = new PreviewLifecycleObserver(renderer);
+            mActivePreviews.put(renderer.getHostToken(), observer);
+
+            renderer.loadAsync();
+            renderer.getHostToken().linkToDeath(observer, 0);
+
+            Bundle result = new Bundle();
+            result.putParcelable(KEY_SURFACE_PACKAGE, renderer.getSurfacePackage());
+
+            Messenger messenger = new Messenger(new Handler(Looper.getMainLooper(), observer));
+            Message msg = Message.obtain();
+            msg.replyTo = messenger;
+            result.putParcelable(KEY_CALLBACK, msg);
+            return result;
+        } catch (Exception e) {
+            Log.e(TAG, "Unable to generate preview", e);
+            if (observer != null) {
+                destroyObserver(observer);
+            }
+            return null;
+        }
+    }
+
+    private synchronized void destroyObserver(PreviewLifecycleObserver observer) {
+        if (observer == null || observer.destroyed) {
+            return;
+        }
+        observer.destroyed = true;
+        Executors.MAIN_EXECUTOR.execute(observer.renderer::destroy);
+        PreviewLifecycleObserver cached = mActivePreviews.get(observer.renderer.getHostToken());
+        if (cached == observer) {
+            mActivePreviews.remove(observer.renderer.getHostToken());
+        }
+    }
+
+    private class PreviewLifecycleObserver implements Handler.Callback, DeathRecipient {
+
+        public final PreviewSurfaceRenderer renderer;
+        public boolean destroyed = false;
+
+        PreviewLifecycleObserver(PreviewSurfaceRenderer renderer) {
+            this.renderer = renderer;
+        }
+
+        @Override
+        public boolean handleMessage(Message message) {
+            destroyObserver(this);
+            return true;
+        }
+
+        @Override
+        public void binderDied() {
+            destroyObserver(this);
+        }
     }
 }
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index f5b6890..5f014db 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -23,7 +23,6 @@
 import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
 import static com.android.launcher3.model.ModelUtils.getMissingHotseatRanks;
 import static com.android.launcher3.model.ModelUtils.sortWorkspaceItemsSpatially;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
 import android.annotation.TargetApi;
 import android.app.Fragment;
@@ -32,7 +31,6 @@
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.Intent;
-import android.content.pm.ShortcutInfo;
 import android.content.res.TypedArray;
 import android.graphics.Color;
 import android.graphics.Rect;
@@ -43,7 +41,6 @@
 import android.os.Looper;
 import android.os.Process;
 import android.util.AttributeSet;
-import android.util.Log;
 import android.view.ContextThemeWrapper;
 import android.view.LayoutInflater;
 import android.view.View;
@@ -57,8 +54,6 @@
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherModel;
-import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
 import com.android.launcher3.WorkspaceLayoutManager;
@@ -67,13 +62,8 @@
 import com.android.launcher3.icons.BaseIconFactory;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.LauncherIcons;
-import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BgDataModel;
-import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
-import com.android.launcher3.model.LoaderResults;
-import com.android.launcher3.model.LoaderTask;
-import com.android.launcher3.model.ModelDelegate;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.model.data.FolderInfo;
@@ -100,13 +90,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentLinkedQueue;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Executor;
-import java.util.concurrent.FutureTask;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 
 /**
  * Utility class for generating the preview of Launcher for a given InvariantDeviceProfile.
@@ -120,8 +104,6 @@
 public class LauncherPreviewRenderer extends ContextWrapper
         implements ActivityContext, WorkspaceLayoutManager, LayoutInflater.Factory2 {
 
-    private static final String TAG = "LauncherPreviewRenderer";
-
     /**
      * Context used just for preview. It also provides a few objects (e.g. UserCache) just for
      * preview purposes.
@@ -138,9 +120,15 @@
         private final ConcurrentLinkedQueue<LauncherIconsForPreview> mIconPool =
                 new ConcurrentLinkedQueue<>();
 
+        private boolean mDestroyed = false;
+
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
             super(base);
             mIdp = idp;
+            mObjectMap.put(InvariantDeviceProfile.INSTANCE, idp);
+            mObjectMap.put(LauncherAppState.INSTANCE,
+                    new LauncherAppState(this, null /* iconCacheFileName */));
+
         }
 
         @Override
@@ -149,11 +137,9 @@
         }
 
         public void onDestroy() {
-            CustomWidgetManager customWidgetManager = (CustomWidgetManager) mObjectMap.get(
-                    CustomWidgetManager.INSTANCE);
-            if (customWidgetManager != null) {
-                customWidgetManager.onDestroy();
-            }
+            CustomWidgetManager.INSTANCE.get(this).onDestroy();
+            LauncherAppState.INSTANCE.get(this).onTerminate();
+            mDestroyed = true;
         }
 
         /**
@@ -162,17 +148,12 @@
          */
         public <T> T getObject(MainThreadInitializedObject<T> mainThreadInitializedObject,
                 MainThreadInitializedObject.ObjectProvider<T> provider) {
+            if (FeatureFlags.IS_STUDIO_BUILD && mDestroyed) {
+                throw new RuntimeException("Context already destroyed");
+            }
             if (!mAllowedObjects.contains(mainThreadInitializedObject)) {
                 throw new IllegalStateException("Leaking unknown objects");
             }
-            if (mainThreadInitializedObject == LauncherAppState.INSTANCE) {
-                throw new IllegalStateException(
-                        "Should not use MainThreadInitializedObject to initialize this with "
-                                + "PreviewContext");
-            }
-            if (mainThreadInitializedObject == InvariantDeviceProfile.INSTANCE) {
-                return (T) mIdp;
-            }
             if (mObjectMap.containsKey(mainThreadInitializedObject)) {
                 return (T) mObjectMap.get(mainThreadInitializedObject);
             }
@@ -210,7 +191,6 @@
     private final Context mContext;
     private final InvariantDeviceProfile mIdp;
     private final DeviceProfile mDp;
-    private final boolean mMigrated;
     private final Rect mInsets;
     private final WorkspaceItemInfo mWorkspaceItemInfo;
     private final LayoutInflater mHomeElementInflater;
@@ -218,13 +198,12 @@
     private final Hotseat mHotseat;
     private final CellLayout mWorkspace;
 
-    public LauncherPreviewRenderer(Context context, InvariantDeviceProfile idp, boolean migrated) {
+    public LauncherPreviewRenderer(Context context, InvariantDeviceProfile idp) {
         super(context);
         mUiHandler = new Handler(Looper.getMainLooper());
         mContext = context;
         mIdp = idp;
         mDp = idp.getDeviceProfile(context).copy(context);
-        mMigrated = migrated;
 
         // TODO: get correct insets once display cutout API is available.
         mInsets = new Rect();
@@ -265,8 +244,9 @@
     }
 
     /** Populate preview and render it. */
-    public View getRenderedView() {
-        populate();
+    public View getRenderedView(BgDataModel dataModel,
+            Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap) {
+        populate(dataModel, widgetProviderInfoMap);
         return mRootView;
     }
 
@@ -392,38 +372,17 @@
         }
     }
 
-    private void populate() {
-        WorkspaceFetcher fetcher;
-        PreviewContext previewContext = null;
-        if (mMigrated) {
-            previewContext = new PreviewContext(mContext, mIdp);
-            LauncherAppState appForPreview = new LauncherAppState(
-                    previewContext, null /* iconCacheFileName */);
-            fetcher = new WorkspaceItemsInfoFromPreviewFetcher(appForPreview);
-            MODEL_EXECUTOR.execute(fetcher);
-        } else {
-            fetcher = new WorkspaceItemsInfoFetcher();
-            LauncherAppState.getInstance(mContext).getModel().enqueueModelUpdateTask(
-                    (LauncherModel.ModelUpdateTask) fetcher);
-        }
-        WorkspaceResult workspaceResult = fetcher.get();
-        if (previewContext != null) {
-            previewContext.onDestroy();
-        }
-
-        if (workspaceResult == null) {
-            return;
-        }
-
+    private void populate(BgDataModel dataModel,
+            Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap) {
         // Separate the items that are on the current screen, and the other remaining items.
         ArrayList<ItemInfo> currentWorkspaceItems = new ArrayList<>();
         ArrayList<ItemInfo> otherWorkspaceItems = new ArrayList<>();
         ArrayList<LauncherAppWidgetInfo> currentAppWidgets = new ArrayList<>();
         ArrayList<LauncherAppWidgetInfo> otherAppWidgets = new ArrayList<>();
         filterCurrentWorkspaceItems(0 /* currentScreenId */,
-                workspaceResult.mWorkspaceItems, currentWorkspaceItems,
+                dataModel.workspaceItems, currentWorkspaceItems,
                 otherWorkspaceItems);
-        filterCurrentWorkspaceItems(0 /* currentScreenId */, workspaceResult.mAppWidgets,
+        filterCurrentWorkspaceItems(0 /* currentScreenId */, dataModel.appWidgets,
                 currentAppWidgets, otherAppWidgets);
         sortWorkspaceItemsSpatially(mIdp, currentWorkspaceItems);
         for (ItemInfo itemInfo : currentWorkspaceItems) {
@@ -444,12 +403,12 @@
             switch (itemInfo.itemType) {
                 case Favorites.ITEM_TYPE_APPWIDGET:
                 case Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                    if (mMigrated) {
-                        inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo,
-                                workspaceResult.mWidgetProvidersMap);
+                    if (widgetProviderInfoMap != null) {
+                        inflateAndAddWidgets(
+                                (LauncherAppWidgetInfo) itemInfo, widgetProviderInfoMap);
                     } else {
                         inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo,
-                                workspaceResult.mWidgetsModel);
+                                dataModel.widgetsModel);
                     }
                     break;
                 default:
@@ -458,8 +417,10 @@
         }
         IntArray ranks = getMissingHotseatRanks(currentWorkspaceItems,
                 mDp.numShownHotseatIcons);
-        List<ItemInfo> predictions = workspaceResult.mHotseatPredictions == null
-                ? Collections.emptyList() : workspaceResult.mHotseatPredictions.items;
+        FixedContainerItems hotseatpredictions =
+                dataModel.extraItems.get(CONTAINER_HOTSEAT_PREDICTION);
+        List<ItemInfo> predictions = hotseatpredictions == null
+                ? Collections.emptyList() : hotseatpredictions.items;
         int count = Math.min(ranks.size(), predictions.size());
         for (int i = 0; i < count; i++) {
             int rank = ranks.get(i);
@@ -494,109 +455,4 @@
         view.measure(makeMeasureSpec(width, EXACTLY), makeMeasureSpec(height, EXACTLY));
         view.layout(0, 0, width, height);
     }
-
-    private static class WorkspaceItemsInfoFetcher implements LauncherModel.ModelUpdateTask,
-            WorkspaceFetcher {
-
-        private final FutureTask<WorkspaceResult> mTask = new FutureTask<>(this);
-
-        private LauncherAppState mApp;
-        private LauncherModel mModel;
-        private BgDataModel mBgDataModel;
-        private AllAppsList mAllAppsList;
-
-        @Override
-        public void init(LauncherAppState app, LauncherModel model, BgDataModel dataModel,
-                AllAppsList allAppsList, Executor uiExecutor) {
-            mApp = app;
-            mModel = model;
-            mBgDataModel = dataModel;
-            mAllAppsList = allAppsList;
-        }
-
-        @Override
-        public FutureTask<WorkspaceResult> getTask() {
-            return mTask;
-        }
-
-        @Override
-        public void run() {
-            mTask.run();
-        }
-
-        @Override
-        public WorkspaceResult call() throws Exception {
-            if (!mModel.isModelLoaded()) {
-                Log.d(TAG, "Workspace not loaded, loading now");
-                mModel.startLoaderForResults(
-                        new LoaderResults(mApp, mBgDataModel, mAllAppsList, new Callbacks[0]));
-                return null;
-            }
-
-            return new WorkspaceResult(mBgDataModel, mBgDataModel.widgetsModel, null);
-        }
-    }
-
-    private static class WorkspaceItemsInfoFromPreviewFetcher extends LoaderTask implements
-            WorkspaceFetcher {
-
-        private final FutureTask<WorkspaceResult> mTask = new FutureTask<>(this);
-
-        WorkspaceItemsInfoFromPreviewFetcher(LauncherAppState app) {
-            super(app, null, new BgDataModel(), new ModelDelegate(), null);
-        }
-
-        @Override
-        public FutureTask<WorkspaceResult> getTask() {
-            return mTask;
-        }
-
-        @Override
-        public void run() {
-            mTask.run();
-        }
-
-        @Override
-        public WorkspaceResult call() {
-            List<ShortcutInfo> allShortcuts = new ArrayList<>();
-            loadWorkspace(allShortcuts, LauncherSettings.Favorites.PREVIEW_CONTENT_URI,
-                    LauncherSettings.Favorites.SCREEN + " = 0 or "
-                            + LauncherSettings.Favorites.CONTAINER + " = "
-                            + LauncherSettings.Favorites.CONTAINER_HOTSEAT);
-            return new WorkspaceResult(mBgDataModel, null, mWidgetProvidersMap);
-        }
-    }
-
-    private interface WorkspaceFetcher extends Runnable, Callable<WorkspaceResult> {
-        FutureTask<WorkspaceResult> getTask();
-
-        default WorkspaceResult get() {
-            try {
-                return getTask().get(5, TimeUnit.SECONDS);
-            } catch (InterruptedException | ExecutionException | TimeoutException e) {
-                Log.d(TAG, "Error fetching workspace items info", e);
-                return null;
-            }
-        }
-    }
-
-    private static class WorkspaceResult {
-        private final ArrayList<ItemInfo> mWorkspaceItems;
-        private final ArrayList<LauncherAppWidgetInfo> mAppWidgets;
-        private final FixedContainerItems mHotseatPredictions;
-        private final WidgetsModel mWidgetsModel;
-        private final Map<ComponentKey, AppWidgetProviderInfo> mWidgetProvidersMap;
-
-        private WorkspaceResult(BgDataModel dataModel,
-                WidgetsModel widgetsModel,
-                Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap) {
-            synchronized (dataModel) {
-                mWorkspaceItems = dataModel.workspaceItems;
-                mAppWidgets = dataModel.appWidgets;
-                mHotseatPredictions = dataModel.extraItems.get(CONTAINER_HOTSEAT_PREDICTION);
-                mWidgetsModel = widgetsModel;
-                mWidgetProvidersMap = widgetProviderInfoMap;
-            }
-        }
-    }
 }
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 6193570..8c39eae 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -21,32 +21,47 @@
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
 import android.app.WallpaperColors;
+import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.hardware.display.DisplayManager;
 import android.os.Bundle;
-import android.os.Handler;
 import android.os.IBinder;
-import android.os.Looper;
-import android.os.Message;
-import android.os.Messenger;
+import android.util.Log;
 import android.view.ContextThemeWrapper;
 import android.view.Display;
 import android.view.SurfaceControlViewHost;
+import android.view.SurfaceControlViewHost.SurfacePackage;
 import android.view.View;
 import android.view.animation.AccelerateDecelerateInterpolator;
 
+import androidx.annotation.UiThread;
+import androidx.annotation.WorkerThread;
+
 import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
+import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
+import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.GridSizeMigrationTask;
 import com.android.launcher3.model.GridSizeMigrationTaskV2;
+import com.android.launcher3.model.LoaderTask;
+import com.android.launcher3.model.ModelDelegate;
+import com.android.launcher3.model.ModelPreload;
+import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.widget.LocalColorExtractor;
 
+import java.util.ArrayList;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 /** Render preview using surface view. */
 @SuppressWarnings("NewApi")
-public class PreviewSurfaceRenderer implements IBinder.DeathRecipient {
+public class PreviewSurfaceRenderer {
+
+    private static final String TAG = "PreviewSurfaceRenderer";
 
     private static final int FADE_IN_ANIMATION_DURATION = 200;
 
@@ -54,8 +69,6 @@
     private static final String KEY_VIEW_WIDTH = "width";
     private static final String KEY_VIEW_HEIGHT = "height";
     private static final String KEY_DISPLAY_ID = "display_id";
-    private static final String KEY_SURFACE_PACKAGE = "surface_package";
-    private static final String KEY_CALLBACK = "callback";
     private static final String KEY_COLORS = "wallpaper_colors";
 
     private final Context mContext;
@@ -65,10 +78,13 @@
     private final int mHeight;
     private final Display mDisplay;
     private final WallpaperColors mWallpaperColors;
+    private final RunnableList mOnDestroyCallbacks = new RunnableList();
 
-    private SurfaceControlViewHost mSurfaceControlViewHost;
+    private final SurfaceControlViewHost mSurfaceControlViewHost;
 
-    PreviewSurfaceRenderer(Context context, Bundle bundle) {
+    private boolean mDestroyed = false;
+
+    public PreviewSurfaceRenderer(Context context, Bundle bundle) throws Exception {
         mContext = context;
 
         String gridName = bundle.getString("name");
@@ -77,106 +93,97 @@
             gridName = InvariantDeviceProfile.getCurrentGridName(context);
         }
         mWallpaperColors = bundle.getParcelable(KEY_COLORS);
-
         mIdp = new InvariantDeviceProfile(context, gridName);
 
         mHostToken = bundle.getBinder(KEY_HOST_TOKEN);
         mWidth = bundle.getInt(KEY_VIEW_WIDTH);
         mHeight = bundle.getInt(KEY_VIEW_HEIGHT);
+        mDisplay = context.getSystemService(DisplayManager.class)
+                .getDisplay(bundle.getInt(KEY_DISPLAY_ID));
 
-        final DisplayManager displayManager = (DisplayManager) context.getSystemService(
-                Context.DISPLAY_SERVICE);
-        mDisplay = displayManager.getDisplay(bundle.getInt(KEY_DISPLAY_ID));
+        mSurfaceControlViewHost = MAIN_EXECUTOR
+                .submit(() -> new SurfaceControlViewHost(mContext, mDisplay, mHostToken))
+                .get(5, TimeUnit.SECONDS);
+        mOnDestroyCallbacks.add(mSurfaceControlViewHost::release);
     }
 
-    /** Handle a received surface view request. */
-    Bundle render() {
-        if (mSurfaceControlViewHost != null) {
-            binderDied();
+    public IBinder getHostToken() {
+        return mHostToken;
+    }
+
+    public SurfacePackage getSurfacePackage() {
+        return mSurfaceControlViewHost.getSurfacePackage();
+    }
+
+    /**
+     * Destroys the preview and all associated data
+     */
+    @UiThread
+    public void destroy() {
+        mDestroyed = true;
+        mOnDestroyCallbacks.executeAllAndDestroy();
+    }
+
+    /**
+     * Generates the preview in background
+     */
+    public void loadAsync() {
+        MODEL_EXECUTOR.execute(this::loadModelData);
+    }
+
+    @WorkerThread
+    private void loadModelData() {
+        final boolean migrated = doGridMigrationIfNecessary();
+
+        final Context inflationContext;
+        if (mWallpaperColors != null) {
+            // Create a themed context, without affecting the main application context
+            Context context = mContext.createDisplayContext(mDisplay);
+            LocalColorExtractor.newInstance(mContext)
+                    .applyColorsOverride(context, mWallpaperColors);
+            inflationContext = new ContextThemeWrapper(context,
+                    Themes.getActivityThemeRes(context, mWallpaperColors.getColorHints()));
+        } else {
+            inflationContext = new ContextThemeWrapper(mContext,  R.style.AppTheme);
         }
 
-        SurfaceControlViewHost.SurfacePackage surfacePackage;
-        try {
-            mSurfaceControlViewHost = MAIN_EXECUTOR
-                    .submit(() -> new SurfaceControlViewHost(mContext, mDisplay, mHostToken))
-                    .get(5, TimeUnit.SECONDS);
-            surfacePackage = mSurfaceControlViewHost.getSurfacePackage();
-            mHostToken.linkToDeath(this, 0);
-        } catch (Exception e) {
-            e.printStackTrace();
-            return null;
-        }
+        if (migrated) {
+            PreviewContext previewContext = new PreviewContext(inflationContext, mIdp);
+            new LoaderTask(
+                    LauncherAppState.getInstance(previewContext),
+                    null,
+                    new BgDataModel(),
+                    new ModelDelegate(), null) {
 
-        MODEL_EXECUTOR.post(() -> {
-            final boolean success = doGridMigrationIfNecessary();
-
-            final Context inflationContext;
-            if (mWallpaperColors != null) {
-                // Workaround to create a themed context
-                Context context = mContext.createDisplayContext(mDisplay);
-                LocalColorExtractor.newInstance(mContext)
-                        .applyColorsOverride(context, mWallpaperColors);
-
-                inflationContext = new ContextThemeWrapper(context,
-                        Themes.getActivityThemeRes(context, mWallpaperColors.getColorHints()));
-            } else {
-                inflationContext = new ContextThemeWrapper(mContext,  R.style.AppTheme);
-            }
-
-            MAIN_EXECUTOR.post(() -> {
-                // If mSurfaceControlViewHost is null due to any reason (e.g. binder died,
-                // happening when user leaves the preview screen before preview rendering finishes),
-                // we should return here.
-                SurfaceControlViewHost host = mSurfaceControlViewHost;
-                if (host == null) {
-                    return;
+                @Override
+                public void run() {
+                    loadWorkspace(new ArrayList<>(), LauncherSettings.Favorites.PREVIEW_CONTENT_URI,
+                            LauncherSettings.Favorites.SCREEN + " = 0 or "
+                                    + LauncherSettings.Favorites.CONTAINER + " = "
+                                    + LauncherSettings.Favorites.CONTAINER_HOTSEAT);
+                    MAIN_EXECUTOR.execute(() -> {
+                        renderView(previewContext, mBgDataModel, mWidgetProvidersMap);
+                        mOnDestroyCallbacks.add(previewContext::onDestroy);
+                    });
                 }
+            }.run();
+        } else {
+            new ModelPreload() {
 
-                View view = new LauncherPreviewRenderer(inflationContext, mIdp, success)
-                        .getRenderedView();
-                // This aspect scales the view to fit in the surface and centers it
-                final float scale = Math.min(mWidth / (float) view.getMeasuredWidth(),
-                        mHeight / (float) view.getMeasuredHeight());
-                view.setScaleX(scale);
-                view.setScaleY(scale);
-                view.setPivotX(0);
-                view.setPivotY(0);
-                view.setTranslationX((mWidth - scale * view.getWidth()) / 2);
-                view.setTranslationY((mHeight - scale * view.getHeight()) / 2);
-                view.setAlpha(0);
-                view.animate().alpha(1)
-                        .setInterpolator(new AccelerateDecelerateInterpolator())
-                        .setDuration(FADE_IN_ANIMATION_DURATION)
-                        .start();
-                host.setView(view, view.getMeasuredWidth(), view.getMeasuredHeight());
-            });
-        });
-
-        Bundle result = new Bundle();
-        result.putParcelable(KEY_SURFACE_PACKAGE, surfacePackage);
-
-        Handler handler = new Handler(Looper.getMainLooper(), message -> {
-            binderDied();
-            return true;
-        });
-        Messenger messenger = new Messenger(handler);
-        Message msg = Message.obtain();
-        msg.replyTo = messenger;
-        result.putParcelable(KEY_CALLBACK, msg);
-        return result;
-    }
-
-    @Override
-    public void binderDied() {
-        if (mSurfaceControlViewHost != null) {
-            MAIN_EXECUTOR.execute(() -> {
-                mSurfaceControlViewHost.release();
-                mSurfaceControlViewHost = null;
-            });
+                @Override
+                public void onComplete(boolean isSuccess) {
+                    if (isSuccess) {
+                        MAIN_EXECUTOR.execute(() ->
+                                renderView(inflationContext, getBgDataModel(), null));
+                    } else {
+                        Log.e(TAG, "Model loading failed");
+                    }
+                }
+            }.start(inflationContext);
         }
-        mHostToken.unlinkToDeath(this, 0);
     }
 
+    @WorkerThread
     private boolean doGridMigrationIfNecessary() {
         boolean needsToMigrate =
                 MULTI_DB_GRID_MIRATION_ALGO.get()
@@ -189,4 +196,29 @@
                 ? GridSizeMigrationTaskV2.migrateGridIfNeeded(mContext, mIdp)
                 : GridSizeMigrationTask.migrateGridIfNeeded(mContext, mIdp);
     }
+
+    @UiThread
+    private void renderView(Context inflationContext, BgDataModel dataModel,
+            Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap) {
+        if (mDestroyed) {
+            return;
+        }
+        View view = new LauncherPreviewRenderer(inflationContext, mIdp)
+                .getRenderedView(dataModel, widgetProviderInfoMap);
+        // This aspect scales the view to fit in the surface and centers it
+        final float scale = Math.min(mWidth / (float) view.getMeasuredWidth(),
+                mHeight / (float) view.getMeasuredHeight());
+        view.setScaleX(scale);
+        view.setScaleY(scale);
+        view.setPivotX(0);
+        view.setPivotY(0);
+        view.setTranslationX((mWidth - scale * view.getWidth()) / 2);
+        view.setTranslationY((mHeight - scale * view.getHeight()) / 2);
+        view.setAlpha(0);
+        view.animate().alpha(1)
+                .setInterpolator(new AccelerateDecelerateInterpolator())
+                .setDuration(FADE_IN_ANIMATION_DURATION)
+                .start();
+        mSurfaceControlViewHost.setView(view, view.getMeasuredWidth(), view.getMeasuredHeight());
+    }
 }
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 297325a..8e0a388 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -131,6 +131,13 @@
     }
 
     /**
+     * Closes the cache DB. This will clear any in-memory cache.
+     */
+    public void close() {
+        mIconDb.close();
+    }
+
+    /**
      * Fetches high-res icon for the provided ItemInfo and updates the caller when done.
      *
      * @return a request ID that can be used to cancel the request.
diff --git a/src/com/android/launcher3/model/ModelPreload.java b/src/com/android/launcher3/model/ModelPreload.java
index 713492b..756b7da 100644
--- a/src/com/android/launcher3/model/ModelPreload.java
+++ b/src/com/android/launcher3/model/ModelPreload.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+
 import android.content.Context;
 import android.util.Log;
 
@@ -52,8 +54,14 @@
     public final void run() {
         mModel.startLoaderForResultsIfNotLoaded(
                 new LoaderResults(mApp, mBgDataModel, mAllAppsList, new Callbacks[0]));
-        Log.d(TAG, "Preload completed : " + mModel.isModelLoaded());
-        onComplete(mModel.isModelLoaded());
+        MODEL_EXECUTOR.post(() -> {
+            Log.d(TAG, "Preload completed : " + mModel.isModelLoaded());
+            onComplete(mModel.isModelLoaded());
+        });
+    }
+
+    public BgDataModel getBgDataModel() {
+        return mBgDataModel;
     }
 
     /**