Ensure app widget ids are restored after database is sanitized.

Previously, it was possible for AppWidgetsRestoredReceiver to
start the restore process before work profile has finished restoring which
resulted in the work profile items being removed from the workspace.

Bug: 131315856
Change-Id: I2f295a1ca91f1996522bcc8052aa139979526e3b

4 files changed, 62 insertions, 16 deletions

diff --git a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
index 0250c363d..d9491419f 100644
--- a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
+++ b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
@@ -8,7 +8,6 @@ import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
 import android.database.Cursor;
-import android.os.Handler;
 import android.util.Log;
 
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -34,16 +33,8 @@ public class AppWidgetsRestoredReceiver extends BroadcastReceiver {
 
             final int[] oldIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS);
             final int[] newIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS);
-            if (oldIds.length == newIds.length) {
-                final PendingResult asyncResult = goAsync();
-                new Handler(LauncherModel.getWorkerLooper())
-                        .postAtFrontOfQueue(new Runnable() {
-                            @Override
-                            public void run() {
-                                restoreAppWidgetIds(context, oldIds, newIds);
-                                asyncResult.finish();
-                            }
-                        });
+            if (oldIds != null && newIds != null && oldIds.length == newIds.length) {
+                RestoreDbTask.setRestoredAppWidgetIds(context, oldIds, newIds);
             } else {
                 Log.e(TAG, "Invalid host restored received");
             }
@@ -54,7 +45,7 @@ public class AppWidgetsRestoredReceiver extends BroadcastReceiver {
      * Updates the app widgets whose id has changed during the restore process.
      */
     @WorkerThread
-    static void restoreAppWidgetIds(Context context, int[] oldWidgetIds, int[] newWidgetIds) {
+    public static void restoreAppWidgetIds(Context context, int[] oldWidgetIds, int[] newWidgetIds) {
         AppWidgetHost appWidgetHost = new LauncherAppWidgetHost(context);
         if (FeatureFlags.GO_DISABLE_WIDGETS) {
             Log.e(TAG, "Skipping widget ID remap as widgets not supported");
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index f830301a3..6ad5c3684 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -158,7 +158,8 @@ public class LauncherProvider extends ContentProvider {
             mOpenHelper = new DatabaseHelper(getContext(), mListenerHandler);
 
             if (RestoreDbTask.isPending(getContext())) {
-                if (!RestoreDbTask.performRestore(mOpenHelper, new BackupManager(getContext()))) {
+                if (!RestoreDbTask.performRestore(getContext(), mOpenHelper,
+                        new BackupManager(getContext()))) {
                     mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
                 }
                 // Set is pending to false irrespective of the result, so that it doesn't get
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 5cfd95cd0..02fc84b9a 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -74,6 +74,7 @@ import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Locale;
+import java.util.StringTokenizer;
 import java.util.concurrent.Executor;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -652,4 +653,23 @@ public final class Utilities {
             return null;
         }
     }
+
+    public static int[] getIntArrayFromString(String tokenized) {
+        StringTokenizer tokenizer = new StringTokenizer(tokenized, ",");
+        int[] array = new int[tokenizer.countTokens()];
+        int count = 0;
+        while (tokenizer.hasMoreTokens()) {
+            array[count] = Integer.parseInt(tokenizer.nextToken());
+            count++;
+        }
+        return array;
+    }
+
+    public static String getStringFromIntArray(int[] array) {
+        StringBuilder str = new StringBuilder();
+        for (int value : array) {
+            str.append(value).append(",");
+        }
+        return str.toString();
+    }
 }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 053c493d7..3c0c5fdde 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -16,18 +16,23 @@
 
 package com.android.launcher3.provider;
 
+import static com.android.launcher3.Utilities.getIntArrayFromString;
+import static com.android.launcher3.Utilities.getStringFromIntArray;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 
 import android.app.backup.BackupManager;
 import android.content.ContentValues;
 import android.content.Context;
+import android.content.SharedPreferences;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
-import android.os.Build;
 import android.os.UserHandle;
 import android.util.LongSparseArray;
 import android.util.SparseLongArray;
 
+import androidx.annotation.NonNull;
+
+import com.android.launcher3.AppWidgetsRestoredReceiver;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -53,10 +58,16 @@ public class RestoreDbTask {
     private static final String INFO_COLUMN_NAME = "name";
     private static final String INFO_COLUMN_DEFAULT_VALUE = "dflt_value";
 
-    public static boolean performRestore(DatabaseHelper helper, BackupManager backupManager) {
+    private static final String APPWIDGET_OLD_IDS = "appwidget_old_ids";
+    private static final String APPWIDGET_IDS = "appwidget_ids";
+
+    public static boolean performRestore(Context context, DatabaseHelper helper,
+            BackupManager backupManager) {
         SQLiteDatabase db = helper.getWritableDatabase();
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            new RestoreDbTask().sanitizeDB(helper, db, backupManager);
+            RestoreDbTask task = new RestoreDbTask();
+            task.sanitizeDB(helper, db, backupManager);
+            task.restoreAppWidgetIdsIfExists(context);
             t.commit();
             return true;
         } catch (Exception e) {
@@ -230,4 +241,27 @@ public class RestoreDbTask {
         FileLog.d(TAG, "Restore data received through full backup " + isPending);
         Utilities.getPrefs(context).edit().putBoolean(RESTORE_TASK_PENDING, isPending).commit();
     }
+
+    private void restoreAppWidgetIdsIfExists(Context context) {
+        SharedPreferences prefs = Utilities.getPrefs(context);
+        if (prefs.contains(APPWIDGET_OLD_IDS) && prefs.contains(APPWIDGET_IDS)) {
+            AppWidgetsRestoredReceiver.restoreAppWidgetIds(context,
+                    getIntArrayFromString(prefs.getString(APPWIDGET_OLD_IDS, "")),
+                    getIntArrayFromString(prefs.getString(APPWIDGET_IDS, "")));
+        } else {
+            FileLog.d(TAG, "No app widget ids to restore.");
+        }
+
+        prefs.edit().remove(APPWIDGET_OLD_IDS)
+                .remove(APPWIDGET_IDS).apply();
+    }
+
+    public static void setRestoredAppWidgetIds(Context context, @NonNull int[] oldIds,
+            @NonNull int[] newIds) {
+        Utilities.getPrefs(context).edit()
+                .putString(APPWIDGET_OLD_IDS, getStringFromIntArray(oldIds))
+                .putString(APPWIDGET_IDS, getStringFromIntArray(newIds))
+                .commit();
+    }
+
 }