Add a tracker on background current drains for each uid

On detecting excessive battery usage of a background uid, we may
move the background restriction of the apps running in this uid to
more restrictive levels.

Bug: 200326767
Test: atest FrameworksMockingServicesTests:BackgroundRestrictionTest
Change-Id: I6c2d41e44367a283d8aa9491be683018a80a810c
diff --git a/core/java/android/util/SparseDoubleArray.java b/core/java/android/util/SparseDoubleArray.java
index e8d96d8..cb51f7a 100644
--- a/core/java/android/util/SparseDoubleArray.java
+++ b/core/java/android/util/SparseDoubleArray.java
@@ -124,6 +124,15 @@
     }
 
     /**
+     * Returns the index for which {@link #keyAt} would return the
+     * specified key, or a negative number if the specified
+     * key is not mapped.
+     */
+    public int indexOfKey(int key) {
+        return mValues.indexOfKey(key);
+    }
+
+    /**
      * Given an index in the range <code>0...size()-1</code>, returns
      * the key from the <code>index</code>th key-value mapping that this
      * SparseDoubleArray stores.
@@ -146,6 +155,20 @@
     }
 
     /**
+     * Given an index in the range <code>0...size()-1</code>, sets a new
+     * value for the <code>index</code>th key-value mapping that this
+     * SparseDoubleArray stores.
+     *
+     * <p>For indices outside of the range <code>0...size()-1</code>, the behavior is undefined for
+     * apps targeting {@link android.os.Build.VERSION_CODES#P} and earlier, and an
+     * {@link ArrayIndexOutOfBoundsException} is thrown for apps targeting
+     * {@link android.os.Build.VERSION_CODES#Q} and later.</p>
+     */
+    public void setValueAt(int index, double value) {
+        mValues.setValueAt(index, Double.doubleToRawLongBits(value));
+    }
+
+    /**
      * Removes all key-value mappings from this SparseDoubleArray.
      */
     public void clear() {
diff --git a/core/java/android/util/SparseLongArray.java b/core/java/android/util/SparseLongArray.java
index 7185972..b739e37 100644
--- a/core/java/android/util/SparseLongArray.java
+++ b/core/java/android/util/SparseLongArray.java
@@ -245,6 +245,28 @@
     }
 
     /**
+     * Given an index in the range <code>0...size()-1</code>, sets a new
+     * value for the <code>index</code>th key-value mapping that this
+     * SparseLongArray stores.
+     *
+     * <p>For indices outside of the range <code>0...size()-1</code>, the behavior is undefined for
+     * apps targeting {@link android.os.Build.VERSION_CODES#P} and earlier, and an
+     * {@link ArrayIndexOutOfBoundsException} is thrown for apps targeting
+     * {@link android.os.Build.VERSION_CODES#Q} and later.</p>
+     *
+     * @hide
+     */
+    public void setValueAt(int index, long value) {
+        if (index >= mSize && UtilConfig.sThrowExceptionForUpperArrayOutOfBounds) {
+            // The array might be slightly bigger than mSize, in which case, indexing won't fail.
+            // Check if exception should be thrown outside of the critical path.
+            throw new ArrayIndexOutOfBoundsException(index);
+        }
+
+        mValues[index] = value;
+    }
+
+    /**
      * Returns the index for which {@link #keyAt} would return the
      * specified key, or a negative number if the specified
      * key is not mapped.
diff --git a/services/core/java/android/os/BatteryStatsInternal.java b/services/core/java/android/os/BatteryStatsInternal.java
index e996eb4..d49cc11 100644
--- a/services/core/java/android/os/BatteryStatsInternal.java
+++ b/services/core/java/android/os/BatteryStatsInternal.java
@@ -20,6 +20,7 @@
 import com.android.internal.os.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 
 import java.util.Collection;
+import java.util.List;
 
 /**
  * Battery stats local system service interface. This is used to pass internal data out of
@@ -42,6 +43,17 @@
     public abstract SystemServiceCpuThreadTimes getSystemServiceCpuThreadTimes();
 
     /**
+     * Returns BatteryUsageStats, which contains power attribution data on a per-subsystem
+     * and per-UID basis.
+     *
+     * <p>
+     * Note: This is a slow running method and should be called from non-blocking threads only.
+     * </p>
+     */
+    public abstract List<BatteryUsageStats> getBatteryUsageStats(
+            List<BatteryUsageStatsQuery> queries);
+
+    /**
      * Inform battery stats how many deferred jobs existed when the app got launched and how
      * long ago was the last job execution for the app.
      * @param uid the uid of the app.
diff --git a/services/core/java/com/android/server/am/AppBatteryTracker.java b/services/core/java/com/android/server/am/AppBatteryTracker.java
new file mode 100644
index 0000000..49a22d6
--- /dev/null
+++ b/services/core/java/com/android/server/am/AppBatteryTracker.java
@@ -0,0 +1,589 @@
+/*
+ * Copyright (C) 2022 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.am;
+
+import static android.app.ActivityManager.RESTRICTION_LEVEL_ADAPTIVE_BUCKET;
+import static android.app.ActivityManager.RESTRICTION_LEVEL_BACKGROUND_RESTRICTED;
+import static android.app.ActivityManager.RESTRICTION_LEVEL_RESTRICTED_BUCKET;
+import static android.app.ActivityManager.isLowRamDeviceStatic;
+import static android.app.usage.UsageStatsManager.REASON_MAIN_FORCED_BY_SYSTEM;
+import static android.app.usage.UsageStatsManager.REASON_MAIN_USAGE;
+import static android.app.usage.UsageStatsManager.REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE;
+import static android.app.usage.UsageStatsManager.REASON_SUB_USAGE_USER_INTERACTION;
+import static android.os.BatteryConsumer.POWER_COMPONENT_ANY;
+import static android.os.BatteryConsumer.PROCESS_STATE_BACKGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE;
+
+import static com.android.server.am.ActivityManagerDebugConfig.TAG_AM;
+import static com.android.server.am.ActivityManagerDebugConfig.TAG_WITH_CLASS_NAME;
+import static com.android.server.am.AppRestrictionController.DEVICE_CONFIG_SUBNAMESPACE_PREFIX;
+import static com.android.server.am.BaseAppStateTracker.ONE_DAY;
+import static com.android.server.am.BaseAppStateTracker.ONE_MINUTE;
+
+import android.annotation.NonNull;
+import android.annotation.UserIdInt;
+import android.app.ActivityManager.RestrictionLevel;
+import android.content.Context;
+import android.os.BatteryConsumer;
+import android.os.BatteryUsageStats;
+import android.os.BatteryUsageStatsQuery;
+import android.os.SystemClock;
+import android.os.UidBatteryConsumer;
+import android.os.UserHandle;
+import android.provider.DeviceConfig;
+import android.util.ArraySet;
+import android.util.Slog;
+import android.util.SparseArray;
+import android.util.SparseBooleanArray;
+import android.util.SparseDoubleArray;
+import android.util.TimeUtils;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.util.ArrayUtils;
+import com.android.server.am.AppBatteryTracker.AppBatteryPolicy;
+import com.android.server.am.BaseAppStateTracker.Injector;
+import com.android.server.pm.UserManagerInternal;
+
+import java.lang.reflect.Constructor;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * The battery usage tracker for apps, currently we are focusing on background + FGS battery here.
+ */
+final class AppBatteryTracker extends BaseAppStateTracker<AppBatteryPolicy> {
+    static final String TAG = TAG_WITH_CLASS_NAME ? "AppBatteryTracker" : TAG_AM;
+
+    private static final boolean DEBUG_BACKGROUND_BATTERY_TRACKER = false;
+
+    // As we don't support realtime per-UID battery usage stats yet, we're polling the stats
+    // in a regular time basis.
+    private final long mBatteryUsageStatsPollingIntervalMs;
+
+    // The timestamp when this system_server was started.
+    private long mBootTimestamp;
+
+    static final long BATTERY_USAGE_STATS_POLLING_INTERVAL_MS_LONG = 30 * ONE_MINUTE; // 30 mins
+    static final long BATTERY_USAGE_STATS_POLLING_INTERVAL_MS_DEBUG = 2_000L; // 2s
+
+    static final BatteryConsumer.Dimensions BATT_DIMEN_FG =
+            new BatteryConsumer.Dimensions(POWER_COMPONENT_ANY, PROCESS_STATE_FOREGROUND);
+    static final BatteryConsumer.Dimensions BATT_DIMEN_BG =
+            new BatteryConsumer.Dimensions(POWER_COMPONENT_ANY, PROCESS_STATE_BACKGROUND);
+    static final BatteryConsumer.Dimensions BATT_DIMEN_FGS =
+            new BatteryConsumer.Dimensions(POWER_COMPONENT_ANY, PROCESS_STATE_FOREGROUND_SERVICE);
+
+    private final Runnable mBgBatteryUsageStatsPolling = this::updateBatteryUsageStats;
+
+    /**
+     * This tracks the user ids which are or were active during the last polling window,
+     * the index is the user id, and the value is if it's still running or not by now.
+     */
+    @GuardedBy("mLock")
+    private final SparseBooleanArray mActiveUserIdStates = new SparseBooleanArray();
+
+    // No lock is needed as it's accessed in the handler thread only.
+    private final ArraySet<UserHandle> mTmpUserIds = new ArraySet<>();
+    private final SparseDoubleArray mTmpBatteryConsumptions = new SparseDoubleArray();
+
+    private BatteryUsageStatsQuery mBatteryUsageStatsQuery;
+
+    AppBatteryTracker(Context context, AppRestrictionController controller) {
+        this(context, controller, null, null);
+    }
+
+    AppBatteryTracker(Context context, AppRestrictionController controller,
+            Constructor<? extends Injector<AppBatteryPolicy>> injector,
+            Object outerContext) {
+        super(context, controller, injector, outerContext);
+        if (injector == null) {
+            mBatteryUsageStatsPollingIntervalMs = BATTERY_USAGE_STATS_POLLING_INTERVAL_MS_LONG;
+        } else {
+            mBatteryUsageStatsPollingIntervalMs = BATTERY_USAGE_STATS_POLLING_INTERVAL_MS_DEBUG;
+        }
+        mInjector.setPolicy(new AppBatteryPolicy(mInjector, this));
+    }
+
+    @Override
+    void onSystemReady() {
+        super.onSystemReady();
+        final UserManagerInternal um = mInjector.getUserManagerInternal();
+        final int[] userIds = um.getUserIds();
+        for (int userId : userIds) {
+            if (um.isUserRunning(userId)) {
+                synchronized (mLock) {
+                    mActiveUserIdStates.put(userId, true);
+                }
+            }
+        }
+        mBootTimestamp = mInjector.currentTimeMillis();
+        scheduleBatteryUsageStatsUpdateIfNecessary();
+    }
+
+    private void scheduleBatteryUsageStatsUpdateIfNecessary() {
+        if (mInjector.getPolicy().isEnabled()) {
+            synchronized (mLock) {
+                if (!mBgHandler.hasCallbacks(mBgBatteryUsageStatsPolling)) {
+                    mBgHandler.postDelayed(mBgBatteryUsageStatsPolling,
+                            mBatteryUsageStatsPollingIntervalMs);
+                }
+            }
+        }
+    }
+
+    @Override
+    void onUserStarted(final @UserIdInt int userId) {
+        synchronized (mLock) {
+            mActiveUserIdStates.put(userId, true);
+        }
+    }
+
+    @Override
+    void onUserStopped(final @UserIdInt int userId) {
+        synchronized (mLock) {
+            mActiveUserIdStates.put(userId, false);
+        }
+    }
+
+    @Override
+    void onUserInteractionStarted(String packageName, int uid) {
+        mInjector.getPolicy().onUserInteractionStarted(packageName, uid);
+    }
+
+    @Override
+    void onBackgroundRestrictionChanged(int uid, String pkgName, boolean restricted) {
+        mInjector.getPolicy().onBackgroundRestrictionChanged(uid, pkgName, restricted);
+    }
+
+    private void updateBatteryUsageStats() {
+        final AppBatteryPolicy bgPolicy = mInjector.getPolicy();
+        try {
+            updateBatteryUsageStatsOnce(mTmpBatteryConsumptions, mTmpUserIds);
+            final SparseDoubleArray uidConsumers = mTmpBatteryConsumptions;
+            for (int i = 0, size = uidConsumers.size(); i < size; i++) {
+                final int uid = uidConsumers.keyAt(i);
+                final double bgConsumption = uidConsumers.valueAt(i);
+                final double percentage = bgPolicy.getPercentage(bgConsumption);
+                if (DEBUG_BACKGROUND_BATTERY_TRACKER) {
+                    Slog.i(TAG, String.format("UID consumed %6.3f mAh (or %4.2f%%) in the past %s",
+                            bgConsumption, percentage,
+                            TimeUtils.formatDuration(bgPolicy.mBgCurrentDrainWindowMs)));
+                }
+                bgPolicy.handleUidBatteryConsumption(uid, percentage);
+            }
+        } finally {
+            scheduleBatteryUsageStatsUpdateIfNecessary();
+        }
+    }
+
+    /**
+     * Query the battery usage stats and fill the UID battery usage into the given buffer.
+     */
+    private void updateBatteryUsageStatsOnce(@NonNull SparseDoubleArray buf,
+            @NonNull ArraySet<UserHandle> userIdsHolder) {
+        final AppBatteryPolicy bgPolicy = mInjector.getPolicy();
+        final long now = mInjector.currentTimeMillis();
+        buf.clear();
+        userIdsHolder.clear();
+        synchronized (mLock) {
+            for (int i = mActiveUserIdStates.size() - 1; i >= 0; i--) {
+                userIdsHolder.add(UserHandle.of(mActiveUserIdStates.keyAt(i)));
+                if (!mActiveUserIdStates.valueAt(i)) {
+                    mActiveUserIdStates.removeAt(i);
+                }
+            }
+        }
+        if (DEBUG_BACKGROUND_BATTERY_TRACKER) {
+            Slog.i(TAG, "updateBatteryUsageStats");
+        }
+        BatteryUsageStatsQuery.Builder builder = new BatteryUsageStatsQuery.Builder()
+                .includeProcessStateData()
+                .setMaxStatsAgeMs(0);
+        final BatteryUsageStats stats = updateBatteryUsageStatsOnceInternal(
+                buf, builder, bgPolicy, userIdsHolder);
+        final long curStart = stats != null ? stats.getStatsStartTimestamp() : 0L;
+        final long curDuration = now - curStart;
+        final long windowSize = calcWindowSize(now, bgPolicy);
+        if (curDuration < windowSize) {
+            // No sufficient data, query snapshots instead.
+            builder = new BatteryUsageStatsQuery.Builder()
+                    .includeProcessStateData()
+                    .aggregateSnapshots(now - windowSize, curStart);
+            updateBatteryUsageStatsOnceInternal(buf, builder, bgPolicy, userIdsHolder);
+        } else if (curDuration > windowSize) {
+            final double scale = windowSize * 1.0d / curDuration;
+            for (int i = 0, size = buf.size(); i < size; i++) {
+                buf.setValueAt(i, buf.valueAt(i) * scale);
+            }
+        }
+    }
+
+    private long calcWindowSize(long now, AppBatteryPolicy bgPolicy) {
+        return Math.min(now - mBootTimestamp, bgPolicy.mBgCurrentDrainWindowMs);
+    }
+
+    private BatteryUsageStats updateBatteryUsageStatsOnceInternal(SparseDoubleArray buf,
+            BatteryUsageStatsQuery.Builder builder, AppBatteryPolicy bgPolicy,
+            ArraySet<UserHandle> userIds) {
+        for (int i = 0, size = userIds.size(); i < size; i++) {
+            builder.addUser(userIds.valueAt(i));
+        }
+        final List<BatteryUsageStats> statsList = mInjector.getBatteryStatsInternal()
+                .getBatteryUsageStats(Arrays.asList(builder.build()));
+        if (ArrayUtils.isEmpty(statsList)) {
+            // Shouldn't happen unless in test.
+            return null;
+        }
+        final BatteryUsageStats stats = statsList.get(0);
+        final List<UidBatteryConsumer> uidConsumers = stats.getUidBatteryConsumers();
+        for (UidBatteryConsumer uidConsumer : uidConsumers) {
+            // TODO: b/200326767 - as we are not supporting per proc state attribution yet,
+            // we couldn't distinguish between a real FGS vs. a bound FGS proc state.
+            final int uid = uidConsumer.getUid();
+            final double bgConsumption = bgPolicy.getBgConsumption(uidConsumer);
+            int index = buf.indexOfKey(uid);
+            if (index < 0) {
+                buf.put(uid, bgConsumption);
+            } else {
+                buf.setValueAt(index, buf.valueAt(index) + bgConsumption);
+            }
+        }
+        return stats;
+    }
+
+    private void onCurrentDrainMonitorEnabled(boolean enabled) {
+        if (enabled) {
+            if (!mBgHandler.hasCallbacks(mBgBatteryUsageStatsPolling)) {
+                mBgHandler.postDelayed(mBgBatteryUsageStatsPolling,
+                        mBatteryUsageStatsPollingIntervalMs);
+            }
+        } else {
+            mBgHandler.removeCallbacks(mBgBatteryUsageStatsPolling);
+        }
+    }
+
+    static final class AppBatteryPolicy extends BaseAppStatePolicy<AppBatteryTracker> {
+        /**
+         * Whether or not we should enable the monitoring on background current drains.
+         */
+        static final String KEY_BG_CURRENT_DRAIN_MONITOR_ENABLED =
+                DEVICE_CONFIG_SUBNAMESPACE_PREFIX + "current_drain_monitor_enabled";
+
+        /**
+         * The threshold of the background current drain (in percentage) to the restricted
+         * standby bucket. In conjunction with the {@link #KEY_BG_CURRENT_DRAIN_WINDOW},
+         * the app could be moved to more restricted standby bucket when its background current
+         * drain rate is over this limit.
+         */
+        static final String KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_RESTRICTED_BUCKET =
+                DEVICE_CONFIG_SUBNAMESPACE_PREFIX + "current_drain_threshold_to_restricted_bucket";
+
+        /**
+         * The threshold of the background current drain (in percentage) to the background
+         * restricted level. In conjunction with the {@link #KEY_BG_CURRENT_DRAIN_WINDOW},
+         * the app could be moved to more restricted level when its background current
+         * drain rate is over this limit.
+         */
+        static final String KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_BG_RESTRICTED =
+                DEVICE_CONFIG_SUBNAMESPACE_PREFIX + "current_drain_threshold_to_bg_restricted";
+
+        /**
+         * The background current drain window size. In conjunction with the
+         * {@link #KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_RESTRICTED_BUCKET}, the app could be moved to
+         * more restrictive bucket when its background current drain rate is over this limit.
+         */
+        static final String KEY_BG_CURRENT_DRAIN_WINDOW =
+                DEVICE_CONFIG_SUBNAMESPACE_PREFIX + "current_drain_window";
+
+        /**
+         * Default value to {@link #mBgCurrentDrainMonitorEnabled}.
+         */
+        static final boolean DEFAULT_BG_CURRENT_DRAIN_MONITOR_ENABLED = true;
+
+        /**
+         * Default value to {@link #mBgCurrentDrainRestrictedBucketThreshold}.
+         */
+        static final float DEFAULT_BG_CURRENT_DRAIN_RESTRICTED_BUCKET_THRESHOLD =
+                isLowRamDeviceStatic() ? 4.0f : 2.0f;
+
+        /**
+         * Default value to {@link #mBgCurrentDrainBgRestrictedThreshold}.
+         */
+        static final float DEFAULT_BG_CURRENT_DRAIN_BG_RESTRICTED_THRESHOLD =
+                isLowRamDeviceStatic() ? 8.0f : 4.0f;
+
+        /**
+         * Default value to {@link #mBgCurrentDrainWindowMs}.
+         */
+        static final long DEFAULT_BG_CURRENT_DRAIN_WINDOW_MS = ONE_DAY;
+
+        /**
+         * @see #KEY_BG_CURRENT_DRAIN_MONITOR_ENABLED.
+         */
+        volatile boolean mBgCurrentDrainMonitorEnabled = DEFAULT_BG_CURRENT_DRAIN_MONITOR_ENABLED;
+
+        /**
+         * @see #KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_RESTRICTED_BUCKET.
+         */
+        volatile float mBgCurrentDrainRestrictedBucketThreshold =
+                DEFAULT_BG_CURRENT_DRAIN_RESTRICTED_BUCKET_THRESHOLD;
+
+        /**
+         * @see #KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_BG_RESTRICTED.
+         */
+        volatile float mBgCurrentDrainBgRestrictedThreshold =
+                DEFAULT_BG_CURRENT_DRAIN_BG_RESTRICTED_THRESHOLD;
+
+        /**
+         * @see #KEY_BG_CURRENT_DRAIN_WINDOW.
+         */
+        volatile long mBgCurrentDrainWindowMs = DEFAULT_BG_CURRENT_DRAIN_WINDOW_MS;
+
+        /**
+         * The capacity of the battery when fully charged in mAh.
+         */
+        private int mBatteryFullChargeMah;
+
+        /**
+         * List of the packages with significant background battery usage, key is the UID of
+         * the package and value is an array of the timestamps when the UID is found guilty and
+         * should be moved to the next level of restriction.
+         */
+        @GuardedBy("mLock")
+        private final SparseArray<long[]> mHighBgBatteryPackages = new SparseArray<>();
+
+        @NonNull
+        private final Object mLock;
+
+        private static final int TIME_STAMP_INDEX_RESTRICTED_BUCKET = 0;
+        private static final int TIME_STAMP_INDEX_BG_RESTRICTED = 1;
+        private static final int TIME_STAMP_INDEX_LAST = 2;
+
+        AppBatteryPolicy(@NonNull Injector injector, @NonNull AppBatteryTracker tracker) {
+            super(injector, tracker);
+            mLock = tracker.mLock;
+        }
+
+        @Override
+        public void onPropertiesChanged(String name) {
+            switch (name) {
+                case KEY_BG_CURRENT_DRAIN_MONITOR_ENABLED:
+                    updateCurrentDrainMonitorEnabled();
+                    break;
+                case KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_RESTRICTED_BUCKET:
+                case KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_BG_RESTRICTED:
+                    updateCurrentDrainThreshold();
+                    break;
+                case KEY_BG_CURRENT_DRAIN_WINDOW:
+                    updateCurrentDrainWindow();
+                    break;
+            }
+        }
+
+        private void updateCurrentDrainMonitorEnabled() {
+            final boolean enabled = mBatteryFullChargeMah > 0
+                    && DeviceConfig.getBoolean(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    KEY_BG_CURRENT_DRAIN_MONITOR_ENABLED,
+                    DEFAULT_BG_CURRENT_DRAIN_MONITOR_ENABLED);
+            if (enabled != mBgCurrentDrainMonitorEnabled) {
+                mBgCurrentDrainMonitorEnabled = enabled;
+                mTracker.onCurrentDrainMonitorEnabled(enabled);
+            }
+        }
+
+        private void updateCurrentDrainThreshold() {
+            mBgCurrentDrainRestrictedBucketThreshold = DeviceConfig.getFloat(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_RESTRICTED_BUCKET,
+                    DEFAULT_BG_CURRENT_DRAIN_RESTRICTED_BUCKET_THRESHOLD);
+            mBgCurrentDrainBgRestrictedThreshold = DeviceConfig.getFloat(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_BG_RESTRICTED,
+                    DEFAULT_BG_CURRENT_DRAIN_BG_RESTRICTED_THRESHOLD);
+        }
+
+        private void updateCurrentDrainWindow() {
+            mBgCurrentDrainWindowMs = DeviceConfig.getLong(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    KEY_BG_CURRENT_DRAIN_WINDOW,
+                    DEFAULT_BG_CURRENT_DRAIN_WINDOW_MS);
+        }
+
+        @Override
+        public void onSystemReady() {
+            mBatteryFullChargeMah =
+                    mInjector.getBatteryManagerInternal().getBatteryFullCharge() / 1000;
+            updateCurrentDrainMonitorEnabled();
+            updateCurrentDrainThreshold();
+            updateCurrentDrainWindow();
+        }
+
+        @Override
+        public @RestrictionLevel int getProposedRestrictionLevel(String packageName, int uid) {
+            synchronized (mLock) {
+                final int index = mHighBgBatteryPackages.indexOfKey(uid);
+                if (index < 0) {
+                    // Not found, return adaptive as the default one.
+                    return RESTRICTION_LEVEL_ADAPTIVE_BUCKET;
+                }
+                final long[] ts = mHighBgBatteryPackages.valueAt(index);
+                return ts[TIME_STAMP_INDEX_BG_RESTRICTED] > 0
+                        ? RESTRICTION_LEVEL_BACKGROUND_RESTRICTED
+                        : RESTRICTION_LEVEL_RESTRICTED_BUCKET;
+            }
+        }
+
+        @Override
+        public boolean isEnabled() {
+            return mBgCurrentDrainMonitorEnabled;
+        }
+
+        double getBgConsumption(final UidBatteryConsumer uidConsumer) {
+            return getConsumedPowerNoThrow(uidConsumer, BATT_DIMEN_BG)
+                    + getConsumedPowerNoThrow(uidConsumer, BATT_DIMEN_FGS);
+        }
+
+        double getPercentage(final double consumption) {
+            return consumption / mBatteryFullChargeMah * 100;
+        }
+
+        void handleUidBatteryConsumption(final int uid, final double percentage) {
+            if (shouldExemptUid(uid)) {
+                return;
+            }
+            boolean notifyController = false;
+            boolean excessive = false;
+            synchronized (mLock) {
+                final int curLevel = mTracker.mAppRestrictionController.getRestrictionLevel(uid);
+                if (curLevel >= RESTRICTION_LEVEL_BACKGROUND_RESTRICTED) {
+                    // We're already in the background restricted level, nothing more we could do.
+                    return;
+                }
+                final int index = mHighBgBatteryPackages.indexOfKey(uid);
+                if (index < 0) {
+                    if (percentage >= mBgCurrentDrainRestrictedBucketThreshold) {
+                        // New findings to us, track it and let the controller know.
+                        final long[] ts = new long[TIME_STAMP_INDEX_LAST];
+                        ts[TIME_STAMP_INDEX_RESTRICTED_BUCKET] = SystemClock.elapsedRealtime();
+                        mHighBgBatteryPackages.put(uid, ts);
+                        notifyController = excessive = true;
+                    }
+                } else {
+                    final long[] ts = mHighBgBatteryPackages.valueAt(index);
+                    if (percentage < mBgCurrentDrainRestrictedBucketThreshold) {
+                        // it's actually back to normal, but we don't untrack it until
+                        // explicit user interactions.
+                        notifyController = true;
+                    } else if (percentage >= mBgCurrentDrainBgRestrictedThreshold) {
+                        // If we're in the restricted standby bucket but still seeing high
+                        // current drains, tell the controller again.
+                        if (curLevel == RESTRICTION_LEVEL_RESTRICTED_BUCKET
+                                && ts[TIME_STAMP_INDEX_BG_RESTRICTED] == 0) {
+                            final long now = SystemClock.elapsedRealtime();
+                            if (now > ts[TIME_STAMP_INDEX_RESTRICTED_BUCKET]
+                                    + mBgCurrentDrainWindowMs) {
+                                ts[TIME_STAMP_INDEX_BG_RESTRICTED] = now;
+                                notifyController = excessive = true;
+                            }
+                        }
+                    }
+                }
+            }
+
+            if (excessive) {
+                if (DEBUG_BACKGROUND_BATTERY_TRACKER) {
+                    Slog.i(TAG, "Excessive background current drain " + uid
+                            + String.format(" %.2f%%", percentage) + " over "
+                            + TimeUtils.formatDuration(mBgCurrentDrainWindowMs));
+                }
+                if (notifyController) {
+                    mTracker.mAppRestrictionController.refreshAppRestrictionLevelForUid(
+                            uid, REASON_MAIN_FORCED_BY_SYSTEM,
+                            REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE, true);
+                }
+            } else {
+                if (DEBUG_BACKGROUND_BATTERY_TRACKER) {
+                    Slog.i(TAG, "Background current drain backs to normal " + uid
+                            + String.format(" %.2f%%", percentage) + " over "
+                            + TimeUtils.formatDuration(mBgCurrentDrainWindowMs));
+                }
+                // For now, we're not lifting the restrictions if the bg current drain backs to
+                // normal util an explicit user interaction.
+            }
+        }
+
+        void onUserInteractionStarted(String packageName, int uid) {
+            boolean changed = false;
+            synchronized (mLock) {
+                final int curLevel = mTracker.mAppRestrictionController.getRestrictionLevel(
+                        uid, packageName);
+                if (curLevel == RESTRICTION_LEVEL_BACKGROUND_RESTRICTED) {
+                    // It's a sticky state, user interaction won't change it, still track it.
+                } else {
+                    // Remove the given UID from our tracking list, as user interacted with it.
+                    final int index = mHighBgBatteryPackages.indexOfKey(uid);
+                    if (index >= 0) {
+                        mHighBgBatteryPackages.removeAt(index);
+                        changed = true;
+                    }
+                }
+            }
+            if (changed) {
+                // Request to refresh the app restriction level.
+                mTracker.mAppRestrictionController.refreshAppRestrictionLevelForUid(uid,
+                        REASON_MAIN_USAGE, REASON_SUB_USAGE_USER_INTERACTION, true);
+            }
+        }
+
+        void onBackgroundRestrictionChanged(int uid, String pkgName, boolean restricted) {
+            if (restricted) {
+                return;
+            }
+            synchronized (mLock) {
+                // User has explicitly removed it from background restricted level,
+                // clear the timestamp of the background-restricted
+                final long[] ts = mHighBgBatteryPackages.get(uid);
+                if (ts != null) {
+                    ts[TIME_STAMP_INDEX_BG_RESTRICTED] = 0;
+                }
+            }
+        }
+
+        private double getConsumedPowerNoThrow(final UidBatteryConsumer uidConsumer,
+                final BatteryConsumer.Dimensions dimens) {
+            try {
+                return uidConsumer.getConsumedPower(dimens);
+            } catch (IllegalArgumentException e) {
+                return 0.0d;
+            }
+        }
+
+        /**
+         * Note: The {@link com.android.server.usage.AppStandbyController} has a complete exemption
+         * list and will place exempted packages in exempted bucket eventually.
+         */
+        @Override
+        public boolean shouldExemptUid(int uid) {
+            if (!super.shouldExemptUid(uid)) {
+                // TODO: b/200326767 - Exempt the location/music/pinned apps.
+            }
+            return false;
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/am/AppRestrictionController.java b/services/core/java/com/android/server/am/AppRestrictionController.java
index 85cccee..9d534e0 100644
--- a/services/core/java/com/android/server/am/AppRestrictionController.java
+++ b/services/core/java/com/android/server/am/AppRestrictionController.java
@@ -817,6 +817,11 @@
         return mLock;
     }
 
+    @VisibleForTesting
+    void addAppStateTracker(@NonNull BaseAppStateTracker tracker) {
+        mAppStateTrackers.add(tracker);
+    }
+
     static class BgHandler extends Handler {
         static final int MSG_BACKGROUND_RESTRICTION_CHANGED = 0;
         static final int MSG_APP_RESTRICTION_LEVEL_CHANGED = 1;
@@ -884,6 +889,7 @@
 
         void initAppStateTrackers(AppRestrictionController controller) {
             mAppRestrictionController = controller;
+            controller.mAppStateTrackers.add(new AppBatteryTracker(mContext, controller));
         }
 
         AppRestrictionController getAppRestrictionController() {
diff --git a/services/core/java/com/android/server/am/BaseAppStateTracker.java b/services/core/java/com/android/server/am/BaseAppStateTracker.java
index ee3521c..9166874 100644
--- a/services/core/java/com/android/server/am/BaseAppStateTracker.java
+++ b/services/core/java/com/android/server/am/BaseAppStateTracker.java
@@ -24,6 +24,8 @@
 import android.annotation.UserIdInt;
 import android.app.ActivityManagerInternal;
 import android.content.Context;
+import android.os.BatteryManagerInternal;
+import android.os.BatteryStatsInternal;
 import android.os.Handler;
 import android.util.Slog;
 
@@ -144,6 +146,8 @@
         T mAppStatePolicy;
 
         ActivityManagerInternal mActivityManagerInternal;
+        BatteryManagerInternal mBatteryManagerInternal;
+        BatteryStatsInternal mBatteryStatsInternal;
         DeviceIdleInternal mDeviceIdleInternal;
         UserManagerInternal mUserManagerInternal;
 
@@ -153,6 +157,8 @@
 
         void onSystemReady() {
             mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
+            mBatteryManagerInternal = LocalServices.getService(BatteryManagerInternal.class);
+            mBatteryStatsInternal = LocalServices.getService(BatteryStatsInternal.class);
             mDeviceIdleInternal = LocalServices.getService(DeviceIdleInternal.class);
             mUserManagerInternal = LocalServices.getService(UserManagerInternal.class);
 
@@ -163,6 +169,14 @@
             return mActivityManagerInternal;
         }
 
+        BatteryManagerInternal getBatteryManagerInternal() {
+            return mBatteryManagerInternal;
+        }
+
+        BatteryStatsInternal getBatteryStatsInternal() {
+            return mBatteryStatsInternal;
+        }
+
         T getPolicy() {
             return mAppStatePolicy;
         }
@@ -174,5 +188,12 @@
         UserManagerInternal getUserManagerInternal() {
             return mUserManagerInternal;
         }
+
+        /**
+         * Equivalent to {@link java.lang.System#currentTimeMillis}.
+         */
+        long currentTimeMillis() {
+            return System.currentTimeMillis();
+        }
     }
 }
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 0b92954..c8ad0e8 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -431,6 +431,11 @@
         }
 
         @Override
+        public List<BatteryUsageStats> getBatteryUsageStats(List<BatteryUsageStatsQuery> queries) {
+            return BatteryStatsService.this.getBatteryUsageStats(queries);
+        }
+
+        @Override
         public void noteJobsDeferred(int uid, int numDeferred, long sinceLast) {
             if (DBG) Slog.d(TAG, "Jobs deferred " + uid + ": " + numDeferred + " " + sinceLast);
             BatteryStatsService.this.noteJobsDeferred(uid, numDeferred, sinceLast);
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 0574a0b..4f9fea9 100644
--- a/services/tests/mockingservicestests/src/com/android/server/am/BackgroundRestrictionTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/am/BackgroundRestrictionTest.java
@@ -23,6 +23,7 @@
 import static android.app.usage.UsageStatsManager.REASON_MAIN_FORCED_BY_SYSTEM;
 import static android.app.usage.UsageStatsManager.REASON_MAIN_FORCED_BY_USER;
 import static android.app.usage.UsageStatsManager.REASON_MAIN_USAGE;
+import static android.app.usage.UsageStatsManager.REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE;
 import static android.app.usage.UsageStatsManager.REASON_SUB_FORCED_USER_FLAG_INTERACTION;
 import static android.app.usage.UsageStatsManager.REASON_SUB_USAGE_USER_INTERACTION;
 import static android.app.usage.UsageStatsManager.STANDBY_BUCKET_ACTIVE;
@@ -35,6 +36,9 @@
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
+import static com.android.server.am.AppBatteryTracker.BATT_DIMEN_BG;
+import static com.android.server.am.AppBatteryTracker.BATT_DIMEN_FG;
+import static com.android.server.am.AppBatteryTracker.BATT_DIMEN_FGS;
 import static com.android.server.am.AppRestrictionController.STOCK_PM_FLAGS;
 
 import static org.junit.Assert.assertEquals;
@@ -43,13 +47,16 @@
 import static org.mockito.Mockito.anyBoolean;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.anyObject;
 import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.app.ActivityManagerInternal;
@@ -60,17 +67,23 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
+import android.os.BatteryManagerInternal;
+import android.os.BatteryStatsInternal;
+import android.os.BatteryUsageStats;
 import android.os.Handler;
 import android.os.MessageQueue;
 import android.os.Process;
 import android.os.SystemClock;
+import android.os.UidBatteryConsumer;
 import android.os.UserHandle;
+import android.provider.DeviceConfig;
 import android.util.Log;
 
 import androidx.test.runner.AndroidJUnit4;
 
 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;
@@ -89,6 +102,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -140,17 +154,23 @@
         STANDBY_BUCKET_NEVER,
     };
 
+    private static final int BATTERY_FULL_CHARGE_MAH = 5_000;
+
     @Mock private ActivityManagerInternal mActivityManagerInternal;
     @Mock private AppOpsManager mAppOpsManager;
     @Mock private AppStandbyInternal mAppStandbyInternal;
     @Mock private AppHibernationManagerInternal mAppHibernationInternal;
     @Mock private AppStateTracker mAppStateTracker;
+    @Mock private BatteryManagerInternal mBatteryManagerInternal;
+    @Mock private BatteryStatsInternal mBatteryStatsInternal;
     @Mock private DeviceIdleInternal mDeviceIdleInternal;
     @Mock private IActivityManager mIActivityManager;
     @Mock private UserManagerInternal mUserManagerInternal;
     @Mock private PackageManager mPackageManager;
     @Mock private PackageManagerInternal mPackageManagerInternal;
 
+    private long mCurrentTimeMillis;
+
     @Captor private ArgumentCaptor<AppStateTracker.BackgroundRestrictedAppListener> mFasListenerCap;
     private AppStateTracker.BackgroundRestrictedAppListener mFasListener;
 
@@ -163,6 +183,7 @@
     private Context mContext = getInstrumentation().getTargetContext();
     private TestBgRestrictionInjector mInjector;
     private AppRestrictionController mBgRestrictionController;
+    private AppBatteryTracker mAppBatteryTracker;
 
     @Before
     public void setUp() throws Exception {
@@ -197,6 +218,9 @@
             doReturn(appStandbyInfoList).when(mAppStandbyInternal).getAppStandbyBuckets(userId);
         }
 
+        doReturn(BATTERY_FULL_CHARGE_MAH * 1000).when(mBatteryManagerInternal)
+                .getBatteryFullCharge();
+
         mBgRestrictionController.onSystemReady();
 
         verify(mInjector.getAppStateTracker())
@@ -384,6 +408,288 @@
         listener.verify(timeout, testUid, testPkgName, RESTRICTION_LEVEL_EXEMPTED);
     }
 
