Implement LauncherLayoutSnapshot pulled atom callback.

This change will pull only the workspace items but not the launcher settings such as grid layout etc. Future CLs will address it.

Bug: 181703659
Test: Manual-`statsd_testdrive 10108` http://gpaste/4894958784872448
Change-Id: I0e7d7da62034c411edb26ff9205a81b0e2e891a0
diff --git a/Android.bp b/Android.bp
index 43d28c9..c8d9186 100644
--- a/Android.bp
+++ b/Android.bp
@@ -33,6 +33,9 @@
 
 android_library {
     name: "launcher-aosp-tapl",
+    libs: [
+        "framework-statsd",
+    ],
     static_libs: [
         "androidx.annotation_annotation",
         "androidx.test.runner",
@@ -192,6 +195,9 @@
     resource_dirs: [
         "quickstep/res",
     ],
+    libs: [
+        "framework-statsd",
+    ],
     static_libs: [
         "Launcher3ResLib",
         "SystemUISharedLib",
@@ -261,6 +267,9 @@
     resource_dirs: [
         "quickstep/res",
     ],
+    libs: [
+        "framework-statsd",
+    ],
     static_libs: [
         "SystemUI-statsd",
         "SystemUISharedLib",
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 5769f0b..55a140d 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -25,7 +25,9 @@
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.Utilities.getDevicePrefs;
 import static com.android.launcher3.hybridhotseat.HotseatPredictionModel.convertDataModelToAppTargetBundle;
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
+import android.app.StatsManager;
 import android.app.prediction.AppPredictionContext;
 import android.app.prediction.AppPredictionManager;
 import android.app.prediction.AppPredictor;
@@ -39,12 +41,14 @@
 import android.content.pm.ShortcutInfo;
 import android.os.UserHandle;
 import android.util.Log;
+import android.util.StatsEvent;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.logging.InstanceIdSequence;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
@@ -53,10 +57,10 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.shortcuts.ShortcutKey;
-import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.PersistedItemArray;
 import com.android.quickstep.logging.StatsLogCompatManager;
+import com.android.systemui.shared.system.SysUiStatsLog;
 
 import java.util.Collections;
 import java.util.List;
@@ -85,6 +89,7 @@
 
     private final InvariantDeviceProfile mIDP;
     private final AppEventProducer mAppEventProducer;
+    private final StatsManager mStatsManager;
 
     protected boolean mActive = false;
 
@@ -93,6 +98,7 @@
 
         mIDP = InvariantDeviceProfile.INSTANCE.get(context);
         StatsLogCompatManager.LOGS_CONSUMER.add(mAppEventProducer);
+        mStatsManager = context.getSystemService(StatsManager.class);
     }
 
     @Override
@@ -155,10 +161,60 @@
             additionalSnapshotEvents(instanceId);
             prefs.edit().putLong(LAST_SNAPSHOT_TIME_MILLIS, now).apply();
         }
+
+        // Only register for launcher snapshot logging if this is the primary ModelDelegate
+        // instance, as there will be additional instances that may be destroyed at any time.
+        if (mIsPrimaryInstance) {
+            registerSnapshotLoggingCallback();
+        }
     }
 
     protected void additionalSnapshotEvents(InstanceId snapshotInstanceId){}
 
+    /**
+     * Registers a callback to log launcher workspace layout using Statsd pulled atom.
+     */
+    protected void registerSnapshotLoggingCallback() {
+        if (mStatsManager == null) {
+            Log.d(TAG, "Failed to get StatsManager");
+        }
+
+        try {
+            mStatsManager.setPullAtomCallback(
+                    SysUiStatsLog.LAUNCHER_LAYOUT_SNAPSHOT,
+                    null /* PullAtomMetadata */,
+                    MODEL_EXECUTOR,
+                    (i, eventList) -> {
+                        InstanceId instanceId = new InstanceIdSequence().newInstanceId();
+                        IntSparseArrayMap<ItemInfo> itemsIdMap;
+                        synchronized (mDataModel) {
+                            itemsIdMap = mDataModel.itemsIdMap.clone();
+                        }
+
+                        for (ItemInfo info : itemsIdMap) {
+                            FolderInfo parent = info.container > 0
+                                    ? (FolderInfo) itemsIdMap.get(info.container) : null;
+                            LauncherAtom.ItemInfo itemInfo = info.buildProto(parent);
+                            Log.d(TAG, itemInfo.toString());
+                            StatsEvent statsEvent = StatsLogCompatManager.buildStatsEvent(itemInfo,
+                                    instanceId);
+                            eventList.add(statsEvent);
+                        }
+                        Log.d(TAG,
+                                String.format(
+                                        "Successfully logged %d workspace items with instanceId=%d",
+                                        itemsIdMap.size(), instanceId.getId()));
+                        additionalSnapshotEvents(instanceId);
+                        return StatsManager.PULL_SUCCESS;
+                    }
+            );
+            Log.d(TAG, "Successfully registered for launcher snapshot logging!");
+        } catch (RuntimeException e) {
+            Log.e(TAG, "Failed to register launcher snapshot logging callback with StatsManager",
+                    e);
+        }
+    }
+
     @Override
     public void validateData() {
         super.validateData();
@@ -175,7 +231,9 @@
         super.destroy();
         mActive = false;
         StatsLogCompatManager.LOGS_CONSUMER.remove(mAppEventProducer);
-
+        if (mIsPrimaryInstance) {
+            mStatsManager.clearPullAtomCallback(SysUiStatsLog.LAUNCHER_LAYOUT_SNAPSHOT);
+        }
         destroyPredictors();
     }
 
@@ -221,7 +279,7 @@
     private void registerPredictor(PredictorState state, AppPredictor predictor) {
         state.predictor = predictor;
         state.predictor.registerPredictionUpdates(
-                Executors.MODEL_EXECUTOR, t -> handleUpdate(state, t));
+                MODEL_EXECUTOR, t -> handleUpdate(state, t));
         state.predictor.requestPredictionUpdate();
     }
 
