[StatsPullAtomService] Introduced flag to fix NetworkStatsManager init

- NetworkStatsManager should be ready at
  PHASE_THIRD_PARTY_APPS_CAN_START phase

Bug: 331989853
Test: build, flash & run
Test: adb shell cmd stats print-stats | grep 10001
Change-Id: Ia1a00fbfc78767b4d3b04bc213a2bfbf3e3d8dc5
diff --git a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
index b5df30f..e0247d0 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -61,6 +61,7 @@
 import static com.android.internal.util.FrameworkStatsLog.TIME_ZONE_DETECTOR_STATE__DETECTION_MODE__UNKNOWN;
 import static com.android.server.am.MemoryStatUtil.readMemoryStatFromFilesystem;
 import static com.android.server.stats.Flags.addMobileBytesTransferByProcStatePuller;
+import static com.android.server.stats.Flags.statsPullNetworkStatsManagerInitOrderFix;
 import static com.android.server.stats.pull.IonMemoryUtil.readProcessSystemIonHeapSizesFromDebugfs;
 import static com.android.server.stats.pull.IonMemoryUtil.readSystemIonHeapSizeFromDebugfs;
 import static com.android.server.stats.pull.ProcfsMemoryUtil.getProcessCmdlines;
@@ -355,7 +356,17 @@
     private TelephonyManager mTelephony;
     private UwbManager mUwbManager;
     private SubscriptionManager mSubscriptionManager;
-    private NetworkStatsManager mNetworkStatsManager;
+
+    /**
+     * NetworkStatsManager initialization happens from one thread before any worker thread
+     * is going to access the networkStatsManager instance:
+     * - @initNetworkStatsManager() - initialization happens no worker thread to access are
+     *   active yet
+     * - @initAndRegisterNetworkStatsPullers Network stats dependant pullers can only be
+     *   initialized after service is ready. Worker thread is spawn here only after the
+     *   initialization is completed in a thread safe way (no async access expected)
+     */
+    private NetworkStatsManager mNetworkStatsManager = null;
 
     @GuardedBy("mKernelWakelockLock")
     private KernelWakelockReader mKernelWakelockReader;
@@ -420,6 +431,12 @@
     public static final boolean ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER =
                 addMobileBytesTransferByProcStatePuller();
 
+    /**
+     * Whether or not to enable the mNetworkStatsManager initialization order fix
+     */
+    private static final boolean ENABLE_NETWORK_STATS_MANAGER_INIT_ORDER_FIX =
+                statsPullNetworkStatsManagerInitOrderFix();
+
     // Puller locks
     private final Object mDataBytesTransferLock = new Object();
     private final Object mBluetoothBytesTransferLock = new Object();
@@ -824,6 +841,9 @@
             });
         } else if (phase == PHASE_THIRD_PARTY_APPS_CAN_START) {
             BackgroundThread.getHandler().post(() -> {
+                if (ENABLE_NETWORK_STATS_MANAGER_INIT_ORDER_FIX) {
+                    initNetworkStatsManager();
+                }
                 // Network stats related pullers can only be initialized after service is ready.
                 initAndRegisterNetworkStatsPullers();
                 // For services that are not ready at boot phase PHASE_SYSTEM_SERVICES_READY
@@ -843,7 +863,9 @@
                 mContext.getSystemService(Context.TELEPHONY_SUBSCRIPTION_SERVICE);
         mStatsSubscriptionsListener = new StatsSubscriptionsListener(mSubscriptionManager);
         mStorageManager = (StorageManager) mContext.getSystemService(StorageManager.class);
-        mNetworkStatsManager = mContext.getSystemService(NetworkStatsManager.class);
+        if (!ENABLE_NETWORK_STATS_MANAGER_INIT_ORDER_FIX) {
+            initNetworkStatsManager();
+        }
 
         // Initialize DiskIO
         mStoragedUidIoStatsReader = new StoragedUidIoStatsReader();
@@ -1019,6 +1041,24 @@
         }
     }
 
+    /**
+     * Calling getNetworkStatsManager() before PHASE_THIRD_PARTY_APPS_CAN_START is unexpected
+     * Callers use before PHASE_THIRD_PARTY_APPS_CAN_START stage is not legit
+     */
+    @NonNull
+    private NetworkStatsManager getNetworkStatsManager() {
+        if (ENABLE_NETWORK_STATS_MANAGER_INIT_ORDER_FIX) {
+            if (mNetworkStatsManager == null) {
+                throw new IllegalStateException("NetworkStatsManager is not ready");
+            }
+        }
+        return mNetworkStatsManager;
+    }
+
+    private void initNetworkStatsManager() {
+        mNetworkStatsManager = mContext.getSystemService(NetworkStatsManager.class);
+    }
+
     private void initAndRegisterNetworkStatsPullers() {
         if (DEBUG) {
             Slog.d(TAG, "Registering NetworkStats pullers with statsd");
@@ -1514,11 +1554,11 @@
         //  I/O and also block main thread when polling.
         //  Consider making perfd queries NetworkStatsService directly.
         if (template.getMatchRule() == MATCH_WIFI && template.getSubscriberIds().isEmpty()) {
-            mNetworkStatsManager.forceUpdate();
+            getNetworkStatsManager().forceUpdate();
         }
 
         final android.app.usage.NetworkStats queryNonTaggedStats =
-                mNetworkStatsManager.querySummary(
+                getNetworkStatsManager().querySummary(
                         template, currentTimeInMillis - elapsedMillisSinceBoot - bucketDuration,
                         currentTimeInMillis);
 
@@ -1528,7 +1568,7 @@
         if (!includeTags) return nonTaggedStats;
 
         final android.app.usage.NetworkStats queryTaggedStats =
-                mNetworkStatsManager.queryTaggedSummary(template,
+                getNetworkStatsManager().queryTaggedSummary(template,
                         currentTimeInMillis - elapsedMillisSinceBoot - bucketDuration,
                         currentTimeInMillis);
         final NetworkStats taggedStats =
diff --git a/services/core/java/com/android/server/stats/stats_flags.aconfig b/services/core/java/com/android/server/stats/stats_flags.aconfig
index 101b98e..c479c6d 100644
--- a/services/core/java/com/android/server/stats/stats_flags.aconfig
+++ b/services/core/java/com/android/server/stats/stats_flags.aconfig
@@ -7,4 +7,12 @@
     description: "Adds mobile_bytes_transfer_by_proc_state atom with system server side aggregation"
     bug: "309512867"
     is_fixed_read_only: true
-}
\ No newline at end of file
+}
+
+flag {
+    name: "stats_pull_network_stats_manager_init_order_fix"
+    namespace: "statsd"
+    description: "Fix the mNetworkStatsManager initialization order"
+    bug: "331989853"
+    is_fixed_read_only: true
+}