Migrate network stats from removed users.

When a user is removed, migrate all network stats belonging to that
user into special UID_REMOVED bucket.  Also removes those stats from
kernel to avoid double-counting if another user is created.

Bug: 7194784
Change-Id: I03f1d660fe3754566326b7749cae8068fc224ea9

core/java/android/net/NetworkStats.java

@@ -21,6 +21,7 @@ import android.os.Parcelable;
 import android.os.SystemClock;
 import android.util.SparseBooleanArray;
 
+import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.Objects;
 
 import java.io.CharArrayWriter;
@@ -608,13 +609,13 @@ public class NetworkStats implements Parcelable {
      * Return all rows except those attributed to the requested UID; doesn't
      * mutate the original structure.
      */
-    public NetworkStats withoutUid(int uid) {
+    public NetworkStats withoutUids(int[] uids) {
         final NetworkStats stats = new NetworkStats(elapsedRealtime, 10);
 
         Entry entry = new Entry();
         for (int i = 0; i < size; i++) {
             entry = getValues(i, entry);
-            if (entry.uid != uid) {
+            if (!ArrayUtils.contains(uids, entry.uid)) {
                 stats.addValues(entry);
             }
         }

services/java/com/android/server/net/NetworkStatsCollection.java

@@ -31,6 +31,7 @@ import android.net.TrafficStats;
 import android.text.format.DateUtils;
 import android.util.AtomicFile;
 
+import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.FileRotator;
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.internal.util.Objects;
@@ -431,13 +432,13 @@ public class NetworkStatsCollection implements FileRotator.Reader {
      * moving any {@link NetworkStats#TAG_NONE} series to
      * {@link TrafficStats#UID_REMOVED}.
      */
-    public void removeUid(int uid) {
+    public void removeUids(int[] uids) {
         final ArrayList<Key> knownKeys = Lists.newArrayList();
         knownKeys.addAll(mStats.keySet());
 
         // migrate all UID stats into special "removed" bucket
         for (Key key : knownKeys) {
-            if (key.uid == uid) {
+            if (ArrayUtils.contains(uids, key.uid)) {
                 // only migrate combined TAG_NONE history
                 if (key.tag == TAG_NONE) {
                     final NetworkStatsHistory uidHistory = mStats.get(key);

services/java/com/android/server/net/NetworkStatsRecorder.java

@@ -42,6 +42,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.lang.ref.WeakReference;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Map;
 
@@ -233,23 +234,27 @@ public class NetworkStatsRecorder {
      * Remove the given UID from all {@link FileRotator} history, migrating it
      * to {@link TrafficStats#UID_REMOVED}.
      */
-    public void removeUidLocked(int uid) {
+    public void removeUidsLocked(int[] uids) {
         try {
-            // process all existing data to migrate uid
+            // Rewrite all persisted data to migrate UID stats
-            mRotator.rewriteAll(new RemoveUidRewriter(mBucketDuration, uid));
+            mRotator.rewriteAll(new RemoveUidRewriter(mBucketDuration, uids));
         } catch (IOException e) {
-            Log.wtf(TAG, "problem removing UID " + uid, e);
+            Log.wtf(TAG, "problem removing UIDs " + Arrays.toString(uids), e);
             recoverFromWtf();
         }
 
-        // clear UID from current stats snapshot
+        // Remove any pending stats
+        mPending.removeUids(uids);
+        mSinceBoot.removeUids(uids);
+
+        // Clear UID from current stats snapshot
         if (mLastSnapshot != null) {
-            mLastSnapshot = mLastSnapshot.withoutUid(uid);
+            mLastSnapshot = mLastSnapshot.withoutUids(uids);
         }
 
         final NetworkStatsCollection complete = mComplete != null ? mComplete.get() : null;
         if (complete != null) {
-            complete.removeUid(uid);
+            complete.removeUids(uids);
         }
     }
 
@@ -293,11 +298,11 @@ public class NetworkStatsRecorder {
      */
     public static class RemoveUidRewriter implements FileRotator.Rewriter {
         private final NetworkStatsCollection mTemp;
-        private final int mUid;
+        private final int[] mUids;
 
-        public RemoveUidRewriter(long bucketDuration, int uid) {
+        public RemoveUidRewriter(long bucketDuration, int[] uids) {
             mTemp = new NetworkStatsCollection(bucketDuration);
-            mUid = uid;
+            mUids = uids;
         }
 
         @Override
@@ -309,7 +314,7 @@ public class NetworkStatsRecorder {
         public void read(InputStream in) throws IOException {
             mTemp.read(in);
             mTemp.clearDirty();
-            mTemp.removeUid(mUid);
+            mTemp.removeUids(mUids);
         }
 
         @Override

services/java/com/android/server/net/NetworkStatsService.java

@@ -23,6 +23,7 @@ import static android.Manifest.permission.MODIFY_NETWORK_ACCOUNTING;
 import static android.Manifest.permission.READ_NETWORK_USAGE_HISTORY;
 import static android.content.Intent.ACTION_SHUTDOWN;
 import static android.content.Intent.ACTION_UID_REMOVED;
+import static android.content.Intent.ACTION_USER_REMOVED;
 import static android.content.Intent.EXTRA_UID;
 import static android.net.ConnectivityManager.ACTION_TETHER_STATE_CHANGED;
 import static android.net.ConnectivityManager.CONNECTIVITY_ACTION_IMMEDIATE;
@@ -76,6 +77,8 @@ import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
 import android.net.IConnectivityManager;
 import android.net.INetworkManagementEventObserver;
 import android.net.INetworkStatsService;
@@ -112,6 +115,7 @@ import android.util.Slog;
 import android.util.SparseIntArray;
 import android.util.TrustedTime;
 
+import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.FileRotator;
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.server.EventLogTags;
@@ -122,8 +126,10 @@ import java.io.File;
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.io.PrintWriter;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 
 /**
  * Collect and persist detailed network statistics, and provide this data to
@@ -322,6 +328,10 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         final IntentFilter removedFilter = new IntentFilter(ACTION_UID_REMOVED);
         mContext.registerReceiver(mRemovedReceiver, removedFilter, null, mHandler);
 
+        // listen for user changes to clean stats
+        final IntentFilter userFilter = new IntentFilter(ACTION_USER_REMOVED);
+        mContext.registerReceiver(mUserReceiver, userFilter, null, mHandler);
+
         // persist stats during clean shutdown
         final IntentFilter shutdownFilter = new IntentFilter(ACTION_SHUTDOWN);
         mContext.registerReceiver(mShutdownReceiver, shutdownFilter);
@@ -739,11 +749,34 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         public void onReceive(Context context, Intent intent) {
             // on background handler thread, and UID_REMOVED is protected
             // broadcast.
-            final int uid = intent.getIntExtra(EXTRA_UID, 0);
+
+            final int uid = intent.getIntExtra(EXTRA_UID, -1);
+            if (uid == -1) return;
+
             synchronized (mStatsLock) {
                 mWakeLock.acquire();
                 try {
-                    removeUidLocked(uid);
+                    removeUidsLocked(uid);
+                } finally {
+                    mWakeLock.release();
+                }
+            }
+        }
+    };
+
+    private BroadcastReceiver mUserReceiver = new BroadcastReceiver() {
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            // On background handler thread, and USER_REMOVED is protected
+            // broadcast.
+
+            final int userId = intent.getIntExtra(Intent.EXTRA_USER_HANDLE, -1);
+            if (userId == -1) return;
+
+            synchronized (mStatsLock) {
+                mWakeLock.acquire();
+                try {
+                    removeUserLocked(userId);
                 } finally {
                     mWakeLock.release();
                 }
@@ -1034,16 +1067,38 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
     /**
      * Clean up {@link #mUidRecorder} after UID is removed.
      */
-    private void removeUidLocked(int uid) {
+    private void removeUidsLocked(int... uids) {
-        // perform one last poll before removing
+        if (LOGV) Slog.v(TAG, "removeUidsLocked() for UIDs " + Arrays.toString(uids));
+
+        // Perform one last poll before removing
         performPollLocked(FLAG_PERSIST_ALL);
 
-        mUidRecorder.removeUidLocked(uid);
+        mUidRecorder.removeUidsLocked(uids);
-        mUidTagRecorder.removeUidLocked(uid);
+        mUidTagRecorder.removeUidsLocked(uids);
 
-        // clear kernel stats associated with UID
+        // Clear kernel stats associated with UID
+        for (int uid : uids) {
             resetKernelUidStats(uid);
         }
+    }
+
+    /**
+     * Clean up {@link #mUidRecorder} after user is removed.
+     */
+    private void removeUserLocked(int userId) {
+        if (LOGV) Slog.v(TAG, "removeUserLocked() for userId=" + userId);
+
+        // Build list of UIDs that we should clean up
+        int[] uids = new int[0];
+        final List<ApplicationInfo> apps = mContext.getPackageManager().getInstalledApplications(
+                PackageManager.GET_UNINSTALLED_PACKAGES | PackageManager.GET_DISABLED_COMPONENTS);
+        for (ApplicationInfo app : apps) {
+            final int uid = UserHandle.getUid(userId, app.uid);
+            uids = ArrayUtils.appendInt(uids, uid);
+        }
+
+        removeUidsLocked(uids);
+    }
 
     @Override
     protected void dump(FileDescriptor fd, PrintWriter writer, String[] args) {