Allow limited alarms/jobs when background restricted

When app is background restricted, allow alarms/jobs but they'll be
under the quota control of restricted standby bucket.

This behavior is gated by a feature flag now.

Bug: 200326767
Test: atest CtsJobSchedulerTestCases
Test: atest CtsAlarmManagerTestCases
Test: atest FrameworksMockingServicesTests
Change-Id: I154656c19954a306e8ae05dc50ea708c4de2a739
diff --git a/apex/jobscheduler/service/java/com/android/server/AppStateTrackerImpl.java b/apex/jobscheduler/service/java/com/android/server/AppStateTrackerImpl.java
index d21a0ea..70d5038 100644
--- a/apex/jobscheduler/service/java/com/android/server/AppStateTrackerImpl.java
+++ b/apex/jobscheduler/service/java/com/android/server/AppStateTrackerImpl.java
@@ -18,6 +18,7 @@
 import android.annotation.NonNull;
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal;
+import android.app.ActivityManagerInternal.AppBackgroundRestrictionListener;
 import android.app.AppOpsManager;
 import android.app.AppOpsManager.PackageOps;
 import android.app.IActivityManager;
@@ -280,6 +281,14 @@
         }
     }
 
+    private final AppBackgroundRestrictionListener mAppBackgroundRestrictionListener =
+            new AppBackgroundRestrictionListener() {
+        @Override
+        public void onAutoRestrictedBucketFeatureFlagChanged(boolean autoRestrictedBucket) {
+            mHandler.notifyAutoRestrictedBucketFeatureFlagChanged(autoRestrictedBucket);
+        }
+    };
+
     /**
      * Listener for any state changes that affect any app's eligibility to run.
      */
@@ -370,6 +379,18 @@
         }
 
         /**
+         * Called when toggling the feature flag of moving to restricted standby bucket
+         * automatically on background-restricted.
+         */
+        private void onAutoRestrictedBucketFeatureFlagChanged(AppStateTrackerImpl sender,
+                boolean autoRestrictedBucket) {
+            updateAllJobs();
+            if (autoRestrictedBucket) {
+                unblockAllUnrestrictedAlarms();
+            }
+        }
+
+        /**
          * Called when the job restrictions for multiple UIDs might have changed, so the job
          * scheduler should re-evaluate all restrictions for all jobs.
          */
@@ -499,6 +520,8 @@
                     mFlagsObserver.isForcedAppStandbyForSmallBatteryEnabled();
             mStandbyTracker = new StandbyTracker();
             mAppStandbyInternal.addListener(mStandbyTracker);
+            mActivityManagerInternal.addAppBackgroundRestrictionListener(
+                    mAppBackgroundRestrictionListener);
 
             try {
                 mIActivityManager.registerUidObserver(new UidObserver(),
@@ -802,6 +825,7 @@
         private static final int MSG_USER_REMOVED = 8;
         private static final int MSG_FORCE_APP_STANDBY_FEATURE_FLAG_CHANGED = 9;
         private static final int MSG_EXEMPTED_BUCKET_CHANGED = 10;
+        private static final int MSG_AUTO_RESTRICTED_BUCKET_FEATURE_FLAG_CHANGED = 11;
 
         private static final int MSG_ON_UID_ACTIVE = 12;
         private static final int MSG_ON_UID_GONE = 13;
@@ -849,6 +873,12 @@
             obtainMessage(MSG_EXEMPTED_BUCKET_CHANGED).sendToTarget();
         }
 
+        public void notifyAutoRestrictedBucketFeatureFlagChanged(boolean autoRestrictedBucket) {
+            removeMessages(MSG_AUTO_RESTRICTED_BUCKET_FEATURE_FLAG_CHANGED);
+            obtainMessage(MSG_AUTO_RESTRICTED_BUCKET_FEATURE_FLAG_CHANGED,
+                    autoRestrictedBucket ? 1 : 0, 0).sendToTarget();
+        }
+
         public void doUserRemoved(int userId) {
             obtainMessage(MSG_USER_REMOVED, userId, 0).sendToTarget();
         }
@@ -952,6 +982,13 @@
                     handleUserRemoved(msg.arg1);
                     return;
 
+                case MSG_AUTO_RESTRICTED_BUCKET_FEATURE_FLAG_CHANGED:
+                    final boolean autoRestrictedBucket = msg.arg1 == 1;
+                    for (Listener l : cloneListeners()) {
+                        l.onAutoRestrictedBucketFeatureFlagChanged(sender, autoRestrictedBucket);
+                    }
+                    return;
+
                 case MSG_ON_UID_ACTIVE:
                     handleUidActive(msg.arg1);
                     return;
@@ -1120,7 +1157,12 @@
             if (ArrayUtils.contains(mPowerExemptAllAppIds, appId)) {
                 return false;
             }
-            return (mForcedAppStandbyEnabled && isRunAnyRestrictedLocked(uid, packageName));
+            // If apps will be put into restricted standby bucket automatically on user-forced
+            // app standby, instead of blocking alarms completely, let the restricted standby bucket
+            // policy take care of it.
+            return (mForcedAppStandbyEnabled
+                    && !mActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled()
+                    && isRunAnyRestrictedLocked(uid, packageName));
         }
     }
 
