Wait for device profile reload

Test: Manually tested the option updated after apply
Bug: 348664593
Flag: com.android.systemui.shared.new_customization_picker_ui
Change-Id: I795b95407483a730ee22a184a6db96f64f6ee829
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
index 374c07b..531cdfd 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
@@ -16,11 +16,13 @@
 package com.android.launcher3.graphics;
 
 import static com.android.launcher3.LauncherPrefs.THEMED_ICONS;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.Themes.isThemedIconEnabled;
 
 import android.content.ContentProvider;
 import android.content.ContentValues;
+import android.content.Context;
 import android.content.pm.PackageManager;
 import android.database.Cursor;
 import android.database.MatrixCursor;
@@ -39,8 +41,15 @@
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile.GridOption;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherPrefs;
+import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.Preconditions;
+import com.android.systemui.shared.Flags;
+
+import java.util.concurrent.ExecutionException;
 
 /**
  * Exposes various launcher grid options and allows the caller to change them.
@@ -144,14 +153,20 @@
 
     @Override
     public int update(Uri uri, ContentValues values, String selection, String[] selectionArgs) {
-        switch (uri.getPath()) {
+        String path = uri.getPath();
+        Context context = getContext();
+        if (path == null || context == null) {
+            return 0;
+        }
+        switch (path) {
             case KEY_DEFAULT_GRID: {
                 String gridName = values.getAsString(KEY_NAME);
-                InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(getContext());
+                InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
                 // Verify that this is a valid grid option
                 GridOption match = null;
-                for (GridOption option : idp.parseAllGridOptions(getContext())) {
-                    if (option.name.equals(gridName)) {
+                for (GridOption option : idp.parseAllGridOptions(context)) {
+                    String name = option.name;
+                    if (name != null && name.equals(gridName)) {
                         match = option;
                         break;
                     }
@@ -160,15 +175,23 @@
                     return 0;
                 }
 
-                idp.setCurrentGrid(getContext(), gridName);
-                getContext().getContentResolver().notifyChange(uri, null);
+                idp.setCurrentGrid(context, gridName);
+                if (Flags.newCustomizationPickerUi()) {
+                    try {
+                        // Wait for device profile to be fully reloaded and applied to the launcher
+                        loadModelSync(context);
+                    } catch (ExecutionException | InterruptedException e) {
+                        Log.e(TAG, "Fail to load model", e);
+                    }
+                }
+                context.getContentResolver().notifyChange(uri, null);
                 return 1;
             }
             case ICON_THEMED:
             case SET_ICON_THEMED: {
-                LauncherPrefs.get(getContext())
+                LauncherPrefs.get(context)
                         .put(THEMED_ICONS, values.getAsBoolean(BOOLEAN_VALUE));
-                getContext().getContentResolver().notifyChange(uri, null);
+                context.getContentResolver().notifyChange(uri, null);
                 return 1;
             }
             default:
@@ -176,6 +199,23 @@
         }
     }
 
+    /**
+     * Loads the model in memory synchronously
+     */
+    private void loadModelSync(Context context) throws ExecutionException, InterruptedException {
+        Preconditions.assertNonUiThread();
+        BgDataModel.Callbacks emptyCallbacks = new BgDataModel.Callbacks() { };
+        LauncherModel launcherModel = LauncherAppState.getInstance(context).getModel();
+        MAIN_EXECUTOR.submit(
+                () -> launcherModel.addCallbacksAndLoad(emptyCallbacks)
+        ).get();
+
+        Executors.MODEL_EXECUTOR.submit(() -> { }).get();
+        MAIN_EXECUTOR.submit(
+                () -> launcherModel.removeCallbacks(emptyCallbacks)
+        ).get();
+    }
+
     @Override
     public Bundle call(String method, String arg, Bundle extras) {
         if (getContext().checkPermission("android.permission.BIND_WALLPAPER",
@@ -227,7 +267,7 @@
         }
         observer.destroyed = true;
         observer.renderer.getHostToken().unlinkToDeath(observer, 0);
-        Executors.MAIN_EXECUTOR.execute(observer.renderer::destroy);
+        MAIN_EXECUTOR.execute(observer.renderer::destroy);
         PreviewLifecycleObserver cached = mActivePreviews.get(observer.getIdentifier());
         if (cached == observer) {
             mActivePreviews.remove(observer.getIdentifier());