+    @Test
+    public void testBgCurrentDrainMonitor() throws Exception {
+        final BatteryUsageStats stats = mock(BatteryUsageStats.class);
+        final List<BatteryUsageStats> statsList = Arrays.asList(stats);
+        final int testPkgIndex = 2;
+        final String testPkgName = TEST_PACKAGE_BASE + testPkgIndex;
+        final int testUser = TEST_USER0;
+        final int testUid = UserHandle.getUid(testUser,
+                TEST_PACKAGE_APPID_BASE + testPkgIndex);
+        final int testUid2 = UserHandle.getUid(testUser,
+                TEST_PACKAGE_APPID_BASE + testPkgIndex + 1);
+        final TestAppRestrictionLevelListener listener = new TestAppRestrictionLevelListener();
+        final long timeout =
+                AppBatteryTracker.BATTERY_USAGE_STATS_POLLING_INTERVAL_MS_DEBUG * 2;
+        final long windowMs = 2_000;
+        final float restrictBucketThreshold = 2.0f;
+        final float restrictBucketThresholdMah =
+                BATTERY_FULL_CHARGE_MAH * restrictBucketThreshold / 100.0f;
+        final float bgRestrictedThreshold = 4.0f;
+        final float bgRestrictedThresholdMah =
+                BATTERY_FULL_CHARGE_MAH * bgRestrictedThreshold / 100.0f;
+
+        DeviceConfigSession<Boolean> bgCurrentDrainMonitor = null;
+        DeviceConfigSession<Long> bgCurrentDrainWindow = null;
+        DeviceConfigSession<Float> bgCurrentDrainRestrictedBucketThreshold = null;
+        DeviceConfigSession<Float> bgCurrentDrainBgRestrictedThreshold = null;
+
+        mBgRestrictionController.addAppRestrictionLevelListener(listener);
+
+        setBackgroundRestrict(testPkgName, testUid, false, listener);
+
+        // Verify the current settings.
+        verifyRestrictionLevel(RESTRICTION_LEVEL_ADAPTIVE_BUCKET, testPkgName, testUid);
+
+        final double[] zeros = new double[]{0.0f, 0.0f};
+        final int[] uids = new int[]{testUid, testUid2};
+
+        try {
+            bgCurrentDrainMonitor = new DeviceConfigSession<>(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    AppBatteryPolicy.KEY_BG_CURRENT_DRAIN_MONITOR_ENABLED,
+                    DeviceConfig::getBoolean,
+                    AppBatteryPolicy.DEFAULT_BG_CURRENT_DRAIN_MONITOR_ENABLED);
+            bgCurrentDrainMonitor.set(true);
+
+            bgCurrentDrainWindow = new DeviceConfigSession<>(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    AppBatteryPolicy.KEY_BG_CURRENT_DRAIN_WINDOW,
+                    DeviceConfig::getLong,
+                    AppBatteryPolicy.DEFAULT_BG_CURRENT_DRAIN_WINDOW_MS);
+            bgCurrentDrainWindow.set(windowMs);
+
+            bgCurrentDrainRestrictedBucketThreshold = new DeviceConfigSession<>(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    AppBatteryPolicy.KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_RESTRICTED_BUCKET,
+                    DeviceConfig::getFloat,
+                    AppBatteryPolicy.DEFAULT_BG_CURRENT_DRAIN_BG_RESTRICTED_THRESHOLD);
+            bgCurrentDrainRestrictedBucketThreshold.set(restrictBucketThreshold);
+
+            bgCurrentDrainBgRestrictedThreshold = new DeviceConfigSession<>(
+                    DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                    AppBatteryPolicy.KEY_BG_CURRENT_DRAIN_THRESHOLD_TO_BG_RESTRICTED,
+                    DeviceConfig::getFloat,
+                    AppBatteryPolicy.DEFAULT_BG_CURRENT_DRAIN_BG_RESTRICTED_THRESHOLD);
+            bgCurrentDrainBgRestrictedThreshold.set(bgRestrictedThreshold);
+
+            mCurrentTimeMillis = 10_000L;
+            doReturn(mCurrentTimeMillis - windowMs).when(stats).getStatsStartTimestamp();
+            doReturn(statsList).when(mBatteryStatsInternal).getBatteryUsageStats(anyObject());
+
+            runTestBgCurrentDrainMonitorOnce(listener, stats, uids,
+                    new double[]{restrictBucketThresholdMah - 1, 0},
+                    new double[]{0, restrictBucketThresholdMah - 1}, zeros,
+                    () -> {
+                        doReturn(mCurrentTimeMillis).when(stats).getStatsStartTimestamp();
+                        mCurrentTimeMillis += windowMs + 1;
+                        try {
+                            listener.verify(timeout, testUid, testPkgName,
+                                    RESTRICTION_LEVEL_ADAPTIVE_BUCKET);
+                            fail("There shouldn't be any level change events");
+                        } catch (Exception e) {
+                            // Expected.
+                        }
+                    });
+
+            runTestBgCurrentDrainMonitorOnce(listener, stats, uids,
+                    new double[]{restrictBucketThresholdMah + 1, 0},
+                    new double[]{0, restrictBucketThresholdMah - 1}, zeros,
+                    () -> {
+                        doReturn(mCurrentTimeMillis).when(stats).getStatsStartTimestamp();
+                        mCurrentTimeMillis += windowMs + 1;
+                        // It should have gone to the restricted bucket.
+                        listener.verify(timeout, testUid, testPkgName,
+                                RESTRICTION_LEVEL_RESTRICTED_BUCKET);
+                        verify(mInjector.getAppStandbyInternal()).restrictApp(
+                                eq(testPkgName),
+                                eq(testUser),
+                                anyInt(), anyInt());
+                    });
+
+
+            runTestBgCurrentDrainMonitorOnce(listener, stats, uids,
+                    new double[]{restrictBucketThresholdMah - 1, 0},
+                    new double[]{0, restrictBucketThresholdMah - 1}, zeros,
+                    () -> {
+                        doReturn(mCurrentTimeMillis).when(stats).getStatsStartTimestamp();
+                        mCurrentTimeMillis += windowMs + 1;
+                        // We won't change restriction level until user interactions.
+                        try {
+                            listener.verify(timeout, testUid, testPkgName,
+                                    RESTRICTION_LEVEL_ADAPTIVE_BUCKET);
+                            fail("There shouldn't be any level change events");
+                        } catch (Exception e) {
+                            // Expected.
+                        }
+                        verify(mInjector.getAppStandbyInternal(), never()).setAppStandbyBucket(
+                                eq(testPkgName),
+                                eq(STANDBY_BUCKET_RARE),
+                                eq(testUser),
+                                anyInt(), anyInt());
+                    });
+
+            // Trigger user interaction.
+            runTestBgCurrentDrainMonitorOnce(listener, stats, uids,
+                    new double[]{restrictBucketThresholdMah - 1, 0},
+                    new double[]{0, restrictBucketThresholdMah - 1}, zeros,
+                    () -> {
+                        doReturn(mCurrentTimeMillis).when(stats).getStatsStartTimestamp();
+                        mCurrentTimeMillis += windowMs + 1;
+                        mIdleStateListener.onUserInteractionStarted(testPkgName, testUser);
+                        waitForIdleHandler(mBgRestrictionController.getBackgroundHandler());
+                        // It should have been back to normal.
+                        listener.verify(timeout, testUid, testPkgName,
+                                RESTRICTION_LEVEL_ADAPTIVE_BUCKET);
+                        verify(mInjector.getAppStandbyInternal(), atLeast(1)).maybeUnrestrictApp(
+                                eq(testPkgName),
+                                eq(testUser),
+                                eq(REASON_MAIN_FORCED_BY_SYSTEM),
+                                eq(REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE),
+                                eq(REASON_MAIN_USAGE),
+                                eq(REASON_SUB_USAGE_USER_INTERACTION));
+                    });
+
+            clearInvocations(mInjector.getAppStandbyInternal());
+
+            runTestBgCurrentDrainMonitorOnce(listener, stats, uids,
+                    new double[]{restrictBucketThresholdMah + 1, 0},
+                    new double[]{0, restrictBucketThresholdMah - 1}, zeros,
+                    () -> {
+                        doReturn(mCurrentTimeMillis).when(stats).getStatsStartTimestamp();
+                        mCurrentTimeMillis += windowMs + 1;
+                        // It should have gone to the restricted bucket.
+                        listener.verify(timeout, testUid, testPkgName,
+                                RESTRICTION_LEVEL_RESTRICTED_BUCKET);
+                        verify(mInjector.getAppStandbyInternal(), times(1)).restrictApp(
+                                eq(testPkgName),
+                                eq(testUser),
+                                anyInt(), anyInt());
+                    });
+
+            clearInvocations(mInjector.getAppStandbyInternal());
+            // Drain a bit more, there shouldn't be any level changes.
+            runTestBgCurrentDrainMonitorOnce(listener, stats, uids,
+                    new double[]{restrictBucketThresholdMah + 2, 0},
+                    new double[]{0, restrictBucketThresholdMah - 1}, zeros,
+                    () -> {
+                        doReturn(mCurrentTimeMillis).when(stats).getStatsStartTimestamp();
+                        mCurrentTimeMillis += windowMs + 1;
+                        // We won't change restriction level until user interactions.
+                        try {
+                            listener.verify(timeout, testUid, testPkgName,
+                                    RESTRICTION_LEVEL_ADAPTIVE_BUCKET);
+                            fail("There shouldn't be any level change events");
+                        } catch (Exception e) {
+                            // Expected.
+                        }
+                        verify(mInjector.getAppStandbyInternal(), never()).setAppStandbyBucket(
+                                eq(testPkgName),
+                                eq(STANDBY_BUCKET_RARE),
+                                eq(testUser),
+                                anyInt(), anyInt());
+                    });
+
+            // Sleep a while and set a higher drain
+            Thread.sleep(windowMs);
+            clearInvocations(mInjector.getAppStandbyInternal());
+            clearInvocations(mBgRestrictionController);
+            runTestBgCurrentDrainMonitorOnce(listener, stats, uids,
+                    new double[]{bgRestrictedThresholdMah + 1, 0},
+                    new double[]{0, restrictBucketThresholdMah - 1}, zeros,
+                    () -> {
+                        doReturn(mCurrentTimeMillis).when(stats).getStatsStartTimestamp();
+                        mCurrentTimeMillis += windowMs + 1;
+                        // We won't change restriction level automatically because it needs
+                        // user consent.
+                        try {
+                            listener.verify(timeout, testUid, testPkgName,
+                                    RESTRICTION_LEVEL_BACKGROUND_RESTRICTED);
+                            fail("There shouldn't be level change event like this");
+                        } catch (Exception e) {
+                            // Expected.
+                        }
+                        verify(mInjector.getAppStandbyInternal(), never()).setAppStandbyBucket(
+                                eq(testPkgName),
+                                eq(STANDBY_BUCKET_RARE),
+                                eq(testUser),
+                                anyInt(), anyInt());
+                        // We should have requested to goto background restricted level.
+                        verify(mBgRestrictionController, times(1)).handleRequestBgRestricted(
+                                eq(testPkgName),
+                                eq(testUid));
+                    });
+
+            // Turn ON the FAS for real.
+            setBackgroundRestrict(testPkgName, testUid, true, listener);
+
+            // Verify it's background restricted now.
+            verifyRestrictionLevel(RESTRICTION_LEVEL_BACKGROUND_RESTRICTED, testPkgName, testUid);
+            listener.verify(timeout, testUid, testPkgName, RESTRICTION_LEVEL_BACKGROUND_RESTRICTED);
+
+            // Trigger user interaction.
+            mIdleStateListener.onUserInteractionStarted(testPkgName, testUser);
+            waitForIdleHandler(mBgRestrictionController.getBackgroundHandler());
+
+            listener.mLatchHolder[0] = new CountDownLatch(1);
+            try {
+                listener.verify(timeout, testUid, testPkgName,
+                        RESTRICTION_LEVEL_ADAPTIVE_BUCKET);
+                fail("There shouldn't be level change event like this");
+            } catch (Exception e) {
+                // Expected.
+            }
+
+            // Turn OFF the FAS.
+            listener.mLatchHolder[0] = new CountDownLatch(1);
+            clearInvocations(mInjector.getAppStandbyInternal());
+            clearInvocations(mBgRestrictionController);
+            setBackgroundRestrict(testPkgName, testUid, false, listener);
+
+            // It'll go back to restricted bucket because it used to behave poorly.
+            listener.verify(timeout, testUid, testPkgName, RESTRICTION_LEVEL_RESTRICTED_BUCKET);
+            verifyRestrictionLevel(RESTRICTION_LEVEL_RESTRICTED_BUCKET, testPkgName, testUid);
+        } finally {
+            closeIfNotNull(bgCurrentDrainMonitor);
+            closeIfNotNull(bgCurrentDrainWindow);
+            closeIfNotNull(bgCurrentDrainRestrictedBucketThreshold);
+            closeIfNotNull(bgCurrentDrainBgRestrictedThreshold);
+        }
+    }
+
+    private void closeIfNotNull(DeviceConfigSession<?> config) throws Exception {
+        if (config != null) {
+            config.close();
+        }
+    }
+
+    private interface RunnableWithException {
+        void run() throws Exception;
+    }
+
+    private void runTestBgCurrentDrainMonitorOnce(TestAppRestrictionLevelListener listener,
+            BatteryUsageStats stats, int[] uids, double[] bg, double[] fgs, double[] fg,
+            RunnableWithException runnable) throws Exception {
+        listener.mLatchHolder[0] = new CountDownLatch(1);
+        ArrayList<UidBatteryConsumer> consumers = new ArrayList<>();
+        for (int i = 0; i < uids.length; i++) {
+            consumers.add(mockUidBatteryConsumer(uids[i], bg[i], fgs[i], fg[i]));
+
+        }
+        doReturn(consumers).when(stats).getUidBatteryConsumers();
+        runnable.run();
+    }
+
+    private UidBatteryConsumer mockUidBatteryConsumer(int uid, double bg, double fgs, double fg) {
+        UidBatteryConsumer uidConsumer = mock(UidBatteryConsumer.class);
+        doReturn(uid).when(uidConsumer).getUid();
+        doReturn(bg).when(uidConsumer).getConsumedPower(eq(BATT_DIMEN_BG));
+        doReturn(fgs).when(uidConsumer).getConsumedPower(eq(BATT_DIMEN_FGS));
+        doReturn(fg).when(uidConsumer).getConsumedPower(eq(BATT_DIMEN_FG));
+        return uidConsumer;
+    }
+
     private void setBackgroundRestrict(String pkgName, int uid, boolean restricted,
             TestAppRestrictionLevelListener listener) throws Exception {
         Log.i(TAG, "Setting background restrict to " + restricted + " for " + pkgName + " " + uid);
@@ -451,6 +757,15 @@
 
         @Override
         void initAppStateTrackers(AppRestrictionController controller) {
+            try {
+                mAppBatteryTracker = new AppBatteryTracker(mContext, controller,
+                        TestAppBatteryTrackerInjector.class.getDeclaredConstructor(
+                                BackgroundRestrictionTest.class),
+                        BackgroundRestrictionTest.this);
+                controller.addAppStateTracker(mAppBatteryTracker);
+            } catch (NoSuchMethodException e) {
+                // Won't happen.
+            }
         }
 
         @Override
@@ -512,6 +827,16 @@
         }
 
         @Override
+        BatteryManagerInternal getBatteryManagerInternal() {
+            return BackgroundRestrictionTest.this.mBatteryManagerInternal;
+        }
+
+        @Override
+        BatteryStatsInternal getBatteryStatsInternal() {
+            return BackgroundRestrictionTest.this.mBatteryStatsInternal;
+        }
+
+        @Override
         DeviceIdleInternal getDeviceIdleInternal() {
             return BackgroundRestrictionTest.this.mDeviceIdleInternal;
         }
@@ -520,5 +845,13 @@
         UserManagerInternal getUserManagerInternal() {
             return BackgroundRestrictionTest.this.mUserManagerInternal;
         }
+
+        @Override
+        long currentTimeMillis() {
+            return BackgroundRestrictionTest.this.mCurrentTimeMillis;
+        }
+    }
+
+    private class TestAppBatteryTrackerInjector extends TestBaseTrackerInjector<AppBatteryPolicy> {
     }
 }