@@ -1161,7 +1203,12 @@
                     || ArrayUtils.contains(mTempExemptAppIds, appId)) {
                 return false;
             }
-            if (mForcedAppStandbyEnabled && isRunAnyRestrictedLocked(uid, packageName)) {
+            // If apps will be put into restricted standby bucket automatically on user-forced
+            // app standby, instead of blocking jobs completely, let the restricted standby bucket
+            // policy take care of it.
+            if (mForcedAppStandbyEnabled
+                    && !mActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled()
+                    && isRunAnyRestrictedLocked(uid, packageName)) {
                 return true;
             }
             if (hasForegroundExemption) {
diff --git a/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java b/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java
index 0ceab35..65d7121 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java
@@ -19,6 +19,7 @@
 import static com.android.server.job.JobSchedulerService.NEVER_INDEX;
 import static com.android.server.job.JobSchedulerService.sElapsedRealtimeClock;
 
+import android.app.ActivityManagerInternal;
 import android.os.SystemClock;
 import android.os.UserHandle;
 import android.util.ArraySet;
@@ -59,6 +60,7 @@
     static final int KNOWN_ACTIVE = 1;
     static final int KNOWN_INACTIVE = 2;
 
+    private final ActivityManagerInternal mActivityManagerInternal;
     private final AppStateTrackerImpl mAppStateTracker;
 
     private final UpdateJobFunctor mUpdateJobFunctor = new UpdateJobFunctor();
@@ -66,6 +68,8 @@
     public BackgroundJobsController(JobSchedulerService service) {
         super(service);
 
+        mActivityManagerInternal = (ActivityManagerInternal) Objects.requireNonNull(
+                LocalServices.getService(ActivityManagerInternal.class));
         mAppStateTracker = (AppStateTrackerImpl) Objects.requireNonNull(
                 LocalServices.getService(AppStateTracker.class));
         mAppStateTracker.addListener(mForceAppStandbyListener);
@@ -216,7 +220,8 @@
         }
         boolean didChange =
                 jobStatus.setBackgroundNotRestrictedConstraintSatisfied(nowElapsed, canRun,
-                        !mAppStateTracker.isRunAnyInBackgroundAppOpsAllowed(uid, packageName));
+                        !mActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled()
+                        && !mAppStateTracker.isRunAnyInBackgroundAppOpsAllowed(uid, packageName));
         didChange |= jobStatus.setUidActive(isActive);
         return didChange;
     }
diff --git a/core/java/android/app/ActivityManagerInternal.java b/core/java/android/app/ActivityManagerInternal.java
index ae36646..70e95a1 100644
--- a/core/java/android/app/ActivityManagerInternal.java
+++ b/core/java/android/app/ActivityManagerInternal.java
@@ -724,4 +724,35 @@
      * Get the restriction level of the given package for given user id.
      */
     public abstract @RestrictionLevel int getRestrictionLevel(String pkg, @UserIdInt int userId);
