From 31b9785e9810604784b75322980cd3826b0186ad Mon Sep 17 00:00:00 2001
From: Jing Ji <jji@google.com>
Date: Wed, 16 Mar 2022 22:32:01 -0700
Subject: [PATCH] Consolidate the bg exemption state changes from the pkgs with
 same UID

As the background battery exemption tracker is per-UID basis, while
the exemption state changes are per UID/package.

Bug: 203105544
Test: atest FrameworksMockingServicesTests:BackgroundRestrictionTest
Change-Id: Id928e868137e641443305aa364c4672280b9ae71
---
 .../server/am/AppBatteryExemptionTracker.java | 82 +++++++++++++++++--
 1 file changed, 77 insertions(+), 5 deletions(-)

diff --git a/services/core/java/com/android/server/am/AppBatteryExemptionTracker.java b/services/core/java/com/android/server/am/AppBatteryExemptionTracker.java
index d1cf0049d1468..6f11b0001c7a2 100644
--- a/services/core/java/com/android/server/am/AppBatteryExemptionTracker.java
+++ b/services/core/java/com/android/server/am/AppBatteryExemptionTracker.java
@@ -26,10 +26,13 @@ import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
 import android.os.SystemClock;
+import android.util.ArrayMap;
 import android.util.Pair;
 import android.util.Slog;
+import android.util.SparseArray;
 
 import com.android.internal.R;
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.am.AppBatteryExemptionTracker.AppBatteryExemptionPolicy;
 import com.android.server.am.AppBatteryExemptionTracker.UidBatteryStates;
@@ -65,6 +68,11 @@ final class AppBatteryExemptionTracker
     // As it's a UID-based tracker, anywhere which requires a package name, use this default name.
     static final String DEFAULT_NAME = "";
 
+    // As it's a UID-based tracker, while the state change event it receives could be
+    // in the combination of UID + package name, we'd have to leverage each package's state.
+    @GuardedBy("mLock")
+    private UidProcessMap<Integer> mUidPackageStates = new UidProcessMap<>();
+
     AppBatteryExemptionTracker(Context context, AppRestrictionController controller) {
         this(context, controller, null, null);
     }
@@ -103,12 +111,75 @@ final class AppBatteryExemptionTracker
                 .getUidBatteryUsage(uid);
         final int stateTypeIndex = stateTypeToIndex(stateType);
         synchronized (mLock) {
-            UidBatteryStates pkg = mPkgEvents.get(uid, DEFAULT_NAME);
-            if (pkg == null) {
-                pkg = createAppStateEvents(uid, DEFAULT_NAME);
-                mPkgEvents.put(uid, DEFAULT_NAME, pkg);
+            final SparseArray<ArrayMap<String, Integer>> map = mUidPackageStates.getMap();
+            ArrayMap<String, Integer> pkgsStates = map.get(uid);
+            if (pkgsStates == null) {
+                pkgsStates = new ArrayMap<>();
+                map.put(uid, pkgsStates);
             }
-            pkg.addEvent(start, now, batteryUsage, stateTypeIndex);
+            int states = 0;
+            int indexOfPkg = pkgsStates.indexOfKey(packageName);
+            if (indexOfPkg >= 0) {
+                states = pkgsStates.valueAt(indexOfPkg);
+            } else {
+                pkgsStates.put(packageName, 0);
+                indexOfPkg = pkgsStates.indexOfKey(packageName);
+            }
+            boolean addEvent = false;
+            if (start) {
+                // Check if there is another package within this UID with this type of event start.
+                boolean alreadyStarted = false;
+                for (int i = pkgsStates.size() - 1; i >= 0; i--) {
+                    final int s = pkgsStates.valueAt(i);
+                    if ((s & stateType) != 0) {
+                        alreadyStarted = true;
+                        break;
+                    }
+                }
+                pkgsStates.setValueAt(indexOfPkg, states | stateType);
+                if (!alreadyStarted) {
+                    // This is the first package within this UID with this type of event start.
+                    addEvent = true;
+                }
+            } else {
+                states &= ~stateType;
+                pkgsStates.setValueAt(indexOfPkg, states);
+                boolean allStopped = true;
+                for (int i = pkgsStates.size() - 1; i >= 0; i--) {
+                    final int s = pkgsStates.valueAt(i);
+                    if ((s & stateType) != 0) {
+                        allStopped = false;
+                        break;
+                    }
+                }
+                if (allStopped) {
+                    // None of the packages in this UID has an active event of this type.
+                    addEvent = true;
+                }
+                if (states == 0) { // None of the states of this package are active, prune it.
+                    pkgsStates.removeAt(indexOfPkg);
+                    if (pkgsStates.size() == 0) {
+                        map.remove(uid);
+                    }
+                }
+            }
+            if (addEvent) {
+                UidBatteryStates pkg = mPkgEvents.get(uid, DEFAULT_NAME);
+                if (pkg == null) {
+                    pkg = createAppStateEvents(uid, DEFAULT_NAME);
+                    mPkgEvents.put(uid, DEFAULT_NAME, pkg);
+                }
+                pkg.addEvent(start, now, batteryUsage, stateTypeIndex);
+            }
+        }
+    }
+
+    @VisibleForTesting
+    @Override
+    void reset() {
+        super.reset();
+        synchronized (mLock) {
+            mUidPackageStates.clear();
         }
     }
 
@@ -116,6 +187,7 @@ final class AppBatteryExemptionTracker
         if (!enabled) {
             synchronized (mLock) {
                 mPkgEvents.clear();
+                mUidPackageStates.clear();
             }
         }
     }