@@ -236,7 +294,7 @@
     private void registerWidgetsPredictor(AppPredictor predictor) {
         mWidgetsRecommendationState.predictor = predictor;
         mWidgetsRecommendationState.predictor.registerPredictionUpdates(
-                Executors.MODEL_EXECUTOR, targets -> {
+                MODEL_EXECUTOR, targets -> {
                     if (mWidgetsRecommendationState.setTargets(targets)) {
                         // No diff, skip
                         return;
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index 6575996..7893f8d 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -31,9 +31,11 @@
 
 import android.content.Context;
 import android.util.Log;
+import android.util.StatsEvent;
 import android.view.View;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
 import androidx.slice.SliceItem;
 
@@ -135,6 +137,36 @@
     }
 
     /**
+     * Builds {@link StatsEvent} from {@link LauncherAtom.ItemInfo}. Used for pulled atom callback
+     * implementation.
+     */
+    public static StatsEvent buildStatsEvent(LauncherAtom.ItemInfo info,
+            @Nullable InstanceId instanceId) {
+        return SysUiStatsLog.buildStatsEvent(
+                SysUiStatsLog.LAUNCHER_LAYOUT_SNAPSHOT, // atom ID,
+                LAUNCHER_WORKSPACE_SNAPSHOT.getId(), // event_id = 1;
+                info.getAttribute().getNumber() * ATTRIBUTE_MULTIPLIER
+                        + info.getItemCase().getNumber(), // item_id = 2;
+                instanceId == null ? 0 : instanceId.getId(), //instance_id = 3;
+                0, //uid = 4 [(is_uid) = true];
+                getPackageName(info), // package_name = 5;
+                getComponentName(info), // component_name = 6;
+                getGridX(info, false), //grid_x = 7 [default = -1];
+                getGridY(info, false), //grid_y = 8 [default = -1];
+                getPageId(info), // page_id = 9 [default = -2];
+                getGridX(info, true), //grid_x_parent = 10 [default = -1];
+                getGridY(info, true), //grid_y_parent = 11 [default = -1];
+                getParentPageId(info), //page_id_parent = 12 [default = -2];
+                getHierarchy(info), // container_id = 13;
+                info.getIsWork(), // is_work_profile = 14;
+                info.getAttribute().getNumber(), // attribute_id = 15;
+                getCardinality(info), // cardinality = 16;
+                info.getWidget().getSpanX(), // span_x = 17 [default = 1];
+                info.getWidget().getSpanY() // span_y = 18 [default = 1];
+        );
+    }
+
+    /**
      * Helps to construct and write statsd compatible log message.
      */
     private static class StatsCompatLogger implements StatsLogger {
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 86217d7..10023b4 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -138,10 +138,11 @@
         mContext = context;
 
         mInvariantDeviceProfile = InvariantDeviceProfile.INSTANCE.get(context);
-        mIconProvider =  new IconProvider(context, Themes.isThemedIconEnabled(context));
+        mIconProvider = new IconProvider(context, Themes.isThemedIconEnabled(context));
         mIconCache = new IconCache(mContext, mInvariantDeviceProfile,
                 iconCacheFileName, mIconProvider);
-        mModel = new LauncherModel(context, this, mIconCache, new AppFilter(mContext));
+        mModel = new LauncherModel(context, this, mIconCache, new AppFilter(mContext),
+                iconCacheFileName != null);
         mOnTerminateCallback.add(mIconCache::close);
     }
 
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 9ebec0a..f38f662 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -126,10 +126,12 @@
         }
     };
 
-    LauncherModel(Context context, LauncherAppState app, IconCache iconCache, AppFilter appFilter) {
+    LauncherModel(Context context, LauncherAppState app, IconCache iconCache, AppFilter appFilter,
+            boolean isPrimaryInstance) {
         mApp = app;
         mBgAllAppsList = new AllAppsList(iconCache, appFilter);
-        mModelDelegate = ModelDelegate.newInstance(context, app, mBgAllAppsList, mBgDataModel);
+        mModelDelegate = ModelDelegate.newInstance(context, app, mBgAllAppsList, mBgDataModel,
+                isPrimaryInstance);
     }
 
     public ModelDelegate getModelDelegate() {
diff --git a/src/com/android/launcher3/model/ModelDelegate.java b/src/com/android/launcher3/model/ModelDelegate.java
index 13ec1ec..765141a 100644
--- a/src/com/android/launcher3/model/ModelDelegate.java
+++ b/src/com/android/launcher3/model/ModelDelegate.java
@@ -40,19 +40,21 @@
      * Creates and initializes a new instance of the delegate
      */
     public static ModelDelegate newInstance(
-            Context context, LauncherAppState app, AllAppsList appsList, BgDataModel dataModel) {
+            Context context, LauncherAppState app, AllAppsList appsList, BgDataModel dataModel,
+            boolean isPrimaryInstance) {
         ModelDelegate delegate = Overrides.getObject(
                 ModelDelegate.class, context, R.string.model_delegate_class);
-
         delegate.mApp = app;
         delegate.mAppsList = appsList;
         delegate.mDataModel = dataModel;
+        delegate.mIsPrimaryInstance = isPrimaryInstance;
         return delegate;
     }
 
     protected LauncherAppState mApp;
     protected AllAppsList mAppsList;
     protected BgDataModel mDataModel;
+    protected boolean mIsPrimaryInstance;
 
     public ModelDelegate() { }