+
+    /**
+     * Get whether or not apps would be put into restricted standby bucket automatically
+     * when it's background-restricted.
+     */
+    public abstract boolean isBgAutoRestrictedBucketFeatureFlagEnabled();
+
+    /**
+     * A listener interface, which will be notified on background restriction changes.
+     */
+    public interface AppBackgroundRestrictionListener {
+        /**
+         * Called when the background restriction level of given uid/package is changed.
+         */
+        default void onRestrictionLevelChanged(int uid, String packageName,
+                @RestrictionLevel int newLevel) {
+        }
+
+        /**
+         * Called when toggling the feature flag of moving to restricted standby bucket
+         * automatically on background-restricted.
+         */
+        default void onAutoRestrictedBucketFeatureFlagChanged(boolean autoRestrictedBucket) {
+        }
+    }
+
+    /**
+     * Register the background restriction listener callback.
+     */
+    public abstract void addAppBackgroundRestrictionListener(
+            @NonNull AppBackgroundRestrictionListener listener);
 }
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index a9263fe..57b5aab 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -16838,6 +16838,17 @@
         public @RestrictionLevel int getRestrictionLevel(String pkg, @UserIdInt int userId) {
             return mAppRestrictionController.getRestrictionLevel(pkg, userId);
         }
+
+        @Override
+        public boolean isBgAutoRestrictedBucketFeatureFlagEnabled() {
+            return mAppRestrictionController.isBgAutoRestrictedBucketFeatureFlagEnabled();
+        }
+
+        @Override
+        public void addAppBackgroundRestrictionListener(
+                @NonNull ActivityManagerInternal.AppBackgroundRestrictionListener listener) {
+            mAppRestrictionController.addAppBackgroundRestrictionListener(listener);
+        }
     }
 
     long inputDispatchingTimedOut(int pid, final boolean aboveSystem, String reason) {
diff --git a/services/core/java/com/android/server/am/AppRestrictionController.java b/services/core/java/com/android/server/am/AppRestrictionController.java
index 9d534e0..aa24a34 100644
--- a/services/core/java/com/android/server/am/AppRestrictionController.java
+++ b/services/core/java/com/android/server/am/AppRestrictionController.java
@@ -34,6 +34,7 @@
 import static android.app.usage.UsageStatsManager.REASON_SUB_FORCED_SYSTEM_FLAG_UNDEFINED;
 import static android.app.usage.UsageStatsManager.REASON_SUB_FORCED_USER_FLAG_INTERACTION;
 import static android.app.usage.UsageStatsManager.REASON_SUB_MASK;
+import static android.app.usage.UsageStatsManager.REASON_SUB_USAGE_SYSTEM_UPDATE;
 import static android.app.usage.UsageStatsManager.REASON_SUB_USAGE_USER_INTERACTION;
 import static android.app.usage.UsageStatsManager.STANDBY_BUCKET_ACTIVE;
 import static android.app.usage.UsageStatsManager.STANDBY_BUCKET_EXEMPTED;
@@ -55,6 +56,7 @@
 import android.annotation.UserIdInt;
 import android.app.ActivityManager;
 import android.app.ActivityManager.RestrictionLevel;
+import android.app.ActivityManagerInternal.AppBackgroundRestrictionListener;
 import android.app.ActivityThread;
 import android.app.AppOpsManager;
 import android.app.IActivityManager;
@@ -62,11 +64,13 @@
 import android.app.usage.AppStandbyInfo;
 import android.app.usage.UsageStatsManager;
 import android.content.BroadcastReceiver;
+import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
+import android.database.ContentObserver;
 import android.net.Uri;
 import android.os.Handler;
 import android.os.HandlerThread;
@@ -78,6 +82,7 @@
 import android.provider.DeviceConfig;
 import android.provider.DeviceConfig.OnPropertiesChangedListener;
 import android.provider.DeviceConfig.Properties;
+import android.provider.Settings.Global;
 import android.util.Slog;
 import android.util.SparseArrayMap;
 import android.util.TimeUtils;
@@ -99,6 +104,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArraySet;
+import java.util.function.Consumer;
 
 /**
  * This class tracks various state of the apps and mutates their restriction levels accordingly.
@@ -126,7 +132,7 @@
     @GuardedBy("mLock")
     private final RestrictionSettings mRestrictionSettings = new RestrictionSettings();
 
-    private final CopyOnWriteArraySet<AppRestrictionLevelListener> mRestrictionLevelListeners =
+    private final CopyOnWriteArraySet<AppBackgroundRestrictionListener> mRestrictionListeners =
             new CopyOnWriteArraySet<>();
 
     /**
@@ -312,6 +318,13 @@
             }
         }
 
+        @GuardedBy("mLock")
+        void forEachUidLocked(@NonNull Consumer<Integer> consumer) {
+            for (int i = mRestrictionLevels.numMaps() - 1; i >= 0; i--) {
+                consumer.accept(mRestrictionLevels.keyAt(i));
+            }
+        }
+
         void removeUser(@UserIdInt int userId) {
             synchronized (mLock) {
                 for (int i = mRestrictionLevels.numMaps() - 1; i >= 0; i--) {
@@ -353,18 +366,80 @@
         }
     }
 
-    private final OnPropertiesChangedListener mOnDeviceConfigChangedListener =
-            new OnPropertiesChangedListener() {
-                @Override
-                public void onPropertiesChanged(Properties properties) {
-                    for (String name : properties.getKeyset()) {
-                        if (name == null || !name.startsWith(DEVICE_CONFIG_SUBNAMESPACE_PREFIX)) {
-                            return;
-                        }
-                        AppRestrictionController.this.onPropertiesChanged(name);
-                    }
+    final class ConstantsObserver extends ContentObserver implements
+            OnPropertiesChangedListener {
+        /**
+         * Whether or not to set the app to restricted standby bucket automatically
+         * when it's background-restricted.
+         */
+        static final String KEY_BG_AUTO_RESTRICTED_BUCKET_ON_BG_RESTRICTION =
+                    DEVICE_CONFIG_SUBNAMESPACE_PREFIX + "auto_restricted_bucket_on_bg_restricted";
+
+        static final boolean DEFAULT_BG_AUTO_RESTRICTED_BUCKET_ON_BG_RESTRICTION = true;
+
+        volatile boolean mBgAutoRestrictedBucket;
+
+        volatile boolean mRestrictedBucketEnabled;
+
+        ConstantsObserver(Handler handler) {
+            super(handler);
+        }
+
+        @Override
+        public void onPropertiesChanged(Properties properties) {
+            for (String name : properties.getKeyset()) {
+                if (name == null || !name.startsWith(DEVICE_CONFIG_SUBNAMESPACE_PREFIX)) {
+                    return;
                 }
-            };
+                switch (name) {
+                    case KEY_BG_AUTO_RESTRICTED_BUCKET_ON_BG_RESTRICTION:
+                        updateBgAutoRestrictedBucketChanged();
+                        break;
+                }
+                AppRestrictionController.this.onPropertiesChanged(name);
+            }
+        }
+
+        @Override
+        public void onChange(boolean selfChange) {
+            updateSettings();
+        }
+
+        public void start() {
+            final ContentResolver cr = mContext.getContentResolver();
+            cr.registerContentObserver(Global.getUriFor(Global.ENABLE_RESTRICTED_BUCKET),
+                    false, this);
+            updateSettings();
+            updateDeviceConfig();
+        }
+
+        void updateSettings() {
+            mRestrictedBucketEnabled = isRestrictedBucketEnabled();
+        }
+
+        private boolean isRestrictedBucketEnabled() {
+            return Global.getInt(mContext.getContentResolver(),
+                    Global.ENABLE_RESTRICTED_BUCKET,
+                    Global.DEFAULT_ENABLE_RESTRICTED_BUCKET) == 1;
+        }
+
+        void updateDeviceConfig() {
+            updateBgAutoRestrictedBucketChanged();
+        }
+
+        private void updateBgAutoRestrictedBucketChanged() {
+            boolean oldValue = mBgAutoRestrictedBucket;
+            mBgAutoRestrictedBucket = DeviceConfig.getBoolean(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    KEY_BG_AUTO_RESTRICTED_BUCKET_ON_BG_RESTRICTION,
+                    DEFAULT_BG_AUTO_RESTRICTED_BUCKET_ON_BG_RESTRICTION);
+            if (oldValue != mBgAutoRestrictedBucket) {
+                dispatchAutoRestrictedBucketFeatureFlagChanged(mBgAutoRestrictedBucket);
+            }
+        }
+    }
+
+    private final ConstantsObserver mConstantsObserver;
 
     private final AppStateTracker.BackgroundRestrictedAppListener mBackgroundRestrictionListener =
             new AppStateTracker.BackgroundRestrictedAppListener() {
@@ -422,20 +497,11 @@
             };
 
     /**
-     * A listener interface, which will be notified on restriction level changes.
+     * Register the background restriction listener callback.
      */
