Added puller impl for mobile data transfered per uid by proc_state

Bug: 309512867
Test: adb shell cmd stats print-stats | grep 10201

Change-Id: Ie1c25766d1d32a7e3814710354eb242785df0c09

7 files changed, 415 insertions, 1 deletions

diff --git a/services/core/Android.bp b/services/core/Android.bp
index 9375fb14a6d7..a54a48a7e84e 100644
--- a/services/core/Android.bp
+++ b/services/core/Android.bp
@@ -214,6 +214,7 @@ java_library_static {
         "backup_flags_lib",
         "policy_flags_lib",
         "net_flags_lib",
+        "stats_flags_lib",
     ],
     javac_shard_size: 50,
     javacflags: [
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index 3de224addf95..a6c35b6e5608 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -472,6 +472,8 @@ import com.android.server.pm.pkg.SELinuxUtil;
 import com.android.server.pm.snapshot.PackageDataSnapshot;
 import com.android.server.power.stats.BatteryStatsImpl;
 import com.android.server.sdksandbox.SdkSandboxManagerLocal;
+import com.android.server.stats.pull.StatsPullAtomService;
+import com.android.server.stats.pull.StatsPullAtomServiceInternal;
 import com.android.server.uri.GrantUri;
 import com.android.server.uri.NeededUriGrants;
 import com.android.server.uri.UriGrantsManagerInternal;
@@ -1308,6 +1310,8 @@ public class ActivityManagerService extends IActivityManager.Stub
      */
     final BatteryStatsService mBatteryStatsService;
 
+    StatsPullAtomServiceInternal mStatsPullAtomServiceInternal;
+
     /**
      * Information about component usage
      */
@@ -16556,6 +16560,21 @@ public class ActivityManagerService extends IActivityManager.Stub
                 final @ProcessCapability int capability) {
         mBatteryStatsService.noteUidProcessState(uid, state);
         mAppOpsService.updateUidProcState(uid, state, capability);
+        if (StatsPullAtomService.ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER) {
+            try {
+                if (mStatsPullAtomServiceInternal == null) {
+                    mStatsPullAtomServiceInternal = LocalServices.getService(
+                            StatsPullAtomServiceInternal.class);
+                }
+                if (mStatsPullAtomServiceInternal != null) {
+                    mStatsPullAtomServiceInternal.noteUidProcessState(uid, state);
+                } else {
+                    Slog.d(TAG, "StatsPullAtomService not ready yet");
+                }
+            } catch (Exception e) {
+                Slog.e(TAG, "Exception during logging uid proc state change event", e);
+            }
+        }
         if (mTrackingAssociations) {
             for (int i1=0, N1=mAssociations.size(); i1<N1; i1++) {
                 ArrayMap<ComponentName, SparseArray<ArrayMap<String, Association>>> targetComponents
diff --git a/services/core/java/com/android/server/stats/Android.bp b/services/core/java/com/android/server/stats/Android.bp
new file mode 100644
index 000000000000..e597c3a6b6e6
--- /dev/null
+++ b/services/core/java/com/android/server/stats/Android.bp
@@ -0,0 +1,12 @@
+aconfig_declarations {
+    name: "stats_flags",
+    package: "com.android.server.stats",
+    srcs: [
+        "stats_flags.aconfig",
+    ],
+}
+
+java_aconfig_library {
+    name: "stats_flags_lib",
+    aconfig_declarations: "stats_flags",
+}
diff --git a/services/core/java/com/android/server/stats/pull/AggregatedMobileDataStatsPuller.java b/services/core/java/com/android/server/stats/pull/AggregatedMobileDataStatsPuller.java
new file mode 100644
index 000000000000..0de73a5a30f6
--- /dev/null
+++ b/services/core/java/com/android/server/stats/pull/AggregatedMobileDataStatsPuller.java
@@ -0,0 +1,285 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.stats.pull;
+
+import android.app.ActivityManager;
+import android.app.StatsManager;
+import android.app.usage.NetworkStatsManager;
+import android.net.NetworkStats;
+import android.os.Handler;
+import android.os.HandlerThread;
+import android.os.Trace;
+import android.util.ArrayMap;
+import android.util.Slog;
+import android.util.SparseIntArray;
+import android.util.StatsEvent;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.util.FrameworkStatsLog;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Aggregates Mobile Data Usage by process state per uid
+ */
+class AggregatedMobileDataStatsPuller {
+    private static final String TAG = "AggregatedMobileDataStatsPuller";
+
+    private static final boolean DEBUG = false;
+
+    private static class UidProcState {
+
+        private final int mUid;
+        private final int mState;
+
+        UidProcState(int uid, int state) {
+            mUid = uid;
+            mState = state;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof UidProcState key)) return false;
+            return mUid == key.mUid && mState == key.mState;
+        }
+
+        @Override
+        public int hashCode() {
+            int result = mUid;
+            result = 31 * result + mState;
+            return result;
+        }
+
+        public int getUid() {
+            return mUid;
+        }
+
+        public int getState() {
+            return mState;
+        }
+
+    }
+
+    private static class MobileDataStats {
+        private long mRxPackets = 0;
+        private long mTxPackets = 0;
+        private long mRxBytes = 0;
+        private long mTxBytes = 0;
+
+        public long getRxPackets() {
+            return mRxPackets;
+        }
+
+        public long getTxPackets() {
+            return mTxPackets;
+        }
+
+        public long getRxBytes() {
+            return mRxBytes;
+        }
+
+        public long getTxBytes() {
+            return mTxBytes;
+        }
+
+        public void addRxPackets(long rxPackets) {
+            mRxPackets += rxPackets;
+        }
+
+        public void addTxPackets(long txPackets) {
+            mTxPackets += txPackets;
+        }
+
+        public void addRxBytes(long rxBytes) {
+            mRxBytes += rxBytes;
+        }
+
+        public void addTxBytes(long txBytes) {
+            mTxBytes += txBytes;
+        }
+
+        public boolean isEmpty() {
+            return mRxPackets == 0 && mTxPackets == 0 && mRxBytes == 0 && mTxBytes == 0;
+        }
+    }
+
+    private final Object mLock = new Object();
+    @GuardedBy("mLock")
+    private final Map<UidProcState, MobileDataStats> mUidStats;
+
+    private final SparseIntArray mUidPreviousState;
+
+    private NetworkStats mLastMobileUidStats = new NetworkStats(0, -1);
+
+    private final NetworkStatsManager mNetworkStatsManager;
+
+    private final Handler mMobileDataStatsHandler;
+
+    AggregatedMobileDataStatsPuller(NetworkStatsManager networkStatsManager) {
+        mUidStats = new ArrayMap<>();
+        mUidPreviousState = new SparseIntArray();
+
+        mNetworkStatsManager = networkStatsManager;
+
+        if (mNetworkStatsManager != null) {
+            updateNetworkStats(mNetworkStatsManager);
+        }
+
+        HandlerThread mMobileDataStatsHandlerThread = new HandlerThread("MobileDataStatsHandler");
+        mMobileDataStatsHandlerThread.start();
+        mMobileDataStatsHandler = new Handler(mMobileDataStatsHandlerThread.getLooper());
+    }
+
+    public void noteUidProcessState(int uid, int state, long unusedElapsedRealtime,
+                                    long unusedUptime) {
+        mMobileDataStatsHandler.post(
+                () -> {
+                    noteUidProcessStateImpl(uid, state);
+                });
+    }
+
+    public int pullDataBytesTransfer(List<StatsEvent> data) {
+        synchronized (mLock) {
+            return pullDataBytesTransferLocked(data);
+        }
+    }
+
+    @GuardedBy("mLock")
+    private MobileDataStats getUidStatsForPreviousStateLocked(int uid) {
+        final int previousState = mUidPreviousState.get(uid, ActivityManager.PROCESS_STATE_UNKNOWN);
+        if (DEBUG && previousState == ActivityManager.PROCESS_STATE_UNKNOWN) {
+            Slog.d(TAG, "getUidStatsForPreviousStateLocked() no prev state info for uid "
+                    + uid + ". Tracking stats with ActivityManager.PROCESS_STATE_UNKNOWN");
+        }
+
+        final UidProcState statsKey = new UidProcState(uid, previousState);
+        MobileDataStats stats;
+        if (mUidStats.containsKey(statsKey)) {
+            stats = mUidStats.get(statsKey);
+        } else {
+            stats = new MobileDataStats();
+            mUidStats.put(statsKey, stats);
+        }
+        return stats;
+    }
+
+    private void noteUidProcessStateImpl(int uid, int state) {
+        // noteUidProcessStateLocked can be called back to back several times while
+        // the updateNetworkStatsLocked loops over several stats for multiple uids
+        // and during the first call in a batch of proc state change event it can
+        // contain info for uid with unknown previous state yet which can happen due to a few
+        // reasons:
+        // - app was just started
+        // - app was started before the ActivityManagerService
+        // as result stats would be created with state == ActivityManager.PROCESS_STATE_UNKNOWN
+        if (mNetworkStatsManager != null) {
+            updateNetworkStats(mNetworkStatsManager);
+        } else {
+            Slog.w(TAG, "noteUidProcessStateLocked() can not get mNetworkStatsManager");
+        }
+        mUidPreviousState.put(uid, state);
+    }
+
+    private void updateNetworkStats(NetworkStatsManager networkStatsManager) {
+        if (DEBUG) {
+            if (Trace.isTagEnabled(Trace.TRACE_TAG_SYSTEM_SERVER)) {
+                Trace.traceBegin(Trace.TRACE_TAG_SYSTEM_SERVER, TAG + "-updateNetworkStats");
+            }
+        }
+
+        final NetworkStats latestStats = networkStatsManager.getMobileUidStats();
+        if (isEmpty(latestStats)) {
+            if (DEBUG) {
+                Slog.w(TAG, "getMobileUidStats() failed");
+                Trace.traceEnd(Trace.TRACE_TAG_SYSTEM_SERVER);
+            }
+            return;
+        }
+        NetworkStats delta = latestStats.subtract(mLastMobileUidStats);
+        mLastMobileUidStats = latestStats;
+
+        if (!isEmpty(delta)) {
+            updateNetworkStatsDelta(delta);
+        } else if (DEBUG) {
+            Slog.w(TAG, "updateNetworkStats() no delta");
+        }
+        if (DEBUG) {
+            Trace.traceEnd(Trace.TRACE_TAG_SYSTEM_SERVER);
+        }
+    }
+
+    private void updateNetworkStatsDelta(NetworkStats delta) {
+        synchronized (mLock) {
+            for (NetworkStats.Entry entry : delta) {
+                if (entry.getRxPackets() == 0 && entry.getTxPackets() == 0) {
+                    continue;
+                }
+                MobileDataStats stats = getUidStatsForPreviousStateLocked(entry.getUid());
+                stats.addTxBytes(entry.getTxBytes());
+                stats.addRxBytes(entry.getRxBytes());
+                stats.addTxPackets(entry.getTxPackets());
+                stats.addRxPackets(entry.getRxPackets());
+            }
+        }
+    }
+
+    @GuardedBy("mLock")
+    private int pullDataBytesTransferLocked(List<StatsEvent> pulledData) {
+        if (DEBUG) {
+            Slog.d(TAG, "pullDataBytesTransferLocked() start");
+        }
+        for (Map.Entry<UidProcState, MobileDataStats> uidStats : mUidStats.entrySet()) {
+            if (!uidStats.getValue().isEmpty()) {
+                MobileDataStats stats = uidStats.getValue();
+                pulledData.add(FrameworkStatsLog.buildStatsEvent(
+                        FrameworkStatsLog.MOBILE_BYTES_TRANSFER_BY_PROC_STATE,
+                        uidStats.getKey().getUid(),
+                        ActivityManager.processStateAmToProto(uidStats.getKey().getState()),
+                        stats.getRxBytes(),
+                        stats.getRxPackets(),
+                        stats.getTxBytes(),
+                        stats.getTxPackets()));
+            }
+        }
+        if (DEBUG) {
+            Slog.d(TAG,
+                    "pullDataBytesTransferLocked() done. results count " + pulledData.size());
+        }
+        if (!pulledData.isEmpty()) {
+            return StatsManager.PULL_SUCCESS;
+        }
+        return StatsManager.PULL_SKIP;
+    }
+
+    private static boolean isEmpty(NetworkStats stats) {
+        long totalRxPackets = 0;
+        long totalTxPackets = 0;
+        for (NetworkStats.Entry entry : stats) {
+            if (entry.getRxPackets() == 0 && entry.getTxPackets() == 0) {
+                continue;
+            }
+            totalRxPackets += entry.getRxPackets();
+            totalTxPackets += entry.getTxPackets();
+            // at least one non empty entry located
+            break;
+        }
+        final long totalPackets = totalRxPackets + totalTxPackets;
+        return totalPackets == 0;
+    }
+}
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 e876241f385e..285bcc328c0c 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -59,6 +59,7 @@ import static com.android.internal.util.FrameworkStatsLog.TIME_ZONE_DETECTOR_STA
 import static com.android.internal.util.FrameworkStatsLog.TIME_ZONE_DETECTOR_STATE__DETECTION_MODE__TELEPHONY;
 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.pull.IonMemoryUtil.readProcessSystemIonHeapSizesFromDebugfs;
 import static com.android.server.stats.pull.IonMemoryUtil.readSystemIonHeapSizeFromDebugfs;
 import static com.android.server.stats.pull.ProcfsMemoryUtil.getProcessCmdlines;
@@ -409,6 +410,15 @@ public class StatsPullAtomService extends SystemService {
     @GuardedBy("mKeystoreLock")
     private IKeystoreMetrics mIKeystoreMetrics;
 
+    private AggregatedMobileDataStatsPuller mAggregatedMobileDataStatsPuller = null;
+
+    /**
+     * Whether or not to enable the new puller with aggregation by process state per uid on a
+     * system server side.
+     */
+    public static final boolean ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER =
+                addMobileBytesTransferByProcStatePuller();
+
     // Puller locks
     private final Object mDataBytesTransferLock = new Object();
     private final Object mBluetoothBytesTransferLock = new Object();
@@ -469,6 +479,20 @@ public class StatsPullAtomService extends SystemService {
         mContext = context;
     }
 
+    private final class StatsPullAtomServiceInternalImpl extends StatsPullAtomServiceInternal {
+
+        @Override
+        public void noteUidProcessState(int uid, int state) {
+            if (ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER
+                    && mAggregatedMobileDataStatsPuller != null) {
+                final long elapsedRealtime = SystemClock.elapsedRealtime();
+                final long uptime = SystemClock.uptimeMillis();
+                mAggregatedMobileDataStatsPuller.noteUidProcessState(uid, state, elapsedRealtime,
+                        uptime);
+            }
+        }
+    }
+
     private native void initializeNativePullers();
 
     /**
@@ -486,6 +510,11 @@ public class StatsPullAtomService extends SystemService {
             }
             try {
                 switch (atomTag) {
+                    case FrameworkStatsLog.MOBILE_BYTES_TRANSFER_BY_PROC_STATE:
+                        if (ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER
+                                && mAggregatedMobileDataStatsPuller != null) {
+                            return mAggregatedMobileDataStatsPuller.pullDataBytesTransfer(data);
+                        }
                     case FrameworkStatsLog.WIFI_BYTES_TRANSFER:
                     case FrameworkStatsLog.WIFI_BYTES_TRANSFER_BY_FG_BG:
                     case FrameworkStatsLog.MOBILE_BYTES_TRANSFER:
@@ -776,7 +805,10 @@ public class StatsPullAtomService extends SystemService {
 
     @Override
     public void onStart() {
-        // no op
+        if (ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER) {
+            LocalServices.addService(StatsPullAtomServiceInternal.class,
+                    new StatsPullAtomServiceInternalImpl());
+        }
     }
 
     @Override
@@ -811,6 +843,9 @@ public class StatsPullAtomService extends SystemService {
         mStatsSubscriptionsListener = new StatsSubscriptionsListener(mSubscriptionManager);
         mStorageManager = (StorageManager) mContext.getSystemService(StorageManager.class);
         mNetworkStatsManager = mContext.getSystemService(NetworkStatsManager.class);
+
+        initMobileDataStatsPuller();
+
         // Initialize DiskIO
         mStoragedUidIoStatsReader = new StoragedUidIoStatsReader();
 
@@ -972,6 +1007,18 @@ public class StatsPullAtomService extends SystemService {
         registerCachedAppsHighWatermarkPuller();
     }
 
+    private void initMobileDataStatsPuller() {
+        if (DEBUG) {
+            Slog.d(TAG,
+                    "ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER = "
+                            + ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER);
+        }
+        if (ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER) {
+            mAggregatedMobileDataStatsPuller =
+                    new AggregatedMobileDataStatsPuller(mNetworkStatsManager);
+        }
+    }
+
     private void initAndRegisterNetworkStatsPullers() {
         if (DEBUG) {
             Slog.d(TAG, "Registering NetworkStats pullers with statsd");
@@ -1013,6 +1060,9 @@ public class StatsPullAtomService extends SystemService {
         registerWifiBytesTransferBackground();
         registerMobileBytesTransfer();
         registerMobileBytesTransferBackground();
+        if (ENABLE_MOBILE_DATA_STATS_AGGREGATED_PULLER) {
+            registerMobileBytesTransferByProcState();
+        }
         registerBytesTransferByTagAndMetered();
         registerDataUsageBytesTransfer();
         registerOemManagedBytesTransfer();
@@ -1021,6 +1071,13 @@ public class StatsPullAtomService extends SystemService {
         }
     }
 
+    private void registerMobileBytesTransferByProcState() {
+        final int tagId = FrameworkStatsLog.MOBILE_BYTES_TRANSFER_BY_PROC_STATE;
+        PullAtomMetadata metadata =
+                new PullAtomMetadata.Builder().setAdditiveFields(new int[] {3, 4, 5, 6}).build();
+        mStatsManager.setPullAtomCallback(tagId, metadata, DIRECT_EXECUTOR, mStatsCallbackImpl);
+    }
+
     private void initAndRegisterDeferredPullers() {
         mUwbManager = mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_UWB)
             ? mContext.getSystemService(UwbManager.class) : null;
diff --git a/services/core/java/com/android/server/stats/pull/StatsPullAtomServiceInternal.java b/services/core/java/com/android/server/stats/pull/StatsPullAtomServiceInternal.java
new file mode 100644
index 000000000000..06adbfc65abc
--- /dev/null
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomServiceInternal.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.stats.pull;
+
+/**
+ * System-server internal interface to the {@link StatsPullAtomService}.
+ *
+ * @hide Only for use within the system server.
+ */
+public abstract class StatsPullAtomServiceInternal {
+
+    /**
+     * @param state Process state from ActivityManager.java.
+     */
+    public abstract void noteUidProcessState(int uid, int state);
+
+}
diff --git a/services/core/java/com/android/server/stats/stats_flags.aconfig b/services/core/java/com/android/server/stats/stats_flags.aconfig
new file mode 100644
index 000000000000..5101a6982fe1
--- /dev/null
+++ b/services/core/java/com/android/server/stats/stats_flags.aconfig
@@ -0,0 +1,9 @@
+package: "com.android.server.stats"
+
+flag {
+    name: "add_mobile_bytes_transfer_by_proc_state_puller"
+    namespace: "statsd"
+    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