Allow multiple standby buckets to be set in one IPC

SystemApi to pass a map of packagename to bucket mapping.
Reduces the number of IPC calls to be made by bucketing
ML app.

Bug: 63527785
Test: atest CtsAppUsageHostTestCases
Change-Id: I4172fd30cc310092a1f182ae26267de9148009b7

7 files changed, 168 insertions, 10 deletions

diff --git a/api/system-current.txt b/api/system-current.txt
index 33fa246a9f51..7e514fd59c07 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -613,7 +613,9 @@ package android.app.usage {
 
   public final class UsageStatsManager {
     method public int getAppStandbyBucket(java.lang.String);
+    method public java.util.Map<java.lang.String, java.lang.Integer> getAppStandbyBuckets();
     method public void setAppStandbyBucket(java.lang.String, int);
+    method public void setAppStandbyBuckets(java.util.Map<java.lang.String, java.lang.Integer>);
     method public void whitelistAppTemporarily(java.lang.String, long, android.os.UserHandle);
     field public static final int STANDBY_BUCKET_EXEMPTED = 5; // 0x5
     field public static final int STANDBY_BUCKET_NEVER = 50; // 0x32
diff --git a/core/java/android/app/usage/IUsageStatsManager.aidl b/core/java/android/app/usage/IUsageStatsManager.aidl
index 4fbbdf2a9281..f089c127d03f 100644
--- a/core/java/android/app/usage/IUsageStatsManager.aidl
+++ b/core/java/android/app/usage/IUsageStatsManager.aidl
@@ -19,6 +19,8 @@ package android.app.usage;
 import android.app.usage.UsageEvents;
 import android.content.pm.ParceledListSlice;
 
+import java.util.Map;
+
 /**
  * System private API for talking with the UsageStatsManagerService.
  *
@@ -38,4 +40,6 @@ interface IUsageStatsManager {
             in String[] annotations, String action);
     int getAppStandbyBucket(String packageName, String callingPackage, int userId);
     void setAppStandbyBucket(String packageName, int bucket, int userId);
+    Map getAppStandbyBuckets(String callingPackage, int userId);
+    void setAppStandbyBuckets(in Map appBuckets, int userId);
 }
diff --git a/core/java/android/app/usage/UsageStatsManager.java b/core/java/android/app/usage/UsageStatsManager.java
index d614b20a0788..1fc45c9fdb94 100644
--- a/core/java/android/app/usage/UsageStatsManager.java
+++ b/core/java/android/app/usage/UsageStatsManager.java
@@ -324,11 +324,11 @@ public final class UsageStatsManager {
      * state of the app based on app usage patterns. Standby buckets determine how much an app will
      * be restricted from running background tasks such as jobs, alarms and certain PendingIntent
      * callbacks.
-     * Restrictions increase progressively from {@link #STANDBY_BUCKET_ACTIVE} to
+     * <p>Restrictions increase progressively from {@link #STANDBY_BUCKET_ACTIVE} to
      * {@link #STANDBY_BUCKET_RARE}, with {@link #STANDBY_BUCKET_ACTIVE} being the least
      * restrictive. The battery level of the device might also affect the restrictions.
      *
-     * @return the current standby bucket of the calling app.
+     * @return the current standby bucket of the calling app. One of STANDBY_BUCKET_* constants.
      */
     public @StandbyBuckets int getAppStandbyBucket() {
         try {
@@ -359,7 +359,13 @@ public final class UsageStatsManager {
 
     /**
      * {@hide}
-     * Changes the app standby state to the provided bucket.
+     * Changes an app's standby bucket to the provided value. The caller can only set the standby
+     * bucket for a different app than itself.
+     * @param packageName the package name of the app to set the bucket for. A SecurityException
+     *                    will be thrown if the package name is that of the caller.
+     * @param bucket the standby bucket to set it to, which should be one of STANDBY_BUCKET_*.
+     *               Setting a standby bucket outside of the range of STANDBY_BUCKET_ACTIVE to
+     *               STANDBY_BUCKET_NEVER will result in a SecurityException.
      */
     @SystemApi
     @RequiresPermission(android.Manifest.permission.CHANGE_APP_IDLE_STATE)
@@ -373,6 +379,39 @@ public final class UsageStatsManager {
 
     /**
      * {@hide}
+     * Returns the current standby bucket of every app that has a bucket assigned to it.
+     * The caller must hold the permission android.permission.PACKAGE_USAGE_STATS. The key of the
+     * returned Map is the package name and the value is the bucket assigned to the package.
+     * @see #getAppStandbyBucket()
+     */
+    @SystemApi
+    @RequiresPermission(android.Manifest.permission.PACKAGE_USAGE_STATS)
+    public Map<String, Integer> getAppStandbyBuckets() {
+        try {
+            return (Map<String, Integer>) mService.getAppStandbyBuckets(
+                    mContext.getOpPackageName(), mContext.getUserId());
+        } catch (RemoteException e) {
+        }
+        return Collections.EMPTY_MAP;
+    }
+
+    /**
+     * {@hide}
+     * Changes the app standby bucket for multiple apps at once. The Map is keyed by the package
+     * name and the value is one of STANDBY_BUCKET_*.
+     * @param appBuckets a map of package name to bucket value.
+     */
+    @SystemApi
+    @RequiresPermission(android.Manifest.permission.CHANGE_APP_IDLE_STATE)
+    public void setAppStandbyBuckets(Map<String, Integer> appBuckets) {
+        try {
+            mService.setAppStandbyBuckets(appBuckets, mContext.getUserId());
+        } catch (RemoteException e) {
+        }
+    }
+
+    /**
+     * {@hide}
      * Temporarily whitelist the specified app for a short duration. This is to allow an app
      * receiving a high priority message to be able to access the network and acquire wakelocks
      * even if the device is in power-save mode or the app is currently considered inactive.
diff --git a/services/core/java/com/android/server/am/ActivityManagerShellCommand.java b/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
index 979323fc8408..54938eb05cf5 100644
--- a/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
+++ b/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
@@ -69,7 +69,9 @@ import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static android.app.ActivityManager.RESIZE_MODE_SYSTEM;
 import static android.app.ActivityManager.RESIZE_MODE_USER;
@@ -1867,10 +1869,24 @@ final class ActivityManagerShellCommand extends ShellCommand {
         String value = getNextArgRequired();
         int bucket = bucketNameToBucketValue(value);
         if (bucket < 0) return -1;
+        boolean multiple = peekNextArg() != null;
+
 
         IUsageStatsManager usm = IUsageStatsManager.Stub.asInterface(ServiceManager.getService(
                 Context.USAGE_STATS_SERVICE));
-        usm.setAppStandbyBucket(packageName, bucketNameToBucketValue(value), userId);
+        if (!multiple) {
+            usm.setAppStandbyBucket(packageName, bucketNameToBucketValue(value), userId);
+        } else {
+            HashMap<String, Integer> buckets = new HashMap<>();
+            buckets.put(packageName, bucket);
+            while ((packageName = getNextArg()) != null) {
+                value = getNextArgRequired();
+                bucket = bucketNameToBucketValue(value);
+                if (bucket < 0) continue;
+                buckets.put(packageName, bucket);
+            }
+            usm.setAppStandbyBuckets(buckets, userId);
+        }
         return 0;
     }
 
@@ -1886,12 +1902,21 @@ final class ActivityManagerShellCommand extends ShellCommand {
                 return -1;
             }
         }
-        String packageName = getNextArgRequired();
+        String packageName = getNextArg();
 
         IUsageStatsManager usm = IUsageStatsManager.Stub.asInterface(ServiceManager.getService(
                 Context.USAGE_STATS_SERVICE));
-        int bucket = usm.getAppStandbyBucket(packageName, null, userId);
-        pw.println(bucket);
+        if (packageName != null) {
+            int bucket = usm.getAppStandbyBucket(packageName, null, userId);
+            pw.println(bucket);
+        } else {
+            Map<String, Integer> buckets = (Map<String, Integer>) usm.getAppStandbyBuckets(
+                    SHELL_PACKAGE_NAME, userId);
+            for (Map.Entry<String, Integer> entry: buckets.entrySet()) {
+                pw.print(entry.getKey()); pw.print(": ");
+                pw.println(entry.getValue());
+            }
+        }
         return 0;
     }
 
diff --git a/services/usage/java/com/android/server/usage/AppIdleHistory.java b/services/usage/java/com/android/server/usage/AppIdleHistory.java
index 6ac4b36a6ad7..620922ac1363 100644
--- a/services/usage/java/com/android/server/usage/AppIdleHistory.java
+++ b/services/usage/java/com/android/server/usage/AppIdleHistory.java
@@ -42,6 +42,8 @@ import java.io.FileOutputStream;
 import java.io.FileReader;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * Keeps track of recent active state changes in apps.
@@ -334,6 +336,16 @@ public class AppIdleHistory {
         return appUsageHistory.currentBucket;
     }
 
+    public Map<String, Integer> getAppStandbyBuckets(int userId, long elapsedRealtime) {
+        ArrayMap<String, AppUsageHistory> userHistory = getUserHistory(userId);
+        int size = userHistory.size();
+        HashMap<String, Integer> buckets = new HashMap<>(size);
+        for (int i = 0; i < size; i++) {
+            buckets.put(userHistory.keyAt(i), userHistory.valueAt(i).currentBucket);
+        }
+        return buckets;
+    }
+
     public String getAppStandbyReason(String packageName, int userId, long elapsedRealtime) {
         ArrayMap<String, AppUsageHistory> userHistory = getUserHistory(userId);
         AppUsageHistory appUsageHistory =
diff --git a/services/usage/java/com/android/server/usage/AppStandbyController.java b/services/usage/java/com/android/server/usage/AppStandbyController.java
index d8086bb4a92d..46efbd059d52 100644
--- a/services/usage/java/com/android/server/usage/AppStandbyController.java
+++ b/services/usage/java/com/android/server/usage/AppStandbyController.java
@@ -84,6 +84,7 @@ import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Manages the standby state of an app, listening to various events.
@@ -774,13 +775,23 @@ public class AppStandbyController {
             return STANDBY_BUCKET_ACTIVE;
         }
 
-        return mAppIdleHistory.getAppStandbyBucket(packageName, userId, elapsedRealtime);
+        synchronized (mAppIdleLock) {
+            return mAppIdleHistory.getAppStandbyBucket(packageName, userId, elapsedRealtime);
+        }
+    }
+
+    public Map<String, Integer> getAppStandbyBuckets(int userId, long elapsedRealtime) {
+        synchronized (mAppIdleLock) {
+            return mAppIdleHistory.getAppStandbyBuckets(userId, elapsedRealtime);
+        }
     }
 
     void setAppStandbyBucket(String packageName, int userId, @StandbyBuckets int newBucket,
             String reason, long elapsedRealtime) {
-        mAppIdleHistory.setAppStandbyBucket(packageName, userId, elapsedRealtime, newBucket,
-                reason);
+        synchronized (mAppIdleLock) {
+            mAppIdleHistory.setAppStandbyBucket(packageName, userId, elapsedRealtime, newBucket,
+                    reason);
+        }
         maybeInformListeners(packageName, userId, elapsedRealtime,
                 newBucket);
     }
diff --git a/services/usage/java/com/android/server/usage/UsageStatsService.java b/services/usage/java/com/android/server/usage/UsageStatsService.java
index 0572771ab6e0..15284d5e3dec 100644
--- a/services/usage/java/com/android/server/usage/UsageStatsService.java
+++ b/services/usage/java/com/android/server/usage/UsageStatsService.java
@@ -68,7 +68,10 @@ import java.io.FileDescriptor;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 
 /**
  * A service that collects, aggregates, and persists application usage data.
@@ -741,6 +744,68 @@ public class UsageStatsService extends SystemService implements
         }
 
         @Override
+        public Map getAppStandbyBuckets(String callingPackageName, int userId) {
+            final int callingUid = Binder.getCallingUid();
+            try {
+                userId = ActivityManager.getService().handleIncomingUser(
+                        Binder.getCallingPid(), callingUid, userId, false, false,
+                        "getAppStandbyBucket", null);
+            } catch (RemoteException re) {
+                throw re.rethrowFromSystemServer();
+            }
+            if (!hasPermission(callingPackageName)) {
+                throw new SecurityException(
+                        "Don't have permission to query app standby bucket");
+            }
+            final long token = Binder.clearCallingIdentity();
+            try {
+                return mAppStandby.getAppStandbyBuckets(userId,
+                        SystemClock.elapsedRealtime());
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+        }
+
+        @Override
+        public void setAppStandbyBuckets(Map appBuckets, int userId) {
+            getContext().enforceCallingPermission(Manifest.permission.CHANGE_APP_IDLE_STATE,
+                    "No permission to change app standby state");
+
+            final int callingUid = Binder.getCallingUid();
+            try {
+                userId = ActivityManager.getService().handleIncomingUser(
+                        Binder.getCallingPid(), callingUid, userId, false, true,
+                        "setAppStandbyBucket", null);
+            } catch (RemoteException re) {
+                throw re.rethrowFromSystemServer();
+            }
+            final long token = Binder.clearCallingIdentity();
+            try {
+                final long elapsedRealtime = SystemClock.elapsedRealtime();
+                Map<String, Integer> buckets = (Map<String, Integer>) appBuckets;
+                for (Map.Entry<String, Integer> entry: buckets.entrySet()) {
+                    String packageName = entry.getKey();
+                    int bucket = entry.getValue();
+                    if (bucket < UsageStatsManager.STANDBY_BUCKET_ACTIVE
+                            || bucket > UsageStatsManager.STANDBY_BUCKET_NEVER) {
+                        throw new IllegalArgumentException(
+                                "Cannot set the standby bucket to " + bucket);
+                    }
+                    // Caller cannot set their own standby state
+                    if (mPackageManagerInternal.getPackageUid(packageName,
+                            PackageManager.MATCH_ANY_USER, userId) == callingUid) {
+                        throw new IllegalArgumentException("Cannot set your own standby bucket");
+                    }
+                    mAppStandby.setAppStandbyBucket(packageName, userId, bucket,
+                            UsageStatsManager.REASON_PREDICTED + ":" + callingUid,
+                            elapsedRealtime);
+                }
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+        }
+
+        @Override
         public void whitelistAppTemporarily(String packageName, long duration, int userId)
                 throws RemoteException {
             StringBuilder reason = new StringBuilder(32);