-    public interface AppRestrictionLevelListener {
-        /**
-         * Called when the restriction level of given uid/package is changed.
-         */
-        void onRestrictionLevelChanged(int uid, String packageName, @RestrictionLevel int newLevel);
-    }
-
-    /**
-     * Register the restriction level listener callback.
-     */
-    public void addAppRestrictionLevelListener(@NonNull AppRestrictionLevelListener listener) {
-        mRestrictionLevelListeners.add(listener);
+    public void addAppBackgroundRestrictionListener(
+            @NonNull AppBackgroundRestrictionListener listener) {
+        mRestrictionListeners.add(listener);
     }
 
     AppRestrictionController(final Context context) {
@@ -448,13 +514,14 @@
         mBgHandlerThread = new HandlerThread("bgres-controller");
         mBgHandlerThread.start();
         mBgHandler = new BgHandler(mBgHandlerThread.getLooper(), injector);
+        mConstantsObserver = new ConstantsObserver(mBgHandler);
         injector.initAppStateTrackers(this);
     }
 
     void onSystemReady() {
         DeviceConfig.addOnPropertiesChangedListener(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
-                ActivityThread.currentApplication().getMainExecutor(),
-                mOnDeviceConfigChangedListener);
+                ActivityThread.currentApplication().getMainExecutor(), mConstantsObserver);
+        mConstantsObserver.start();
         initRestrictionStates();
         registerForUidObservers();
         registerForSystemBroadcasts();
@@ -559,7 +626,8 @@
                         .isAppBackgroundRestricted(uid, packageName)) {
                     return RESTRICTION_LEVEL_BACKGROUND_RESTRICTED;
                 }
-                level = standbyBucket == STANDBY_BUCKET_RESTRICTED
+                level = mConstantsObserver.mRestrictedBucketEnabled
+                        && standbyBucket == STANDBY_BUCKET_RESTRICTED
                         ? RESTRICTION_LEVEL_RESTRICTED_BUCKET
                         : RESTRICTION_LEVEL_ADAPTIVE_BUCKET;
                 if (calcTrackers) {
@@ -594,9 +662,13 @@
      */
     private @RestrictionLevel int calcAppRestrictionLevelFromTackers(int uid, String packageName) {
         @RestrictionLevel int level = RESTRICTION_LEVEL_UNKNOWN;
+        final boolean isRestrictedBucketEnabled = mConstantsObserver.mRestrictedBucketEnabled;
         for (int i = mAppStateTrackers.size() - 1; i >= 0; i--) {
             @RestrictionLevel int l = mAppStateTrackers.get(i).getPolicy()
                     .getProposedRestrictionLevel(packageName, uid);
+            if (!isRestrictedBucketEnabled && l == RESTRICTION_LEVEL_RESTRICTED_BUCKET) {
+                l = RESTRICTION_LEVEL_ADAPTIVE_BUCKET;
+            }
             level = Math.max(level, l);
         }
         return level;
@@ -660,6 +732,10 @@
         final AppStandbyInternal appStandbyInternal = mInjector.getAppStandbyInternal();
         if (level >= RESTRICTION_LEVEL_RESTRICTED_BUCKET
                 && curLevel < RESTRICTION_LEVEL_RESTRICTED_BUCKET) {
+            if (!mConstantsObserver.mRestrictedBucketEnabled
+                    || !mConstantsObserver.mBgAutoRestrictedBucket) {
+                return;
+            }
             // Moving the app standby bucket to restricted in the meanwhile.
             if (DEBUG_BG_RESTRICTION_CONTROLLER
                     && level == RESTRICTION_LEVEL_BACKGROUND_RESTRICTED) {
@@ -736,10 +812,36 @@
 
     private void dispatchAppRestrictionLevelChanges(int uid, String pkgName,
             @RestrictionLevel int newLevel) {
-        mRestrictionLevelListeners.forEach(
+        mRestrictionListeners.forEach(
                 l -> l.onRestrictionLevelChanged(uid, pkgName, newLevel));
     }
 
+    private void dispatchAutoRestrictedBucketFeatureFlagChanged(boolean newValue) {
+        final AppStandbyInternal appStandbyInternal = mInjector.getAppStandbyInternal();
+        final ArrayList<Runnable> pendingTasks = new ArrayList<>();
+        synchronized (mLock) {
+            mRestrictionSettings.forEachUidLocked(uid -> {
+                mRestrictionSettings.forEachPackageInUidLocked(uid, (pkgName, level, reason) -> {
+                    if (level == RESTRICTION_LEVEL_BACKGROUND_RESTRICTED) {
+                        pendingTasks.add(newValue
+                                ? () -> appStandbyInternal.restrictApp(pkgName,
+                                UserHandle.getUserId(uid), reason & REASON_MAIN_MASK,
+                                reason & REASON_SUB_MASK)
+                                : () -> appStandbyInternal.maybeUnrestrictApp(pkgName,
+                                UserHandle.getUserId(uid), reason & REASON_MAIN_MASK,
+                                reason & REASON_SUB_MASK, REASON_MAIN_USAGE,
+                                REASON_SUB_USAGE_SYSTEM_UPDATE));
+                    }
+                });
+            });
+        }
+        for (int i = 0; i < pendingTasks.size(); i++) {
+            pendingTasks.get(i).run();
+        }
+        mRestrictionListeners.forEach(
+                l -> l.onAutoRestrictedBucketFeatureFlagChanged(newValue));
+    }
+
     private void handleAppStandbyBucketChanged(int bucket, String packageName,
             @UserIdInt int userId) {
         final int uid = mInjector.getPackageManagerInternal().getPackageUid(
@@ -1064,6 +1166,10 @@
         mRestrictionSettings.removeUid(uid);
     }
 
+    boolean isBgAutoRestrictedBucketFeatureFlagEnabled() {
+        return mConstantsObserver.mBgAutoRestrictedBucket;
+    }
+
     private void onPropertiesChanged(String name) {
         for (int i = 0, size = mAppStateTrackers.size(); i < size; i++) {
             mAppStateTrackers.get(i).onPropertiesChanged(name);
diff --git a/services/tests/mockingservicestests/src/com/android/server/AppStateTrackerTest.java b/services/tests/mockingservicestests/src/com/android/server/AppStateTrackerTest.java
index 40b3664..9b04ae4 100644
--- a/services/tests/mockingservicestests/src/com/android/server/AppStateTrackerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/AppStateTrackerTest.java
@@ -344,6 +344,30 @@
         callStart(instance);
 
         assertFalse(instance.isForceAllAppsStandbyEnabled());
+
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(false);
+
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false},
+                false);
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false},
+                true);
+        areAlarmsRestrictedByBatterySaver(instance,
+                new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false});
+
+        // Toggle the auto restricted bucket feature flag on bg restriction, shouldn't make a
+        // difference.
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(true);
+
         areJobsRestricted(instance,
                 new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
                 new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
@@ -364,6 +388,9 @@
 
         assertTrue(instance.isForceAllAppsStandbyEnabled());
 
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(false);
+
         areJobsRestricted(instance,
                 new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
                 new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
@@ -379,6 +406,29 @@
                 new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
                 new boolean[] {true, true, true, false});
 
+        // Toggle the auto restricted bucket feature flag on bg restriction, shouldn't make a
+        // difference.
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(true);
+
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {true, true, true, false},
+                false);
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false},
+                true);
+        areAlarmsRestrictedByBatterySaver(instance,
+                new int[] {UID_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {true, true, true, false});
+
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(false);
+
         // Toggle the foreground state.
 
         assertFalse(instance.isUidActive(UID_1));
@@ -500,9 +550,35 @@
                 new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
                 new boolean[] {true, false, false, true, false});
 
+        // Toggle the auto restricted bucket feature flag on bg restriction.
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(true);
+
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false, false},
+                false);
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false, false},
+                true);
+
+        areAlarmsRestrictedByBatterySaver(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false, false});
+        areAlarmsRestrictedByFAS(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false, false});
+
         // Toggle power saver, should still be the same.
         mPowerSaveMode = true;
         mPowerSaveObserver.accept(getPowerSaveState());
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(false);
 
         areJobsRestricted(instance,
                 new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
@@ -524,9 +600,36 @@
                 new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
                 new boolean[] {true, false, false, true, false});
 
+        // Toggle the auto restricted bucket feature flag on bg restriction.
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(true);
+
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {true, true, true, true, false},
+                false);
+        areJobsRestricted(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false, false},
+                true);
+
+        areAlarmsRestrictedByBatterySaver(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {true, true, true, true, false});
+        areAlarmsRestrictedByFAS(instance,
+                new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
+                new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
+                new boolean[] {false, false, false, false, false});
+
         mPowerSaveMode = false;
         mPowerSaveObserver.accept(getPowerSaveState());
 
+        when(mMockIActivityManagerInternal.isBgAutoRestrictedBucketFeatureFlagEnabled())
+                .thenReturn(false);
+
         areJobsRestricted(instance,
                 new int[] {UID_1, UID_10_1, UID_2, UID_10_2, Process.SYSTEM_UID},
                 new String[] {PACKAGE_1, PACKAGE_1, PACKAGE_2, PACKAGE_2, PACKAGE_SYSTEM},
diff --git a/services/tests/mockingservicestests/src/com/android/server/am/BackgroundRestrictionTest.java b/services/tests/mockingservicestests/src/com/android/server/am/BackgroundRestrictionTest.java
index 4f9fea9..1d031e1 100644
--- a/services/tests/mockingservicestests/src/com/android/server/am/BackgroundRestrictionTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/am/BackgroundRestrictionTest.java
@@ -60,6 +60,7 @@
 import static org.mockito.Mockito.verify;
 
 import android.app.ActivityManagerInternal;
+import android.app.ActivityManagerInternal.AppBackgroundRestrictionListener;
 import android.app.AppOpsManager;
 import android.app.IActivityManager;
 import android.app.IUidObserver;
@@ -84,7 +85,6 @@
 import com.android.server.AppStateTracker;
 import com.android.server.DeviceIdleInternal;
 import com.android.server.am.AppBatteryTracker.AppBatteryPolicy;
-import com.android.server.am.AppRestrictionController.AppRestrictionLevelListener;
 import com.android.server.apphibernation.AppHibernationManagerInternal;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.usage.AppStandbyInternal;
@@ -272,7 +272,7 @@
         final TestAppRestrictionLevelListener listener = new TestAppRestrictionLevelListener();
         final long timeout = 1_000; // ms
 
-        mBgRestrictionController.addAppRestrictionLevelListener(listener);
+        mBgRestrictionController.addAppBackgroundRestrictionListener(listener);
 
         setBackgroundRestrict(testPkgName, testUid, false, listener);
 
@@ -368,7 +368,7 @@
         final TestAppRestrictionLevelListener listener = new TestAppRestrictionLevelListener();
         final long timeout = 1_000; // ms
 
-        mBgRestrictionController.addAppRestrictionLevelListener(listener);
+        mBgRestrictionController.addAppBackgroundRestrictionListener(listener);
 
         setBackgroundRestrict(testPkgName, testUid, false, listener);
 
@@ -435,7 +435,7 @@
         DeviceConfigSession<Float> bgCurrentDrainRestrictedBucketThreshold = null;
         DeviceConfigSession<Float> bgCurrentDrainBgRestrictedThreshold = null;
 
-        mBgRestrictionController.addAppRestrictionLevelListener(listener);
+        mBgRestrictionController.addAppBackgroundRestrictionListener(listener);
 
         setBackgroundRestrict(testPkgName, testUid, false, listener);
 
@@ -699,7 +699,7 @@
         waitForIdleHandler(mBgRestrictionController.getBackgroundHandler());
     }
 
-    private class TestAppRestrictionLevelListener implements AppRestrictionLevelListener {
+    private class TestAppRestrictionLevelListener implements AppBackgroundRestrictionListener {
         final CountDownLatch[] mLatchHolder = new CountDownLatch[1];
         final int[] mUidHolder = new int[1];
         final String[] mPkgNameHolder = new String[1];