Merge "MeasuredEnergyStats - remove unused accumulate parameter" into sc-dev
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
index fc2a409..040a116 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
@@ -887,8 +887,10 @@
     }
 
     @Override
-    public void onUserStarting(@NonNull TargetUser user) {
+    public void onUserUnlocked(@NonNull TargetUser user) {
         synchronized (mLock) {
+            // Note that the user has started after its unlocked instead of when the user
+            // actually starts because the storage won't be decrypted until unlock.
             mStartedUsers = ArrayUtils.appendInt(mStartedUsers, user.getUserIdentifier());
         }
         // Let's kick any outstanding jobs for this user.
@@ -896,12 +898,6 @@
     }
 
     @Override
-    public void onUserUnlocking(@NonNull TargetUser user) {
-        // Let's kick any outstanding jobs for this user.
-        mHandler.obtainMessage(MSG_CHECK_JOB).sendToTarget();
-    }
-
-    @Override
     public void onUserStopping(@NonNull TargetUser user) {
         synchronized (mLock) {
             mStartedUsers = ArrayUtils.removeInt(mStartedUsers, user.getUserIdentifier());
diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index c6efaac..2b6f336 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -47,6 +47,7 @@
             POWER_COMPONENT_SYSTEM_SERVICES,
             POWER_COMPONENT_SENSORS,
             POWER_COMPONENT_GNSS,
+            POWER_COMPONENT_WIFI,
             POWER_COMPONENT_WAKELOCK,
             POWER_COMPONENT_SCREEN,
             POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS,
@@ -66,6 +67,7 @@
     public static final int POWER_COMPONENT_MOBILE_RADIO = 8;
     public static final int POWER_COMPONENT_SENSORS = 9;
     public static final int POWER_COMPONENT_GNSS = 10;
+    public static final int POWER_COMPONENT_WIFI = 11;
     public static final int POWER_COMPONENT_WAKELOCK = 12;
     public static final int POWER_COMPONENT_SCREEN = 13;
     // Power that is re-attributed to other battery consumers. For example, for System Server
@@ -94,6 +96,7 @@
             TIME_COMPONENT_MOBILE_RADIO,
             TIME_COMPONENT_SENSORS,
             TIME_COMPONENT_GNSS,
+            TIME_COMPONENT_WIFI,
             TIME_COMPONENT_WAKELOCK,
             TIME_COMPONENT_SCREEN,
     })
@@ -112,6 +115,7 @@
     public static final int TIME_COMPONENT_MOBILE_RADIO = 8;
     public static final int TIME_COMPONENT_SENSORS = 9;
     public static final int TIME_COMPONENT_GNSS = 10;
+    public static final int TIME_COMPONENT_WIFI = 11;
     public static final int TIME_COMPONENT_WAKELOCK = 12;
     public static final int TIME_COMPONENT_SCREEN = 13;
 
diff --git a/core/java/com/android/internal/os/UsageBasedPowerEstimator.java b/core/java/com/android/internal/os/UsageBasedPowerEstimator.java
index 5910b61..fd52014 100644
--- a/core/java/com/android/internal/os/UsageBasedPowerEstimator.java
+++ b/core/java/com/android/internal/os/UsageBasedPowerEstimator.java
@@ -32,6 +32,10 @@
         mAveragePowerMahPerMs = averagePowerMilliAmp / MILLIS_IN_HOUR;
     }
 
+    public boolean isSupported() {
+        return mAveragePowerMahPerMs != 0;
+    }
+
     /**
      * Given a {@link BatteryStats.Timer}, returns the accumulated duration.
      */
diff --git a/core/java/com/android/internal/os/WifiPowerCalculator.java b/core/java/com/android/internal/os/WifiPowerCalculator.java
index 63763f7..98f613f 100644
--- a/core/java/com/android/internal/os/WifiPowerCalculator.java
+++ b/core/java/com/android/internal/os/WifiPowerCalculator.java
@@ -15,8 +15,13 @@
  */
 package com.android.internal.os;
 
+import android.os.BatteryConsumer;
 import android.os.BatteryStats;
+import android.os.BatteryUsageStats;
+import android.os.BatteryUsageStatsQuery;
 import android.os.Process;
+import android.os.SystemBatteryConsumer;
+import android.os.UidBatteryConsumer;
 import android.os.UserHandle;
 import android.util.Log;
 import android.util.SparseArray;
@@ -30,25 +35,93 @@
 public class WifiPowerCalculator extends PowerCalculator {
     private static final boolean DEBUG = BatteryStatsHelper.DEBUG;
     private static final String TAG = "WifiPowerCalculator";
-    private final double mIdleCurrentMa;
-    private final double mTxCurrentMa;
-    private final double mRxCurrentMa;
-    private final PowerProfile mPowerProfile;
+    private final UsageBasedPowerEstimator mIdlePowerEstimator;
+    private final UsageBasedPowerEstimator mTxPowerEstimator;
+    private final UsageBasedPowerEstimator mRxPowerEstimator;
+    private final UsageBasedPowerEstimator mPowerOnPowerEstimator;
+    private final UsageBasedPowerEstimator mScanPowerEstimator;
+    private final UsageBasedPowerEstimator mBatchScanPowerEstimator;
     private final boolean mHasWifiPowerController;
-    private double mTotalAppPowerDrain = 0;
-    private long mTotalAppRunningTime = 0;
-    private WifiPowerEstimator mWifiPowerEstimator;
-    private boolean mHasWifiPowerReporting;
+    private final double mWifiPowerPerPacket;
+
+    private static class PowerDurationAndTraffic {
+        public double powerMah;
+        public long durationMs;
+
+        public long wifiRxPackets;
+        public long wifiTxPackets;
+        public long wifiRxBytes;
+        public long wifiTxBytes;
+    }
 
     public WifiPowerCalculator(PowerProfile profile) {
-        mPowerProfile = profile;
+        mPowerOnPowerEstimator = new UsageBasedPowerEstimator(
+                profile.getAveragePower(PowerProfile.POWER_WIFI_ON));
+        mScanPowerEstimator = new UsageBasedPowerEstimator(
+                profile.getAveragePower(PowerProfile.POWER_WIFI_SCAN));
+        mBatchScanPowerEstimator = new UsageBasedPowerEstimator(
+                profile.getAveragePower(PowerProfile.POWER_WIFI_BATCHED_SCAN));
+        mIdlePowerEstimator = new UsageBasedPowerEstimator(
+                profile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_IDLE));
+        mTxPowerEstimator = new UsageBasedPowerEstimator(
+                profile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_TX));
+        mRxPowerEstimator = new UsageBasedPowerEstimator(
+                profile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX));
 
-        mIdleCurrentMa = profile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_IDLE);
-        mTxCurrentMa = profile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_TX);
-        mRxCurrentMa = profile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX);
+        mWifiPowerPerPacket = getWifiPowerPerPacket(profile);
 
-        mHasWifiPowerController = mIdleCurrentMa != 0 && mTxCurrentMa != 0 && mRxCurrentMa != 0;
-        mWifiPowerEstimator = new WifiPowerEstimator(mPowerProfile);
+        mHasWifiPowerController =
+                mIdlePowerEstimator.isSupported() && mTxPowerEstimator.isSupported()
+                        && mRxPowerEstimator.isSupported();
+    }
+
+    @Override
+    public void calculate(BatteryUsageStats.Builder builder, BatteryStats batteryStats,
+            long rawRealtimeUs, long rawUptimeUs, BatteryUsageStatsQuery query) {
+
+        // batteryStats.hasWifiActivityReporting can change if we get energy data at a later point,
+        // so always check this field.
+        final boolean hasWifiPowerReporting =
+                mHasWifiPowerController && batteryStats.hasWifiActivityReporting();
+
+        final SystemBatteryConsumer.Builder systemBatteryConsumerBuilder =
+                builder.getOrCreateSystemBatteryConsumerBuilder(
+                        SystemBatteryConsumer.DRAIN_TYPE_WIFI);
+
+        long totalAppDurationMs = 0;
+        double totalAppPowerMah = 0;
+        final PowerDurationAndTraffic powerDurationAndTraffic = new PowerDurationAndTraffic();
+        final SparseArray<UidBatteryConsumer.Builder> uidBatteryConsumerBuilders =
+                builder.getUidBatteryConsumerBuilders();
+        for (int i = uidBatteryConsumerBuilders.size() - 1; i >= 0; i--) {
+            final UidBatteryConsumer.Builder app = uidBatteryConsumerBuilders.valueAt(i);
+            calculateApp(powerDurationAndTraffic, app.getBatteryStatsUid(), rawRealtimeUs,
+                    BatteryStats.STATS_SINCE_CHARGED,
+                    hasWifiPowerReporting);
+
+            totalAppDurationMs += powerDurationAndTraffic.durationMs;
+            totalAppPowerMah += powerDurationAndTraffic.powerMah;
+
+            app.setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_WIFI,
+                    powerDurationAndTraffic.durationMs);
+            app.setConsumedPower(BatteryConsumer.POWER_COMPONENT_WIFI,
+                    powerDurationAndTraffic.powerMah);
+
+            if (app.getUid() == Process.WIFI_UID) {
+                systemBatteryConsumerBuilder.addUidBatteryConsumer(app);
+                app.excludeFromBatteryUsageStats();
+            }
+        }
+
+        calculateRemaining(powerDurationAndTraffic, batteryStats, rawRealtimeUs,
+                BatteryStats.STATS_SINCE_CHARGED,
+                hasWifiPowerReporting, totalAppDurationMs, totalAppPowerMah);
+
+        systemBatteryConsumerBuilder
+                .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_WIFI,
+                        powerDurationAndTraffic.durationMs)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_WIFI,
+                        powerDurationAndTraffic.powerMah);
     }
 
     /**
@@ -64,100 +137,151 @@
 
         // batteryStats.hasWifiActivityReporting can change if we get energy data at a later point,
         // so always check this field.
-        mHasWifiPowerReporting = mHasWifiPowerController && batteryStats.hasWifiActivityReporting();
+        final boolean hasWifiPowerReporting =
+                mHasWifiPowerController && batteryStats.hasWifiActivityReporting();
 
-        super.calculate(sippers, batteryStats, rawRealtimeUs, rawUptimeUs, statsType, asUsers);
+        final BatterySipper bs = new BatterySipper(BatterySipper.DrainType.WIFI, null, 0);
 
-        BatterySipper bs = new BatterySipper(BatterySipper.DrainType.WIFI, null, 0);
-        calculateRemaining(bs, batteryStats, rawRealtimeUs, rawUptimeUs, statsType);
-
+        long totalAppDurationMs = 0;
+        double totalAppPowerMah = 0;
+        final PowerDurationAndTraffic powerDurationAndTraffic = new PowerDurationAndTraffic();
         for (int i = sippers.size() - 1; i >= 0; i--) {
-            BatterySipper app = sippers.get(i);
-            if (app.getUid() == Process.WIFI_UID) {
-                if (DEBUG) Log.d(TAG, "WiFi adding sipper " + app + ": cpu=" + app.cpuTimeMs);
-                app.isAggregated = true;
-                bs.add(app);
+            final BatterySipper app = sippers.get(i);
+            if (app.drainType == BatterySipper.DrainType.APP) {
+                calculateApp(powerDurationAndTraffic, app.uidObj, rawRealtimeUs, statsType,
+                        hasWifiPowerReporting);
+
+                totalAppDurationMs += powerDurationAndTraffic.durationMs;
+                totalAppPowerMah += powerDurationAndTraffic.powerMah;
+
+                app.wifiPowerMah = powerDurationAndTraffic.powerMah;
+                app.wifiRunningTimeMs = powerDurationAndTraffic.durationMs;
+                app.wifiRxBytes = powerDurationAndTraffic.wifiRxBytes;
+                app.wifiRxPackets = powerDurationAndTraffic.wifiRxPackets;
+                app.wifiTxBytes = powerDurationAndTraffic.wifiTxBytes;
+                app.wifiTxPackets = powerDurationAndTraffic.wifiTxPackets;
+                if (app.getUid() == Process.WIFI_UID) {
+                    if (DEBUG) Log.d(TAG, "WiFi adding sipper " + app + ": cpu=" + app.cpuTimeMs);
+                    app.isAggregated = true;
+                    bs.add(app);
+                }
             }
         }
+
+        calculateRemaining(powerDurationAndTraffic, batteryStats, rawRealtimeUs, statsType,
+                hasWifiPowerReporting, totalAppDurationMs, totalAppPowerMah);
+
+        bs.wifiRunningTimeMs += powerDurationAndTraffic.durationMs;
+        bs.wifiPowerMah += powerDurationAndTraffic.powerMah;
+
         if (bs.sumPower() > 0) {
             sippers.add(bs);
         }
     }
 
-    @Override
-    protected void calculateApp(BatterySipper app, BatteryStats.Uid u, long rawRealtimeUs,
-            long rawUptimeUs, int statsType) {
-        if (!mHasWifiPowerReporting) {
-            mWifiPowerEstimator.calculateApp(app, u, rawRealtimeUs, rawUptimeUs, statsType);
-            return;
-        }
-
-        final BatteryStats.ControllerActivityCounter counter = u.getWifiControllerActivity();
-        if (counter == null) {
-            return;
-        }
-
-        final long idleTime = counter.getIdleTimeCounter().getCountLocked(statsType);
-        final long txTime = counter.getTxTimeCounters()[0].getCountLocked(statsType);
-        final long rxTime = counter.getRxTimeCounter().getCountLocked(statsType);
-        app.wifiRunningTimeMs = idleTime + rxTime + txTime;
-        mTotalAppRunningTime += app.wifiRunningTimeMs;
-
-        app.wifiPowerMah =
-                ((idleTime * mIdleCurrentMa) + (txTime * mTxCurrentMa) + (rxTime * mRxCurrentMa))
-                        / (1000 * 60 * 60);
-        mTotalAppPowerDrain += app.wifiPowerMah;
-
-        app.wifiRxPackets = u.getNetworkActivityPackets(BatteryStats.NETWORK_WIFI_RX_DATA,
+    private void calculateApp(PowerDurationAndTraffic powerDurationAndTraffic, BatteryStats.Uid u,
+            long rawRealtimeUs,
+            int statsType, boolean hasWifiPowerReporting) {
+        powerDurationAndTraffic.wifiRxPackets = u.getNetworkActivityPackets(
+                BatteryStats.NETWORK_WIFI_RX_DATA,
                 statsType);
-        app.wifiTxPackets = u.getNetworkActivityPackets(BatteryStats.NETWORK_WIFI_TX_DATA,
+        powerDurationAndTraffic.wifiTxPackets = u.getNetworkActivityPackets(
+                BatteryStats.NETWORK_WIFI_TX_DATA,
                 statsType);
-        app.wifiRxBytes = u.getNetworkActivityBytes(BatteryStats.NETWORK_WIFI_RX_DATA,
+        powerDurationAndTraffic.wifiRxBytes = u.getNetworkActivityBytes(
+                BatteryStats.NETWORK_WIFI_RX_DATA,
                 statsType);
-        app.wifiTxBytes = u.getNetworkActivityBytes(BatteryStats.NETWORK_WIFI_TX_DATA,
+        powerDurationAndTraffic.wifiTxBytes = u.getNetworkActivityBytes(
+                BatteryStats.NETWORK_WIFI_TX_DATA,
                 statsType);
 
-        if (DEBUG && app.wifiPowerMah != 0) {
-            Log.d(TAG, "UID " + u.getUid() + ": idle=" + idleTime + "ms rx=" + rxTime + "ms tx=" +
-                    txTime + "ms power=" + formatCharge(app.wifiPowerMah));
+        if (hasWifiPowerReporting) {
+            final BatteryStats.ControllerActivityCounter counter = u.getWifiControllerActivity();
+            if (counter != null) {
+                final long idleTime = counter.getIdleTimeCounter().getCountLocked(statsType);
+                final long txTime = counter.getTxTimeCounters()[0].getCountLocked(statsType);
+                final long rxTime = counter.getRxTimeCounter().getCountLocked(statsType);
+
+                powerDurationAndTraffic.durationMs = idleTime + rxTime + txTime;
+                powerDurationAndTraffic.powerMah = mIdlePowerEstimator.calculatePower(idleTime)
+                        + mTxPowerEstimator.calculatePower(txTime)
+                        + mRxPowerEstimator.calculatePower(rxTime);
+
+                if (DEBUG && powerDurationAndTraffic.powerMah != 0) {
+                    Log.d(TAG, "UID " + u.getUid() + ": idle=" + idleTime + "ms rx=" + rxTime
+                            + "ms tx=" + txTime + "ms power=" + formatCharge(
+                            powerDurationAndTraffic.powerMah));
+                }
+            }
+        } else {
+            final double wifiPacketPower = (
+                    powerDurationAndTraffic.wifiRxPackets + powerDurationAndTraffic.wifiTxPackets)
+                    * mWifiPowerPerPacket;
+            final long wifiRunningTime = u.getWifiRunningTime(rawRealtimeUs, statsType) / 1000;
+            final long wifiScanTimeMs = u.getWifiScanTime(rawRealtimeUs, statsType) / 1000;
+            long batchScanTimeMs = 0;
+            for (int bin = 0; bin < BatteryStats.Uid.NUM_WIFI_BATCHED_SCAN_BINS; bin++) {
+                batchScanTimeMs += u.getWifiBatchedScanTime(bin, rawRealtimeUs, statsType) / 1000;
+            }
+
+            powerDurationAndTraffic.durationMs = wifiRunningTime;
+            powerDurationAndTraffic.powerMah = wifiPacketPower
+                    + mPowerOnPowerEstimator.calculatePower(wifiRunningTime)
+                    + mScanPowerEstimator.calculatePower(wifiScanTimeMs)
+                    + mBatchScanPowerEstimator.calculatePower(batchScanTimeMs);
+
+            if (DEBUG && powerDurationAndTraffic.powerMah != 0) {
+                Log.d(TAG, "UID " + u.getUid() + ": power=" + formatCharge(
+                        powerDurationAndTraffic.powerMah));
+            }
         }
     }
 
-    private void calculateRemaining(BatterySipper app, BatteryStats stats, long rawRealtimeUs,
-            long rawUptimeUs, int statsType) {
-        if (!mHasWifiPowerReporting) {
-            mWifiPowerEstimator.calculateRemaining(app, stats, rawRealtimeUs, rawUptimeUs,
-                    statsType);
-            return;
+    private void calculateRemaining(PowerDurationAndTraffic powerDurationAndTraffic,
+            BatteryStats stats, long rawRealtimeUs,
+            int statsType, boolean hasWifiPowerReporting, long totalAppDurationMs,
+            double totalAppPowerMah) {
+        long totalDurationMs;
+        double totalPowerMah;
+        if (hasWifiPowerReporting) {
+            final BatteryStats.ControllerActivityCounter counter =
+                    stats.getWifiControllerActivity();
+
+            final long idleTimeMs = counter.getIdleTimeCounter().getCountLocked(statsType);
+            final long txTimeMs = counter.getTxTimeCounters()[0].getCountLocked(statsType);
+            final long rxTimeMs = counter.getRxTimeCounter().getCountLocked(statsType);
+
+            totalDurationMs = idleTimeMs + rxTimeMs + txTimeMs;
+
+            totalPowerMah =
+                    counter.getPowerCounter().getCountLocked(statsType) / (double) (1000 * 60 * 60);
+            if (totalPowerMah == 0) {
+                // Some controllers do not report power drain, so we can calculate it here.
+                totalPowerMah = mIdlePowerEstimator.calculatePower(idleTimeMs)
+                        + mTxPowerEstimator.calculatePower(txTimeMs)
+                        + mRxPowerEstimator.calculatePower(rxTimeMs);
+            }
+        } else {
+            totalDurationMs = stats.getGlobalWifiRunningTime(rawRealtimeUs, statsType) / 1000;
+            totalPowerMah = mPowerOnPowerEstimator.calculatePower(totalDurationMs);
         }
 
-        final BatteryStats.ControllerActivityCounter counter = stats.getWifiControllerActivity();
-
-        final long idleTimeMs = counter.getIdleTimeCounter().getCountLocked(statsType);
-        final long txTimeMs = counter.getTxTimeCounters()[0].getCountLocked(statsType);
-        final long rxTimeMs = counter.getRxTimeCounter().getCountLocked(statsType);
-
-        app.wifiRunningTimeMs = Math.max(0,
-                (idleTimeMs + rxTimeMs + txTimeMs) - mTotalAppRunningTime);
-
-        double powerDrainMah = counter.getPowerCounter().getCountLocked(statsType)
-                / (double) (1000 * 60 * 60);
-        if (powerDrainMah == 0) {
-            // Some controllers do not report power drain, so we can calculate it here.
-            powerDrainMah = ((idleTimeMs * mIdleCurrentMa) + (txTimeMs * mTxCurrentMa)
-                    + (rxTimeMs * mRxCurrentMa)) / (1000 * 60 * 60);
-        }
-        app.wifiPowerMah = Math.max(0, powerDrainMah - mTotalAppPowerDrain);
+        powerDurationAndTraffic.durationMs = Math.max(0, totalDurationMs - totalAppDurationMs);
+        powerDurationAndTraffic.powerMah = Math.max(0, totalPowerMah - totalAppPowerMah);
 
         if (DEBUG) {
-            Log.d(TAG, "left over WiFi power: " + formatCharge(app.wifiPowerMah));
+            Log.d(TAG, "left over WiFi power: " + formatCharge(powerDurationAndTraffic.powerMah));
         }
     }
 
-    @Override
-    public void reset() {
-        mTotalAppPowerDrain = 0;
-        mTotalAppRunningTime = 0;
-        mWifiPowerEstimator.reset();
+    /**
+     * Return estimated power per Wi-Fi packet in mAh/packet where 1 packet = 2 KB.
+     */
+    private static double getWifiPowerPerPacket(PowerProfile profile) {
+        // TODO(b/179392913): Extract average bit rates from system
+        final long wifiBps = 1000000;
+        final double averageWifiActivePower =
+                profile.getAveragePower(PowerProfile.POWER_WIFI_ACTIVE) / 3600;
+        return averageWifiActivePower / (((double) wifiBps) / 8 / 2048);
     }
 }
diff --git a/core/java/com/android/internal/os/WifiPowerEstimator.java b/core/java/com/android/internal/os/WifiPowerEstimator.java
deleted file mode 100644
index d0a105c..0000000
--- a/core/java/com/android/internal/os/WifiPowerEstimator.java
+++ /dev/null
@@ -1,126 +0,0 @@
-/*
- * Copyright (C) 2015 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.internal.os;
-
-import android.os.BatteryStats;
-import android.os.Process;
-import android.os.UserHandle;
-import android.util.Log;
-import android.util.SparseArray;
-
-import java.util.List;
-
-/**
- * Estimates WiFi power usage based on timers in BatteryStats.
- */
-public class WifiPowerEstimator extends PowerCalculator {
-    private static final boolean DEBUG = BatteryStatsHelper.DEBUG;
-    private static final String TAG = "WifiPowerEstimator";
-    private final double mWifiPowerPerPacket;
-    private final double mWifiPowerOn;
-    private final double mWifiPowerScan;
-    private final double mWifiPowerBatchScan;
-    private long mTotalAppWifiRunningTimeMs = 0;
-
-    public WifiPowerEstimator(PowerProfile profile) {
-        mWifiPowerPerPacket = getWifiPowerPerPacket(profile);
-        mWifiPowerOn = profile.getAveragePower(PowerProfile.POWER_WIFI_ON);
-        mWifiPowerScan = profile.getAveragePower(PowerProfile.POWER_WIFI_SCAN);
-        mWifiPowerBatchScan = profile.getAveragePower(PowerProfile.POWER_WIFI_BATCHED_SCAN);
-    }
-
-    @Override
-    public void calculate(List<BatterySipper> sippers, BatteryStats batteryStats,
-            long rawRealtimeUs, long rawUptimeUs, int statsType, SparseArray<UserHandle> asUsers) {
-        super.calculate(sippers, batteryStats, rawRealtimeUs, rawUptimeUs, statsType, asUsers);
-
-        BatterySipper bs = new BatterySipper(BatterySipper.DrainType.WIFI, null, 0);
-        calculateRemaining(bs, batteryStats, rawRealtimeUs, rawUptimeUs, statsType);
-
-        for (int i = sippers.size() - 1; i >= 0; i--) {
-            BatterySipper app = sippers.get(i);
-            if (app.getUid() == Process.WIFI_UID) {
-                if (DEBUG) Log.d(TAG, "WiFi adding sipper " + app + ": cpu=" + app.cpuTimeMs);
-                app.isAggregated = true;
-                bs.add(app);
-            }
-        }
-        if (bs.sumPower() > 0) {
-            sippers.add(bs);
-        }
-    }
-
-    /**
-     * Return estimated power per Wi-Fi packet in mAh/packet where 1 packet = 2 KB.
-     */
-    private static double getWifiPowerPerPacket(PowerProfile profile) {
-        final long WIFI_BPS = 1000000; // TODO: Extract average bit rates from system
-        final double WIFI_POWER = profile.getAveragePower(PowerProfile.POWER_WIFI_ACTIVE)
-                / 3600;
-        return WIFI_POWER / (((double)WIFI_BPS) / 8 / 2048);
-    }
-
-    @Override
-    protected void calculateApp(BatterySipper app, BatteryStats.Uid u, long rawRealtimeUs,
-                             long rawUptimeUs, int statsType) {
-        app.wifiRxPackets = u.getNetworkActivityPackets(BatteryStats.NETWORK_WIFI_RX_DATA,
-                statsType);
-        app.wifiTxPackets = u.getNetworkActivityPackets(BatteryStats.NETWORK_WIFI_TX_DATA,
-                statsType);
-        app.wifiRxBytes = u.getNetworkActivityBytes(BatteryStats.NETWORK_WIFI_RX_DATA,
-                statsType);
-        app.wifiTxBytes = u.getNetworkActivityBytes(BatteryStats.NETWORK_WIFI_TX_DATA,
-                statsType);
-
-        final double wifiPacketPower = (app.wifiRxPackets + app.wifiTxPackets)
-                * mWifiPowerPerPacket;
-
-        app.wifiRunningTimeMs = u.getWifiRunningTime(rawRealtimeUs, statsType) / 1000;
-        mTotalAppWifiRunningTimeMs += app.wifiRunningTimeMs;
-        final double wifiLockPower = (app.wifiRunningTimeMs * mWifiPowerOn) / (1000*60*60);
-
-        final long wifiScanTimeMs = u.getWifiScanTime(rawRealtimeUs, statsType) / 1000;
-        final double wifiScanPower = (wifiScanTimeMs * mWifiPowerScan) / (1000*60*60);
-
-        double wifiBatchScanPower = 0;
-        for (int bin = 0; bin < BatteryStats.Uid.NUM_WIFI_BATCHED_SCAN_BINS; bin++) {
-            final long batchScanTimeMs =
-                    u.getWifiBatchedScanTime(bin, rawRealtimeUs, statsType) / 1000;
-            final double batchScanPower = (batchScanTimeMs * mWifiPowerBatchScan) / (1000*60*60);
-            wifiBatchScanPower += batchScanPower;
-        }
-
-        app.wifiPowerMah = wifiPacketPower + wifiLockPower + wifiScanPower + wifiBatchScanPower;
-        if (DEBUG && app.wifiPowerMah != 0) {
-            Log.d(TAG, "UID " + u.getUid() + ": power=" + formatCharge(app.wifiPowerMah));
-        }
-    }
-
-    void calculateRemaining(BatterySipper app, BatteryStats stats, long rawRealtimeUs,
-                                   long rawUptimeUs, int statsType) {
-        final long totalRunningTimeMs = stats.getGlobalWifiRunningTime(rawRealtimeUs, statsType)
-                / 1000;
-        final double powerDrain = ((totalRunningTimeMs - mTotalAppWifiRunningTimeMs) * mWifiPowerOn)
-                / (1000*60*60);
-        app.wifiRunningTimeMs = totalRunningTimeMs;
-        app.wifiPowerMah = Math.max(0, powerDrain);
-    }
-
-    @Override
-    public void reset() {
-        mTotalAppWifiRunningTimeMs = 0;
-    }
-}
diff --git a/core/proto/android/server/powerstatsservice.proto b/core/proto/android/server/powerstatsservice.proto
index 4b1ee02..f64e146 100644
--- a/core/proto/android/server/powerstatsservice.proto
+++ b/core/proto/android/server/powerstatsservice.proto
@@ -144,7 +144,7 @@
      */
     optional int64 total_state_entry_count = 3;
     /**
-     * Last time this state was entered. Time in milliseconds since boot
+     * Last time this state was entered. Walltime in milliseconds since Unix epoch.
      */
     optional int64 last_entry_timestamp_ms = 4;
 }
@@ -208,7 +208,7 @@
     /** Unique index identifying the energy consumer. */
     optional int32 id = 1;
 
-    /** Time since device boot(CLOCK_BOOTTIME) in milli-seconds */
+    /** Walltime in milliseconds since Unix epoch */
     optional int64 timestamp_ms = 2;
 
     /** Accumulated energy since device boot in microwatt-seconds (uWs) */
@@ -247,7 +247,7 @@
      */
     optional int32 id = 1;
 
-    /** Time since device boot(CLOCK_BOOTTIME) in milli-seconds */
+    /** Walltime in milliseconds since Unix epoch */
     optional int64 timestamp_ms = 2;
 
     /** Accumulated energy since device boot in microwatt-seconds (uWs) */
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java
index b819d9e..a1beecc 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java
@@ -71,6 +71,7 @@
         UserPowerCalculatorTest.class,
         VideoPowerCalculatorTest.class,
         WakelockPowerCalculatorTest.class,
+        WifiPowerCalculatorTest.class,
 
         com.android.internal.power.MeasuredEnergyStatsTest.class
     })
diff --git a/core/tests/coretests/src/com/android/internal/os/WifiPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/WifiPowerCalculatorTest.java
new file mode 100644
index 0000000..e100545
--- /dev/null
+++ b/core/tests/coretests/src/com/android/internal/os/WifiPowerCalculatorTest.java
@@ -0,0 +1,128 @@
+/*
+ * Copyright (C) 2021 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.internal.os;
+
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.net.NetworkCapabilities;
+import android.net.NetworkStats;
+import android.os.BatteryConsumer;
+import android.os.Process;
+import android.os.SystemBatteryConsumer;
+import android.os.UidBatteryConsumer;
+import android.os.WorkSource;
+import android.os.connectivity.WifiActivityEnergyInfo;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class WifiPowerCalculatorTest {
+    private static final double PRECISION = 0.00001;
+
+    private static final int APP_UID = Process.FIRST_APPLICATION_UID + 42;
+
+    @Rule
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_IDLE, 360.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX, 480.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_TX, 720.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_ON, 360.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_SCAN, 480.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_BATCHED_SCAN, 720.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_ACTIVE, 1080.0);
+
+    @Test
+    public void testPowerControllerBasedModel() {
+        BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+
+        batteryStats.noteNetworkInterfaceForTransports("wifi",
+                new int[]{NetworkCapabilities.TRANSPORT_WIFI});
+
+        NetworkStats networkStats = new NetworkStats(10000, 1)
+                .insertEntry("wifi", APP_UID, 0, 0, 1000, 100, 2000, 20, 100)
+                .insertEntry("wifi", Process.WIFI_UID, 0, 0, 1111, 111, 2222, 22, 111);
+        mStatsRule.setNetworkStats(networkStats);
+
+        WifiActivityEnergyInfo energyInfo = new WifiActivityEnergyInfo(10000,
+                WifiActivityEnergyInfo.STACK_STATE_STATE_ACTIVE, 1000, 2000, 3000, 4000);
+
+        batteryStats.updateWifiState(energyInfo, 1000, 1000);
+
+        WifiPowerCalculator calculator = new WifiPowerCalculator(mStatsRule.getPowerProfile());
+        mStatsRule.apply(calculator);
+
+        UidBatteryConsumer uidConsumer = mStatsRule.getUidBatteryConsumer(APP_UID);
+        assertThat(uidConsumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_WIFI))
+                .isEqualTo(1423);
+        assertThat(uidConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_WIFI))
+                .isWithin(PRECISION).of(0.2214666);
+
+        SystemBatteryConsumer systemConsumer =
+                mStatsRule.getSystemBatteryConsumer(SystemBatteryConsumer.DRAIN_TYPE_WIFI);
+        assertThat(systemConsumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_WIFI))
+                .isEqualTo(5577);
+        assertThat(systemConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_WIFI))
+                .isWithin(PRECISION).of(0.645200);
+    }
+
+    @Test
+    public void testTimerBasedModel() {
+        BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+
+        batteryStats.noteNetworkInterfaceForTransports("wifi",
+                new int[]{NetworkCapabilities.TRANSPORT_WIFI});
+
+        NetworkStats networkStats = new NetworkStats(10000, 1)
+                .insertEntry("wifi", APP_UID, 0, 0, 1000, 100, 2000, 20, 100)
+                .insertEntry("wifi", Process.WIFI_UID, 0, 0, 1111, 111, 2222, 22, 111);
+        mStatsRule.setNetworkStats(networkStats);
+
+        batteryStats.noteWifiScanStartedLocked(APP_UID, 1000, 1000);
+        batteryStats.noteWifiScanStoppedLocked(APP_UID, 2000, 2000);
+        batteryStats.noteWifiRunningLocked(new WorkSource(APP_UID), 3000, 3000);
+        batteryStats.noteWifiStoppedLocked(new WorkSource(APP_UID), 4000, 4000);
+        batteryStats.noteWifiRunningLocked(new WorkSource(Process.WIFI_UID), 1111, 2222);
+        batteryStats.noteWifiStoppedLocked(new WorkSource(Process.WIFI_UID), 3333, 4444);
+
+        // Don't pass WifiActivityEnergyInfo, making WifiPowerCalculator rely exclusively
+        // on the packet counts.
+        batteryStats.updateWifiState(/* energyInfo */ null, 1000, 1000);
+
+        WifiPowerCalculator calculator = new WifiPowerCalculator(mStatsRule.getPowerProfile());
+        mStatsRule.apply(calculator);
+
+        UidBatteryConsumer uidConsumer = mStatsRule.getUidBatteryConsumer(APP_UID);
+        assertThat(uidConsumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_WIFI))
+                .isEqualTo(1000);
+        assertThat(uidConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_WIFI))
+                .isWithin(PRECISION).of(0.8231573);
+
+        SystemBatteryConsumer systemConsumer =
+                mStatsRule.getSystemBatteryConsumer(SystemBatteryConsumer.DRAIN_TYPE_WIFI);
+        assertThat(systemConsumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_WIFI))
+                .isEqualTo(2222);
+        assertThat(systemConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_WIFI))
+                .isWithin(PRECISION).of(0.8759216);
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java
index 4564af4..f118b1e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayLayout.java
@@ -50,6 +50,7 @@
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
+import java.util.Objects;
 
 /**
  * Contains information about the layout-properties of a display. This refers to internal layout
@@ -82,6 +83,31 @@
     private boolean mHasStatusBar = false;
     private int mNavBarFrameHeight = 0;
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof DisplayLayout)) return false;
+        final DisplayLayout other = (DisplayLayout) o;
+        return mUiMode == other.mUiMode
+                && mWidth == other.mWidth
+                && mHeight == other.mHeight
+                && Objects.equals(mCutout, other.mCutout)
+                && mRotation == other.mRotation
+                && mDensityDpi == other.mDensityDpi
+                && Objects.equals(mNonDecorInsets, other.mNonDecorInsets)
+                && Objects.equals(mStableInsets, other.mStableInsets)
+                && mHasNavigationBar == other.mHasNavigationBar
+                && mHasStatusBar == other.mHasStatusBar
+                && mNavBarFrameHeight == other.mNavBarFrameHeight;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(mUiMode, mWidth, mHeight, mCutout, mRotation, mDensityDpi,
+                mNonDecorInsets, mStableInsets, mHasNavigationBar, mHasStatusBar,
+                mNavBarFrameHeight);
+    }
+
     /**
      * Create empty layout.
      */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java
index 66560d3..a52db24 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java
@@ -90,15 +90,15 @@
 
     private final PipSurfaceTransactionHelper mSurfaceTransactionHelper;
 
-    private PipTransitionAnimator mCurrentAnimator;
-
-    private ThreadLocal<AnimationHandler> mSfAnimationHandlerThreadLocal =
+    private final ThreadLocal<AnimationHandler> mSfAnimationHandlerThreadLocal =
             ThreadLocal.withInitial(() -> {
                 AnimationHandler handler = new AnimationHandler();
                 handler.setProvider(new SfVsyncFrameCallbackProvider());
                 return handler;
             });
 
+    private PipTransitionAnimator mCurrentAnimator;
+
     public PipAnimationController(PipSurfaceTransactionHelper helper) {
         mSurfaceTransactionHelper = helper;
     }
@@ -268,6 +268,7 @@
             if (mPipAnimationCallback != null) {
                 mPipAnimationCallback.onPipAnimationEnd(mTaskInfo, tx, this);
             }
+            mTransitionDirection = TRANSITION_DIRECTION_NONE;
         }
 
         @Override
@@ -275,6 +276,7 @@
             if (mPipAnimationCallback != null) {
                 mPipAnimationCallback.onPipAnimationCancel(mTaskInfo, this);
             }
+            mTransitionDirection = TRANSITION_DIRECTION_NONE;
         }
 
         @Override public void onAnimationRepeat(Animator animation) {}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
index 7fff37d..9a584c6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
@@ -65,6 +65,7 @@
 import com.android.wm.shell.pip.PipUtils;
 
 import java.io.PrintWriter;
+import java.util.Objects;
 import java.util.function.Consumer;
 
 /**
@@ -101,9 +102,12 @@
      */
     private final DisplayChangeController.OnDisplayChangingListener mRotationController = (
             int displayId, int fromRotation, int toRotation, WindowContainerTransaction t) -> {
-        if (!mPipTaskOrganizer.isInPip() || mPipTaskOrganizer.isDeferringEnterPipAnimation()) {
-            // Skip if we aren't in PIP or haven't actually entered PIP yet. We still need to update
-            // the display layout in the bounds handler in this case.
+        if (!mPipTaskOrganizer.isInPip()
+                || mPipBoundsState.getDisplayLayout().rotation() == toRotation
+                || mPipTaskOrganizer.isDeferringEnterPipAnimation()) {
+            // Skip if the same rotation has been set or we aren't in PIP or haven't actually
+            // entered PIP yet. We still need to update the display layout in the bounds handler
+            // in this case.
             onDisplayRotationChangedNotInPip(mContext, toRotation);
             // do not forget to update the movement bounds as well.
             updateMovementBounds(mPipBoundsState.getNormalBounds(), true /* fromRotation */,
@@ -378,6 +382,9 @@
     }
 
     private void onDisplayChanged(DisplayLayout layout, boolean saveRestoreSnapFraction) {
+        if (Objects.equals(layout, mPipBoundsState.getDisplayLayout())) {
+            return;
+        }
         Runnable updateDisplayLayout = () -> {
             mPipBoundsState.setDisplayLayout(layout);
             updateMovementBounds(null /* toBounds */,
@@ -476,8 +483,11 @@
             int launcherRotation, int shelfHeight) {
         setShelfHeightLocked(shelfHeight > 0 /* visible */, shelfHeight);
         onDisplayRotationChangedNotInPip(mContext, launcherRotation);
-        return mPipTaskOrganizer.startSwipePipToHome(componentName, activityInfo,
+        final Rect entryBounds = mPipTaskOrganizer.startSwipePipToHome(componentName, activityInfo,
                 pictureInPictureParams);
+        // sync mPipBoundsState with the newly calculated bounds.
+        mPipBoundsState.setNormalBounds(entryBounds);
+        return entryBounds;
     }
 
     private void stopSwipePipToHome(ComponentName componentName, Rect destinationBounds) {
diff --git a/packages/SystemUI/res/layout/qs_tile_label.xml b/packages/SystemUI/res/layout/qs_tile_label.xml
index 571cbbc..1ce87c1 100644
--- a/packages/SystemUI/res/layout/qs_tile_label.xml
+++ b/packages/SystemUI/res/layout/qs_tile_label.xml
@@ -20,7 +20,6 @@
     android:layout_height="wrap_content"
     android:clipChildren="false"
     android:clipToPadding="false"
-    android:minHeight="48dp"
     android:paddingTop="12dp">
     <LinearLayout
         android:id="@+id/label_group"
diff --git a/packages/SystemUI/res/values/config.xml b/packages/SystemUI/res/values/config.xml
index f13fdd3..0893c14 100644
--- a/packages/SystemUI/res/values/config.xml
+++ b/packages/SystemUI/res/values/config.xml
@@ -91,6 +91,9 @@
     <!-- The number of columns in the QuickSettings -->
     <integer name="quick_settings_num_columns">3</integer>
 
+    <!-- The number of columns in the Quick Settings customizer -->
+    <integer name="quick_settings_edit_num_columns">@integer/quick_settings_num_columns</integer>
+
     <!-- The number of rows in the QuickSettings -->
     <integer name="quick_settings_max_rows">3</integer>
 
diff --git a/packages/SystemUI/src/com/android/keyguard/NumPadButton.java b/packages/SystemUI/src/com/android/keyguard/NumPadButton.java
index 8cb1bc4..5db3349 100644
--- a/packages/SystemUI/src/com/android/keyguard/NumPadButton.java
+++ b/packages/SystemUI/src/com/android/keyguard/NumPadButton.java
@@ -69,7 +69,9 @@
 
     @Override
     public boolean onTouchEvent(MotionEvent event) {
-        if (mAnimator != null) mAnimator.start();
+        if (event.getActionMasked() == MotionEvent.ACTION_DOWN) {
+            if (mAnimator != null) mAnimator.start();
+        }
         return super.onTouchEvent(event);
     }
 
diff --git a/packages/SystemUI/src/com/android/keyguard/NumPadKey.java b/packages/SystemUI/src/com/android/keyguard/NumPadKey.java
index a4a781d..e6a9d4f 100644
--- a/packages/SystemUI/src/com/android/keyguard/NumPadKey.java
+++ b/packages/SystemUI/src/com/android/keyguard/NumPadKey.java
@@ -160,10 +160,9 @@
     public boolean onTouchEvent(MotionEvent event) {
         if (event.getActionMasked() == MotionEvent.ACTION_DOWN) {
             doHapticKeyClick();
+            if (mAnimator != null) mAnimator.start();
         }
 
-        if (mAnimator != null) mAnimator.start();
-
         return super.onTouchEvent(event);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java b/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java
index c8edaec..fe76668 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSAnimator.java
@@ -29,7 +29,6 @@
 import com.android.systemui.qs.TouchAnimator.Builder;
 import com.android.systemui.qs.TouchAnimator.Listener;
 import com.android.systemui.qs.dagger.QSScope;
-import com.android.systemui.qs.tileimpl.QSTileBaseView;
 import com.android.systemui.statusbar.FeatureFlags;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.tuner.TunerService.Tunable;
@@ -56,7 +55,8 @@
     private final ArrayList<View> mAllViews = new ArrayList<>();
     /**
      * List of {@link View}s representing Quick Settings that are being animated from the quick QS
-     * position to the normal QS panel.
+     * position to the normal QS panel. These views will only show once the animation is complete,
+     * to prevent overlapping of semi transparent views
      */
     private final ArrayList<View> mQuickQsViews = new ArrayList<>();
     private final QuickQSPanel mQuickQsPanel;
@@ -233,7 +233,6 @@
                 // Quick tiles.
                 QSTileView quickTileView = mQuickQSPanelController.getTileView(tile);
                 if (quickTileView == null) continue;
-                View qqsBgCircle = ((QSTileBaseView) quickTileView).getBgCircle();
 
                 getRelativePosition(loc1, quickTileView.getIcon().getIconView(), view);
                 getRelativePosition(loc2, tileIcon, view);
@@ -255,11 +254,6 @@
                     translationXBuilder.addFloat(tileView, "translationX", -xDiff, 0);
                     translationYBuilder.addFloat(tileView, "translationY", -yDiff, 0);
 
-                    if (mFeatureFlags.isQSLabelsEnabled()) {
-                        firstPageBuilder.addFloat(qqsBgCircle, "alpha", 1, 1, 0);
-                        mAllViews.add(qqsBgCircle);
-                    }
-
                 } else { // These tiles disappear when expanding
                     firstPageBuilder.addFloat(quickTileView, "alpha", 1, 0);
                     translationYBuilder.addFloat(quickTileView, "translationY", 0, yDiff);
@@ -271,7 +265,11 @@
                             translationX);
                 }
 
-                mQuickQsViews.add(tileView.getIconWithBackground());
+                if (mFeatureFlags.isQSLabelsEnabled()) {
+                    mQuickQsViews.add(tileView);
+                } else {
+                    mQuickQsViews.add(tileView.getIconWithBackground());
+                }
                 mAllViews.add(tileView.getIcon());
                 mAllViews.add(quickTileView);
             } else if (mFullRows && isIconInAnimatedRow(count)) {
@@ -362,7 +360,7 @@
         if(view == parent || view == null) return;
         // Ignore tile pages as they can have some offset we don't want to take into account in
         // RTL.
-        if (!(view instanceof PagedTileLayout.TilePage || view instanceof SideLabelTileLayout)) {
+        if (!isAPage(view)) {
             loc1[0] += view.getLeft();
             loc1[1] += view.getTop();
         }
@@ -374,6 +372,16 @@
         getRelativePositionInt(loc1, (View) view.getParent(), parent);
     }
 
+    // Returns true if the view is a possible page in PagedTileLayout
+    private boolean isAPage(View view) {
+        if (view instanceof PagedTileLayout.TilePage) {
+            return true;
+        } else if (view instanceof SideLabelTileLayout) {
+            return !(view instanceof QuickQSPanel.QQSSideLabelTileLayout);
+        }
+        return false;
+    }
+
     public void setPosition(float position) {
         if (mNeedsAnimatorUpdate) {
             updateAnimators();
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
index 91ae571..0563307 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
@@ -112,7 +112,7 @@
     private int mMediaTotalBottomMargin;
     private int mFooterMarginStartHorizontal;
     private Consumer<Boolean> mMediaVisibilityChangedListener;
-    private boolean mSideLabels;
+    protected boolean mSideLabels;
 
     public QSPanel(Context context, AttributeSet attrs) {
         super(context, attrs);
@@ -201,16 +201,20 @@
                 mFooterPageIndicator.setNumPages(((PagedTileLayout) mTileLayout).getNumPages());
             }
 
-            // Allow the UI to be as big as it want's to, we're in a scroll view
-            int newHeight = 10000;
-            int availableHeight = MeasureSpec.getSize(heightMeasureSpec);
-            int excessHeight = newHeight - availableHeight;
-            // Measure with EXACTLY. That way, The content will only use excess height and will
-            // be measured last, after other views and padding is accounted for. This only
-            // works because our Layouts in here remeasure themselves with the exact content
-            // height.
-            heightMeasureSpec = MeasureSpec.makeMeasureSpec(newHeight, MeasureSpec.EXACTLY);
-            ((PagedTileLayout) mTileLayout).setExcessHeight(excessHeight);
+            // In landscape, mTileLayout's parent is not the panel but a view that contains the
+            // tile layout and the media controls.
+            if (((View) mTileLayout).getParent() == this) {
+                // Allow the UI to be as big as it want's to, we're in a scroll view
+                int newHeight = 10000;
+                int availableHeight = MeasureSpec.getSize(heightMeasureSpec);
+                int excessHeight = newHeight - availableHeight;
+                // Measure with EXACTLY. That way, The content will only use excess height and will
+                // be measured last, after other views and padding is accounted for. This only
+                // works because our Layouts in here remeasure themselves with the exact content
+                // height.
+                heightMeasureSpec = MeasureSpec.makeMeasureSpec(newHeight, MeasureSpec.EXACTLY);
+                ((PagedTileLayout) mTileLayout).setExcessHeight(excessHeight);
+            }
         }
         super.onMeasure(widthMeasureSpec, heightMeasureSpec);
 
@@ -758,7 +762,7 @@
                 // Let's use 3 columns to match the current layout
                 int columns;
                 if (mSideLabels) {
-                    columns = horizontal ? 1 : 2;
+                    columns = horizontal ? 2 : 4;
                 } else {
                     columns = horizontal ? 3 : TileLayout.NO_MAX_COLUMNS;
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java
index fcb35e2..d60801e 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java
@@ -18,7 +18,6 @@
 
 import static com.android.systemui.media.dagger.MediaModule.QS_PANEL;
 import static com.android.systemui.qs.QSPanel.QS_SHOW_BRIGHTNESS;
-import static com.android.systemui.qs.dagger.QSFlagsModule.QS_LABELS_FLAG;
 import static com.android.systemui.qs.dagger.QSFragmentModule.QS_USING_MEDIA_PLAYER;
 
 import android.annotation.NonNull;
@@ -66,7 +65,6 @@
 
     private BrightnessMirrorController mBrightnessMirrorController;
     private boolean mGridContentVisible = true;
-    private boolean mQsLabelsFlag;
 
     private final QSPanel.OnConfigurationChangedListener mOnConfigurationChangedListener =
             new QSPanel.OnConfigurationChangedListener() {
@@ -93,7 +91,6 @@
             DumpManager dumpManager, MetricsLogger metricsLogger, UiEventLogger uiEventLogger,
             QSLogger qsLogger, BrightnessController.Factory brightnessControllerFactory,
             BrightnessSlider.Factory brightnessSliderFactory,
-            @Named(QS_LABELS_FLAG) boolean qsLabelsFlag,
             FeatureFlags featureFlags) {
         super(view, qstileHost, qsCustomizerController, usingMediaPlayer, mediaHost,
                 metricsLogger, uiEventLogger, qsLogger, dumpManager, featureFlags);
@@ -108,9 +105,6 @@
         mView.setBrightnessView(mBrightnessSlider.getRootView());
 
         mBrightnessController = brightnessControllerFactory.create(mBrightnessSlider);
-
-        mQsLabelsFlag = qsLabelsFlag;
-        mSideLabels = qsLabelsFlag;
     }
 
     @Override
@@ -329,7 +323,7 @@
         @Override
         public void onTuningChanged(String key, String newValue) {
             if (QS_REMOVE_LABELS.equals(key)) {
-                if (!mQsLabelsFlag) return;
+                if (!mQSLabelFlag) return;
                 boolean newShowLabels = newValue == null || "0".equals(newValue);
                 if (mShowLabels == newShowLabels) return;
                 mShowLabels = newShowLabels;
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java
index 9426e71..f1174fb 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java
@@ -75,7 +75,7 @@
 
     private final QSHost.Callback mQSHostCallback = this::setTiles;
     protected boolean mShowLabels = true;
-    protected boolean mSideLabels;
+    protected boolean mQSLabelFlag;
 
     private final QSPanel.OnConfigurationChangedListener mOnConfigurationChangedListener =
             new QSPanel.OnConfigurationChangedListener() {
@@ -118,11 +118,12 @@
         mQSLogger = qsLogger;
         mDumpManager = dumpManager;
         mFeatureFlags = featureFlags;
+        mQSLabelFlag = featureFlags.isQSLabelsEnabled();
     }
 
     @Override
     protected void onInit() {
-        mView.initialize(mSideLabels);
+        mView.initialize(mQSLabelFlag);
         mQSLogger.logAllTilesChangeListening(mView.isListening(), mView.getDumpableTag(), "");
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java b/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java
index a29ac3b..9b66b59 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanel.java
@@ -69,12 +69,22 @@
 
     @Override
     public TileLayout createRegularTileLayout() {
-        return new QuickQSPanel.HeaderTileLayout(mContext);
+        if (mSideLabels) {
+            return new QQSSideLabelTileLayout(mContext);
+        } else {
+            return new QuickQSPanel.HeaderTileLayout(mContext);
+        }
     }
 
     @Override
     protected QSTileLayout createHorizontalTileLayout() {
-        return new DoubleLineTileLayout(mContext);
+        if (mSideLabels) {
+            TileLayout t = createRegularTileLayout();
+            t.setMaxColumns(2);
+            return t;
+        } else {
+            return new DoubleLineTileLayout(mContext);
+        }
     }
 
     @Override
@@ -331,4 +341,38 @@
             }
         }
     }
+
+    static class QQSSideLabelTileLayout extends SideLabelTileLayout {
+        QQSSideLabelTileLayout(Context context) {
+            super(context, null);
+            setClipChildren(false);
+            setClipToPadding(false);
+            LinearLayout.LayoutParams lp = new LinearLayout.LayoutParams(LayoutParams.MATCH_PARENT,
+                    LayoutParams.WRAP_CONTENT);
+            lp.gravity = Gravity.CENTER_HORIZONTAL;
+            setLayoutParams(lp);
+            setMaxColumns(4);
+        }
+
+        @Override
+        public boolean updateResources() {
+            boolean b = super.updateResources();
+            mMaxAllowedRows = 2;
+            return b;
+        }
+
+        @Override
+        protected void onConfigurationChanged(Configuration newConfig) {
+            super.onConfigurationChanged(newConfig);
+            updateResources();
+        }
+
+        @Override
+        protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
+            // Make sure to always use the correct number of rows. As it's determined by the
+            // columns, just use as many as needed.
+            updateMaxRows(10000, mRecords.size());
+            super.onMeasure(widthMeasureSpec, heightMeasureSpec);
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanelController.java b/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanelController.java
index 383e932..671f8f7 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanelController.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QuickQSPanelController.java
@@ -17,7 +17,6 @@
 package com.android.systemui.qs;
 
 import static com.android.systemui.media.dagger.MediaModule.QUICK_QS_PANEL;
-import static com.android.systemui.qs.dagger.QSFlagsModule.QS_LABELS_FLAG;
 import static com.android.systemui.qs.dagger.QSFragmentModule.QS_USING_MEDIA_PLAYER;
 
 import com.android.internal.logging.MetricsLogger;
@@ -42,8 +41,6 @@
 @QSScope
 public class QuickQSPanelController extends QSPanelControllerBase<QuickQSPanel> {
 
-    private boolean mUseSideLabels;
-
     private final QSPanel.OnConfigurationChangedListener mOnConfigurationChangedListener =
             newConfig -> {
                 int newMaxTiles = getResources().getInteger(R.integer.quick_qs_panel_max_columns);
@@ -58,12 +55,10 @@
             @Named(QS_USING_MEDIA_PLAYER) boolean usingMediaPlayer,
             @Named(QUICK_QS_PANEL) MediaHost mediaHost,
             MetricsLogger metricsLogger, UiEventLogger uiEventLogger, QSLogger qsLogger,
-            DumpManager dumpManager, @Named(QS_LABELS_FLAG) boolean qsLabelsFlag,
-            FeatureFlags featureFlags
+            DumpManager dumpManager, FeatureFlags featureFlags
     ) {
         super(view, qsTileHost, qsCustomizerController, usingMediaPlayer, mediaHost, metricsLogger,
                 uiEventLogger, qsLogger, dumpManager, featureFlags);
-        mUseSideLabels = qsLabelsFlag;
     }
 
     @Override
@@ -104,19 +99,7 @@
                 break;
             }
         }
-        if (mUseSideLabels) {
-            List<QSTile> newTiles = new ArrayList<>();
-            for (int i = 0; i < tiles.size(); i += 2) {
-                newTiles.add(tiles.get(i));
-            }
-            for (int i = 1; i < tiles.size(); i += 2) {
-                newTiles.add(tiles.get(i));
-            }
-            super.setTiles(newTiles, true);
-
-        } else {
-            super.setTiles(tiles, true);
-        }
+        super.setTiles(tiles, !mQSLabelFlag);
     }
 
     /** */
diff --git a/packages/SystemUI/src/com/android/systemui/qs/SideLabelTileLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/SideLabelTileLayout.kt
index 74a7ac1..4de4a78 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/SideLabelTileLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/SideLabelTileLayout.kt
@@ -20,11 +20,13 @@
 import android.util.AttributeSet
 import com.android.systemui.R
 
-open class SideLabelTileLayout(context: Context, attrs: AttributeSet) : TileLayout(context, attrs) {
+open class SideLabelTileLayout(
+    context: Context,
+    attrs: AttributeSet?
+) : TileLayout(context, attrs) {
 
     override fun updateResources(): Boolean {
         return super.updateResources().also {
-            mResourceColumns = 2
             mMaxAllowedRows = 4
             mCellMarginHorizontal = (mCellMarginHorizontal * 1.2).toInt()
             mCellMarginVertical = mCellMarginHorizontal
diff --git a/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java b/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java
index d559e07..14cbf98 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/TileLayout.java
@@ -124,6 +124,7 @@
     public boolean updateResources() {
         final Resources res = mContext.getResources();
         mResourceColumns = Math.max(1, res.getInteger(R.integer.quick_settings_num_columns));
+        updateColumns();
         mMaxCellHeight = mContext.getResources().getDimensionPixelSize(R.dimen.qs_tile_height);
         mCellMarginHorizontal = res.getDimensionPixelSize(R.dimen.qs_tile_margin_horizontal);
         mCellMarginVertical= res.getDimensionPixelSize(R.dimen.qs_tile_margin_vertical);
diff --git a/packages/SystemUI/src/com/android/systemui/qs/customize/TileAdapter.java b/packages/SystemUI/src/com/android/systemui/qs/customize/TileAdapter.java
index 048fdc3..7a91421 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/customize/TileAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/customize/TileAdapter.java
@@ -75,6 +75,8 @@
     private static final int ACTION_ADD = 1;
     private static final int ACTION_MOVE = 2;
 
+    private static final int NUM_COLUMNS_ID = R.integer.quick_settings_edit_num_columns;
+
     private final Context mContext;
 
     private final Handler mHandler = new Handler();
@@ -87,6 +89,7 @@
     private int mEditIndex;
     private int mTileDividerIndex;
     private int mFocusIndex;
+
     private boolean mNeedsFocus;
     private List<String> mCurrentSpecs;
     private List<TileInfo> mOtherTiles;
@@ -109,7 +112,7 @@
         mDecoration = new TileItemDecoration(context);
         mMarginDecoration = new MarginTileDecoration();
         mMinNumTiles = context.getResources().getInteger(R.integer.quick_settings_min_num_tiles);
-        mNumColumns = context.getResources().getInteger(R.integer.quick_settings_num_columns);
+        mNumColumns = context.getResources().getInteger(NUM_COLUMNS_ID);
         mAccessibilityDelegate = new TileAdapterDelegate();
     }
 
@@ -129,7 +132,7 @@
      * @return {@code true} if the number of columns changed, {@code false} otherwise
      */
     public boolean updateNumColumns() {
-        int numColumns = mContext.getResources().getInteger(R.integer.quick_settings_num_columns);
+        int numColumns = mContext.getResources().getInteger(NUM_COLUMNS_ID);
         if (numColumns != mNumColumns) {
             mNumColumns = numColumns;
             return true;
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileView.java b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileView.java
index a699e2e..424aafa 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileView.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileView.java
@@ -105,24 +105,27 @@
 
     @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
+        mLabel.setSingleLine(false);
         super.onMeasure(widthMeasureSpec, heightMeasureSpec);
 
-        // Remeasure view if the primary label requires more then 2 lines or the secondary label
-        // text will be cut off.
-        if (mLabel.getLineCount() > mMaxLabelLines || !TextUtils.isEmpty(mSecondLine.getText())
-                        && mSecondLine.getLineHeight() > mSecondLine.getHeight()) {
-            if (!mLabel.isSingleLine()) {
-                mLabel.setSingleLine();
-                super.onMeasure(widthMeasureSpec, heightMeasureSpec);
-            }
-        } else {
-            if (mLabel.isSingleLine()) {
-                mLabel.setSingleLine(false);
-                super.onMeasure(widthMeasureSpec, heightMeasureSpec);
-            }
+        // Remeasure view if the primary label requires more than mMaxLabelLines lines or the
+        // secondary label text will be cut off.
+        if (shouldLabelBeSingleLine()) {
+            mLabel.setSingleLine();
+            super.onMeasure(widthMeasureSpec, heightMeasureSpec);
         }
     }
 
+    private boolean shouldLabelBeSingleLine() {
+        if (mLabel.getLineCount() > mMaxLabelLines) {
+            return true;
+        } else if (!TextUtils.isEmpty(mSecondLine.getText())
+                && mLabel.getLineCount() > mMaxLabelLines - 1) {
+            return true;
+        }
+        return false;
+    }
+
     @Override
     protected void handleStateChanged(QSTile.State state) {
         super.handleStateChanged(state);
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewHorizontal.kt b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewHorizontal.kt
index dc81b702..c98de8c 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewHorizontal.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewHorizontal.kt
@@ -24,8 +24,8 @@
 import android.graphics.drawable.RippleDrawable
 import android.service.quicksettings.Tile.STATE_ACTIVE
 import android.view.Gravity
-import android.view.View
 import android.widget.LinearLayout
+import android.widget.RelativeLayout
 import com.android.systemui.R
 import com.android.systemui.plugins.qs.QSIconView
 import com.android.systemui.plugins.qs.QSTile
@@ -37,7 +37,6 @@
 ) : QSTileView(context, icon, false) {
 
     private var paintDrawable: PaintDrawable? = null
-    private var divider: View? = null
 
     init {
         orientation = HORIZONTAL
@@ -49,7 +48,12 @@
 
     override fun createLabel() {
         super.createLabel()
-        findViewById<LinearLayout>(R.id.label_group)?.gravity = Gravity.START
+        findViewById<LinearLayout>(R.id.label_group)?.apply {
+            gravity = Gravity.START
+            (layoutParams as? RelativeLayout.LayoutParams)?.apply {
+                removeRule(RelativeLayout.ALIGN_PARENT_TOP)
+            }
+        }
         mLabel.gravity = Gravity.START
         mLabel.textDirection = TEXT_DIRECTION_LOCALE
         mSecondLine.gravity = Gravity.START
@@ -57,7 +61,7 @@
         val padding = context.resources.getDimensionPixelSize(R.dimen.qs_tile_side_label_padding)
         mLabelContainer.setPaddingRelative(0, padding, padding, padding)
         (mLabelContainer.layoutParams as LayoutParams).gravity =
-                Gravity.CENTER_VERTICAL or Gravity.START
+            Gravity.CENTER_VERTICAL or Gravity.START
     }
 
     override fun updateRippleSize() {
@@ -93,7 +97,6 @@
         paintDrawable?.setTint(getCircleColor(state.state))
         mSecondLine.setTextColor(mLabel.textColors)
         mLabelContainer.background = null
-        divider?.backgroundTintList = mLabel.textColors
     }
 
     override fun handleExpand(dualTarget: Boolean) {}
diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java
index 9effc67..bbb2f1a 100644
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java
@@ -161,10 +161,10 @@
     void applyCurrentUserOverlays(
             Map<String, OverlayIdentifier> categoryToPackage,
             FabricatedOverlay[] pendingCreation,
-            Set<UserHandle> userHandles) {
+            int currentUser,
+            Set<UserHandle> managedProfiles) {
         // Disable all overlays that have not been specified in the user setting.
         final Set<String> overlayCategoriesToDisable = new HashSet<>(THEME_CATEGORIES);
-        overlayCategoriesToDisable.removeAll(categoryToPackage.keySet());
         final Set<String> targetPackagesToQuery = overlayCategoriesToDisable.stream()
                 .map(category -> mCategoryToTargetPackage.get(category))
                 .collect(Collectors.toSet());
@@ -175,6 +175,7 @@
                 .filter(o ->
                         mTargetPackageToCategories.get(o.targetPackageName).contains(o.category))
                 .filter(o -> overlayCategoriesToDisable.contains(o.category))
+                .filter(o -> !categoryToPackage.containsValue(new OverlayIdentifier(o.packageName)))
                 .filter(o -> o.isEnabled())
                 .map(o -> new Pair<>(o.category, o.packageName))
                 .collect(Collectors.toList());
@@ -186,17 +187,18 @@
             }
         }
 
-        // Toggle overlays in the order of THEME_CATEGORIES.
+        for (Pair<String, String> packageToDisable : overlaysToDisable) {
+            OverlayIdentifier overlayInfo = new OverlayIdentifier(packageToDisable.second);
+            setEnabled(transaction, overlayInfo, packageToDisable.first, currentUser,
+                    managedProfiles, false);
+        }
+
         for (String category : THEME_CATEGORIES) {
             if (categoryToPackage.containsKey(category)) {
                 OverlayIdentifier overlayInfo = categoryToPackage.get(category);
-                setEnabled(transaction, overlayInfo, category, userHandles, true);
+                setEnabled(transaction, overlayInfo, category, currentUser, managedProfiles, true);
             }
         }
-        for (Pair<String, String> packageToDisable : overlaysToDisable) {
-            OverlayIdentifier overlayInfo = new OverlayIdentifier(packageToDisable.second);
-            setEnabled(transaction, overlayInfo, packageToDisable.first, userHandles, false);
-        }
 
         mExecutor.execute(() -> {
             try {
@@ -213,18 +215,30 @@
     }
 
     private void setEnabled(OverlayManagerTransaction.Builder transaction,
-            OverlayIdentifier identifier, String category, Set<UserHandle> handles,
-            boolean enabled) {
+            OverlayIdentifier identifier, String category, int currentUser,
+            Set<UserHandle> managedProfiles, boolean enabled) {
         if (DEBUG) {
             Log.d(TAG, "setEnabled: " + identifier.getPackageName() + " category: "
                     + category + ": " + enabled);
         }
-        for (UserHandle userHandle : handles) {
-            transaction.setEnabled(identifier, enabled, userHandle.getIdentifier());
-        }
-        if (!handles.contains(UserHandle.SYSTEM) && SYSTEM_USER_CATEGORIES.contains(category)) {
+
+        transaction.setEnabled(identifier, enabled, currentUser);
+        if (currentUser != UserHandle.SYSTEM.getIdentifier()
+                && SYSTEM_USER_CATEGORIES.contains(category)) {
             transaction.setEnabled(identifier, enabled, UserHandle.SYSTEM.getIdentifier());
         }
+
+        // Do not apply Launcher or Theme picker overlays to managed users. Apps are not
+        // installed in there.
+        OverlayInfo overlayInfo = mOverlayManager.getOverlayInfo(identifier, UserHandle.SYSTEM);
+        if (overlayInfo == null || overlayInfo.targetPackageName.equals(mLauncherPackage)
+                || overlayInfo.targetPackageName.equals(mThemePickerPackage)) {
+            return;
+        }
+
+        for (UserHandle userHandle : managedProfiles) {
+            transaction.setEnabled(identifier, enabled, userHandle.getIdentifier());
+        }
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
index 028cbd0..5d02845 100644
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
@@ -55,14 +55,13 @@
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.util.settings.SecureSettings;
 
-import com.google.android.collect.Sets;
-
 import org.json.JSONException;
 import org.json.JSONObject;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Executor;
@@ -151,7 +150,7 @@
             @Override
             public void onReceive(Context context, Intent intent) {
                 if (DEBUG) Log.d(TAG, "Updating overlays for user switch / profile added.");
-                updateThemeOverlays();
+                reevaluateSystemTheme(true /* forceReload */);
             }
         }, filter, mBgExecutor, UserHandle.ALL);
         mSecureSettings.registerContentObserverForUser(
@@ -163,7 +162,7 @@
                             int userId) {
                         if (DEBUG) Log.d(TAG, "Overlay changed for user: " + userId);
                         if (ActivityManager.getCurrentUser() == userId) {
-                            updateThemeOverlays();
+                            reevaluateSystemTheme(true /* forceReload */);
                         }
                     }
                 },
@@ -180,7 +179,7 @@
                     mLockColors = lockColors;
                 }
                 mSystemColors = systemColor;
-                reevaluateSystemTheme();
+                reevaluateSystemTheme(false /* forceReload */);
             });
         });
         if (USE_LOCK_SCREEN_WALLPAPER) {
@@ -192,7 +191,7 @@
                     }
                     // It's possible that the user has a lock screen wallpaper. On this case we'll
                     // end up with different colors after unlocking.
-                    reevaluateSystemTheme();
+                    reevaluateSystemTheme(false /* forceReload */);
                 }
             });
         }
@@ -209,11 +208,11 @@
                     Log.d(TAG, "got new lock colors: " + wallpaperColors + " where: " + which);
                 }
             }
-            reevaluateSystemTheme();
+            reevaluateSystemTheme(false /* forceReload */);
         }, null, UserHandle.USER_ALL);
     }
 
-    private void reevaluateSystemTheme() {
+    private void reevaluateSystemTheme(boolean forceReload) {
         WallpaperColors currentColors =
                 mKeyguardStateController.isShowing() && mLockColors != null
                         ? mLockColors : mSystemColors;
@@ -228,7 +227,8 @@
             accentCandidate = getAccentColor(currentColors);
         }
 
-        if (mMainWallpaperColor == mainColor && mWallpaperAccentColor == accentCandidate) {
+        if (mMainWallpaperColor == mainColor && mWallpaperAccentColor == accentCandidate
+                && !forceReload) {
             return;
         }
 
@@ -309,6 +309,16 @@
             } catch (NumberFormatException e) {
                 Log.w(TAG, "Invalid color definition: " + systemPalette.getPackageName());
             }
+        } else if (!mIsMonetEnabled && systemPalette != null) {
+            try {
+                // It's possible that we flipped the flag off and still have a @ColorInt in the
+                // setting. We need to sanitize the input, otherwise the overlay transaction will
+                // fail.
+                Integer.parseInt(systemPalette.getPackageName().toLowerCase(), 16);
+                categoryToPackage.remove(OVERLAY_CATEGORY_SYSTEM_PALETTE);
+            } catch (NumberFormatException e) {
+                // This is a package name. All good, let's continue
+            }
         }
 
         // Same for accent color.
@@ -322,6 +332,13 @@
             } catch (NumberFormatException e) {
                 Log.w(TAG, "Invalid color definition: " + accentPalette.getPackageName());
             }
+        } else if (!mIsMonetEnabled && accentPalette != null) {
+            try {
+                Integer.parseInt(accentPalette.getPackageName().toLowerCase(), 16);
+                categoryToPackage.remove(OVERLAY_CATEGORY_ACCENT_COLOR);
+            } catch (NumberFormatException e) {
+                // This is a package name. All good, let's continue
+            }
         }
 
         // Compatibility with legacy themes, where full packages were defined, instead of just
@@ -337,10 +354,10 @@
             categoryToPackage.put(OVERLAY_CATEGORY_ACCENT_COLOR, mSecondaryOverlay.getIdentifier());
         }
 
-        Set<UserHandle> userHandles = Sets.newHashSet(UserHandle.of(currentUser));
+        Set<UserHandle> managedProfiles = new HashSet<>();
         for (UserInfo userInfo : mUserManager.getEnabledProfiles(currentUser)) {
             if (userInfo.isManagedProfile()) {
-                userHandles.add(userInfo.getUserHandle());
+                managedProfiles.add(userInfo.getUserHandle());
             }
         }
         if (DEBUG) {
@@ -352,9 +369,10 @@
             mNeedsOverlayCreation = false;
             mThemeManager.applyCurrentUserOverlays(categoryToPackage, new FabricatedOverlay[] {
                     mPrimaryOverlay, mSecondaryOverlay, mNeutralOverlay
-            }, userHandles);
+            }, currentUser, managedProfiles);
         } else {
-            mThemeManager.applyCurrentUserOverlays(categoryToPackage, null, userHandles);
+            mThemeManager.applyCurrentUserOverlays(categoryToPackage, null, currentUser,
+                    managedProfiles);
         }
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.java
index 0dfebab..a2a179e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.java
@@ -121,7 +121,6 @@
                 mQSTileHost, mQSCustomizerController, true, mMediaHost,
                 mQSTileRevealControllerFactory, mDumpManager, mMetricsLogger, mUiEventLogger,
                 mQSLogger, mBrightnessControllerFactory, mToggleSliderViewControllerFactory,
-                /* labelsFlag */ false,
                 mFeatureFlags
         );
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelControllerTest.kt
index 5870200..cb380d5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelControllerTest.kt
@@ -87,7 +87,6 @@
                 uiEventLogger,
                 qsLogger,
                 dumpManager,
-                false,
                 featureFlags
         )
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayApplierTest.java b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayApplierTest.java
index e798207..6067b42 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayApplierTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayApplierTest.java
@@ -32,7 +32,6 @@
 import static com.android.systemui.theme.ThemeOverlayApplier.THEME_CATEGORIES;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.mock;
@@ -88,7 +87,9 @@
     private static final String THEMEPICKER_PACKAGE = "com.android.wallpaper";
     private static final String LAUNCHER_PACKAGE = "com.android.launcher3";
     private static final UserHandle TEST_USER = UserHandle.of(5);
-    private static final Set<UserHandle> TEST_USER_HANDLES = Sets.newHashSet(TEST_USER);
+    private static final UserHandle TEST_USER_MANAGED_PROFILE = UserHandle.of(6);
+    private static final Set<UserHandle> TEST_USER_HANDLES =
+            Sets.newHashSet(TEST_USER_MANAGED_PROFILE);
 
     @Mock
     OverlayManager mOverlayManager;
@@ -159,13 +160,19 @@
                                 THEMEPICKER_PACKAGE, OVERLAY_CATEGORY_ICON_THEME_PICKER, false),
                         createOverlayInfo(TEST_ENABLED_PREFIX + OVERLAY_CATEGORY_ICON_THEME_PICKER,
                                 THEMEPICKER_PACKAGE, OVERLAY_CATEGORY_ICON_THEME_PICKER, true)));
+
+        OverlayInfo launcherTargetInfo = new OverlayInfo("packageName", LAUNCHER_PACKAGE,
+                null, null, "/", 0, 0, 0, false);
+        when(mOverlayManager.getOverlayInfo(any(OverlayIdentifier.class), any()))
+                .thenReturn(launcherTargetInfo);
         clearInvocations(mOverlayManager);
         verify(mDumpManager).registerDumpable(any(), any());
     }
 
     @Test
     public void allCategoriesSpecified_allEnabledExclusively() {
-        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, null, TEST_USER_HANDLES);
+        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, null, TEST_USER.getIdentifier(),
+                TEST_USER_HANDLES);
         verify(mOverlayManager).commit(any());
 
         for (OverlayIdentifier overlayPackage : ALL_CATEGORIES_MAP.values()) {
@@ -176,7 +183,8 @@
 
     @Test
     public void allCategoriesSpecified_sysuiCategoriesAlsoAppliedToSysuiUser() {
-        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, null, TEST_USER_HANDLES);
+        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, null, TEST_USER.getIdentifier(),
+                TEST_USER_HANDLES);
 
         for (Map.Entry<String, OverlayIdentifier> entry : ALL_CATEGORIES_MAP.entrySet()) {
             if (SYSTEM_USER_CATEGORIES.contains(entry.getKey())) {
@@ -192,27 +200,25 @@
     @Test
     public void allCategoriesSpecified_enabledForAllUserHandles() {
         Set<UserHandle> userHandles = Sets.newHashSet(TEST_USER_HANDLES);
-        UserHandle newUserHandle = UserHandle.of(10);
-        userHandles.add(newUserHandle);
-        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, null, userHandles);
+        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, null, TEST_USER.getIdentifier(),
+                userHandles);
 
         for (OverlayIdentifier overlayPackage : ALL_CATEGORIES_MAP.values()) {
             verify(mTransactionBuilder).setEnabled(eq(overlayPackage), eq(true),
                     eq(TEST_USER.getIdentifier()));
-            verify(mTransactionBuilder).setEnabled(eq(overlayPackage), eq(true),
-                    eq(newUserHandle.getIdentifier()));
+            // Not enabled for work profile because the target package is LAUNCHER_PACKAGE
+            verify(mTransactionBuilder, never()).setEnabled(eq(overlayPackage), eq(true),
+                    eq(TEST_USER_MANAGED_PROFILE.getIdentifier()));
         }
     }
 
     @Test
     public void applyCurrentUserOverlays_createsPendingOverlays() {
-        Set<UserHandle> userHandles = Sets.newHashSet(TEST_USER_HANDLES);
-        UserHandle newUserHandle = UserHandle.of(10);
-        userHandles.add(newUserHandle);
-        FabricatedOverlay[] pendingCreation = new FabricatedOverlay[] {
+        FabricatedOverlay[] pendingCreation = new FabricatedOverlay[]{
                 mock(FabricatedOverlay.class)
         };
-        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, pendingCreation, userHandles);
+        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, pendingCreation,
+                TEST_USER.getIdentifier(), TEST_USER_HANDLES);
 
         for (FabricatedOverlay overlay : pendingCreation) {
             verify(mTransactionBuilder).registerFabricatedOverlay(eq(overlay));
@@ -220,20 +226,13 @@
     }
 
     @Test
-    public void allCategoriesSpecified_overlayManagerNotQueried() {
-        mManager.applyCurrentUserOverlays(ALL_CATEGORIES_MAP, null, TEST_USER_HANDLES);
-
-        verify(mOverlayManager, never())
-                .getOverlayInfosForTarget(anyString(), any(UserHandle.class));
-    }
-
-    @Test
     public void someCategoriesSpecified_specifiedEnabled_unspecifiedDisabled() {
         Map<String, OverlayIdentifier> categoryToPackage = new HashMap<>(ALL_CATEGORIES_MAP);
         categoryToPackage.remove(OVERLAY_CATEGORY_ICON_SETTINGS);
         categoryToPackage.remove(OVERLAY_CATEGORY_ICON_ANDROID);
 
-        mManager.applyCurrentUserOverlays(categoryToPackage, null, TEST_USER_HANDLES);
+        mManager.applyCurrentUserOverlays(categoryToPackage, null, TEST_USER.getIdentifier(),
+                TEST_USER_HANDLES);
 
         for (OverlayIdentifier overlayPackage : categoryToPackage.values()) {
             verify(mTransactionBuilder).setEnabled(eq(overlayPackage), eq(true),
@@ -249,7 +248,8 @@
 
     @Test
     public void zeroCategoriesSpecified_allDisabled() {
-        mManager.applyCurrentUserOverlays(Maps.newArrayMap(), null, TEST_USER_HANDLES);
+        mManager.applyCurrentUserOverlays(Maps.newArrayMap(), null, TEST_USER.getIdentifier(),
+                TEST_USER_HANDLES);
 
         for (String category : THEME_CATEGORIES) {
             verify(mTransactionBuilder).setEnabled(
@@ -263,7 +263,8 @@
         Map<String, OverlayIdentifier> categoryToPackage = new HashMap<>(ALL_CATEGORIES_MAP);
         categoryToPackage.put("blah.category", new OverlayIdentifier("com.example.blah.category"));
 
-        mManager.applyCurrentUserOverlays(categoryToPackage, null, TEST_USER_HANDLES);
+        mManager.applyCurrentUserOverlays(categoryToPackage, null, TEST_USER.getIdentifier(),
+                TEST_USER_HANDLES);
 
         verify(mTransactionBuilder, never()).setEnabled(
                 eq(new OverlayIdentifier("com.example.blah.category")), eq(false),
@@ -273,23 +274,6 @@
                 eq(TEST_USER.getIdentifier()));
     }
 
-    @Test
-    public void overlayManagerOnlyQueriedForUnspecifiedPackages() {
-        Map<String, OverlayIdentifier> categoryToPackage = new HashMap<>(ALL_CATEGORIES_MAP);
-        categoryToPackage.remove(OVERLAY_CATEGORY_ICON_SETTINGS);
-
-        mManager.applyCurrentUserOverlays(categoryToPackage, null, TEST_USER_HANDLES);
-
-        verify(mOverlayManager).getOverlayInfosForTarget(SETTINGS_PACKAGE, UserHandle.SYSTEM);
-        verify(mOverlayManager, never()).getOverlayInfosForTarget(ANDROID_PACKAGE,
-                UserHandle.SYSTEM);
-        verify(mOverlayManager, never()).getOverlayInfosForTarget(SYSUI_PACKAGE, UserHandle.SYSTEM);
-        verify(mOverlayManager, never()).getOverlayInfosForTarget(LAUNCHER_PACKAGE,
-                UserHandle.SYSTEM);
-        verify(mOverlayManager, never()).getOverlayInfosForTarget(THEMEPICKER_PACKAGE,
-                UserHandle.SYSTEM);
-    }
-
     private static OverlayInfo createOverlayInfo(String packageName, String targetPackageName,
             String category, boolean enabled) {
         return new OverlayInfo(packageName, null, targetPackageName, null, category, "",
diff --git a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
index aa385ef..d80c40f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
@@ -143,7 +143,7 @@
                 ArgumentCaptor.forClass(Map.class);
 
         verify(mThemeOverlayApplier)
-                .applyCurrentUserOverlays(themeOverlays.capture(), any(), any());
+                .applyCurrentUserOverlays(themeOverlays.capture(), any(), anyInt(), any());
 
         // Assert that we received the colors that we were expecting
         assertThat(themeOverlays.getValue().get(OVERLAY_CATEGORY_SYSTEM_PALETTE))
@@ -175,7 +175,7 @@
                 ArgumentCaptor.forClass(Map.class);
 
         verify(mThemeOverlayApplier)
-                .applyCurrentUserOverlays(themeOverlays.capture(), any(), any());
+                .applyCurrentUserOverlays(themeOverlays.capture(), any(), anyInt(), any());
 
         // Assert that we received the colors that we were expecting
         assertThat(themeOverlays.getValue().get(OVERLAY_CATEGORY_SYSTEM_PALETTE))
@@ -198,7 +198,7 @@
                 ArgumentCaptor.forClass(Map.class);
 
         verify(mThemeOverlayApplier)
-                .applyCurrentUserOverlays(themeOverlays.capture(), any(), any());
+                .applyCurrentUserOverlays(themeOverlays.capture(), any(), anyInt(), any());
 
         // Assert that we received the colors that we were expecting
         assertThat(themeOverlays.getValue().get(OVERLAY_CATEGORY_SYSTEM_PALETTE))
diff --git a/services/core/java/com/android/server/powerstats/PowerStatsLogger.java b/services/core/java/com/android/server/powerstats/PowerStatsLogger.java
index c4f29ea..ef0079e 100644
--- a/services/core/java/com/android/server/powerstats/PowerStatsLogger.java
+++ b/services/core/java/com/android/server/powerstats/PowerStatsLogger.java
@@ -16,6 +16,8 @@
 
 package com.android.server.powerstats;
 
+import static java.lang.System.currentTimeMillis;
+
 import android.content.Context;
 import android.hardware.power.stats.Channel;
 import android.hardware.power.stats.EnergyConsumer;
@@ -26,11 +28,14 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
+import android.os.SystemClock;
 import android.util.AtomicFile;
 import android.util.Slog;
 import android.util.proto.ProtoInputStream;
 import android.util.proto.ProtoOutputStream;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import com.android.server.powerstats.PowerStatsHALWrapper.IPowerStatsHALWrapper;
 import com.android.server.powerstats.ProtoStreamUtils.ChannelUtils;
 import com.android.server.powerstats.ProtoStreamUtils.EnergyConsumerResultUtils;
@@ -61,6 +66,8 @@
     protected static final int MSG_LOG_TO_DATA_STORAGE_LOW_FREQUENCY = 1;
     protected static final int MSG_LOG_TO_DATA_STORAGE_HIGH_FREQUENCY = 2;
 
+    // TODO(b/181240441): Add a listener to update the Wall clock baseline when changed
+    private final long mStartWallTime;
     private final PowerStatsDataStorage mPowerStatsMeterStorage;
     private final PowerStatsDataStorage mPowerStatsModelStorage;
     private final PowerStatsDataStorage mPowerStatsResidencyStorage;
@@ -79,6 +86,8 @@
                 // Log power meter data.
                 EnergyMeasurement[] energyMeasurements =
                     mPowerStatsHALWrapper.readEnergyMeter(new int[0]);
+                EnergyMeasurementUtils.adjustTimeSinceBootToEpoch(energyMeasurements,
+                        mStartWallTime);
                 mPowerStatsMeterStorage.write(
                         EnergyMeasurementUtils.getProtoBytes(energyMeasurements));
                 if (DEBUG) EnergyMeasurementUtils.print(energyMeasurements);
@@ -86,6 +95,8 @@
                 // Log power model data without attribution data.
                 EnergyConsumerResult[] ecrNoAttribution =
                     mPowerStatsHALWrapper.getEnergyConsumed(new int[0]);
+                EnergyConsumerResultUtils.adjustTimeSinceBootToEpoch(ecrNoAttribution,
+                        mStartWallTime);
                 mPowerStatsModelStorage.write(
                         EnergyConsumerResultUtils.getProtoBytes(ecrNoAttribution, false));
                 if (DEBUG) EnergyConsumerResultUtils.print(ecrNoAttribution);
@@ -97,6 +108,8 @@
                 // Log power model data with attribution data.
                 EnergyConsumerResult[] ecrAttribution =
                     mPowerStatsHALWrapper.getEnergyConsumed(new int[0]);
+                EnergyConsumerResultUtils.adjustTimeSinceBootToEpoch(ecrAttribution,
+                        mStartWallTime);
                 mPowerStatsModelStorage.write(
                         EnergyConsumerResultUtils.getProtoBytes(ecrAttribution, true));
                 if (DEBUG) EnergyConsumerResultUtils.print(ecrAttribution);
@@ -108,6 +121,8 @@
                 // Log state residency data.
                 StateResidencyResult[] stateResidencyResults =
                     mPowerStatsHALWrapper.getStateResidency(new int[0]);
+                StateResidencyResultUtils.adjustTimeSinceBootToEpoch(stateResidencyResults,
+                        mStartWallTime);
                 mPowerStatsResidencyStorage.write(
                         StateResidencyResultUtils.getProtoBytes(stateResidencyResults));
                 if (DEBUG) StateResidencyResultUtils.print(stateResidencyResults);
@@ -293,12 +308,19 @@
         return mDeleteResidencyDataOnBoot;
     }
 
+    @VisibleForTesting
+    public long getStartWallTime() {
+        return mStartWallTime;
+    }
+
     public PowerStatsLogger(Context context, File dataStoragePath,
             String meterFilename, String meterCacheFilename,
             String modelFilename, String modelCacheFilename,
             String residencyFilename, String residencyCacheFilename,
             IPowerStatsHALWrapper powerStatsHALWrapper) {
         super(Looper.getMainLooper());
+        mStartWallTime = currentTimeMillis() - SystemClock.elapsedRealtime();
+        if (DEBUG) Slog.d(TAG, "mStartWallTime: " + mStartWallTime);
         mPowerStatsHALWrapper = powerStatsHALWrapper;
         mDataStoragePath = dataStoragePath;
 
diff --git a/services/core/java/com/android/server/powerstats/ProtoStreamUtils.java b/services/core/java/com/android/server/powerstats/ProtoStreamUtils.java
index 11b22a5..746a098 100644
--- a/services/core/java/com/android/server/powerstats/ProtoStreamUtils.java
+++ b/services/core/java/com/android/server/powerstats/ProtoStreamUtils.java
@@ -109,6 +109,18 @@
     }
 
     static class StateResidencyResultUtils {
+        public static void adjustTimeSinceBootToEpoch(StateResidencyResult[] stateResidencyResult,
+                long startWallTime) {
+            for (int i = 0; i < stateResidencyResult.length; i++) {
+                final int stateLength = stateResidencyResult[i].stateResidencyData.length;
+                for (int j = 0; j < stateLength; j++) {
+                    final StateResidency stateResidencyData =
+                            stateResidencyResult[i].stateResidencyData[j];
+                    stateResidencyData.lastEntryTimestampMs += startWallTime;
+                }
+            }
+        }
+
         public static byte[] getProtoBytes(StateResidencyResult[] stateResidencyResult) {
             ProtoOutputStream pos = new ProtoOutputStream();
             packProtoMessage(stateResidencyResult, pos);
@@ -306,6 +318,13 @@
     }
 
     static class EnergyMeasurementUtils {
+        public static void adjustTimeSinceBootToEpoch(EnergyMeasurement[] energyMeasurement,
+                long startWallTime) {
+            for (int i = 0; i < energyMeasurement.length; i++) {
+                energyMeasurement[i].timestampMs += startWallTime;
+            }
+        }
+
         public static byte[] getProtoBytes(EnergyMeasurement[] energyMeasurement) {
             ProtoOutputStream pos = new ProtoOutputStream();
             packProtoMessage(energyMeasurement, pos);
@@ -518,6 +537,13 @@
     }
 
     static class EnergyConsumerResultUtils {
+        public static void adjustTimeSinceBootToEpoch(EnergyConsumerResult[] energyConsumerResult,
+                long startWallTime) {
+            for (int i = 0; i < energyConsumerResult.length; i++) {
+                energyConsumerResult[i].timestampMs += startWallTime;
+            }
+        }
+
         public static byte[] getProtoBytes(EnergyConsumerResult[] energyConsumerResult,
                 boolean includeAttribution) {
             ProtoOutputStream pos = new ProtoOutputStream();
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 71d2b2f..cca85b2 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -2393,15 +2393,9 @@
                 }
             }
 
-            // We may be deferring layout passes at the moment, but since the client is interested
-            // in the new out values right now we need to force a layout.
-            mWindowPlacerLocked.performSurfacePlacement(true /* force */);
-
+            // Create surfaceControl before surface placement otherwise layout will be skipped
+            // (because WS.isGoneForLayout() is true when there is no surface.
             if (shouldRelayout) {
-                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "relayoutWindow: viewVisibility_1");
-
-                result = win.relayoutVisibleWindow(result);
-
                 try {
                     result = createSurfaceControl(outSurfaceControl, result, win, winAnimator);
                 } catch (Exception e) {
@@ -2413,6 +2407,17 @@
                     Binder.restoreCallingIdentity(origId);
                     return 0;
                 }
+            }
+
+            // We may be deferring layout passes at the moment, but since the client is interested
+            // in the new out values right now we need to force a layout.
+            mWindowPlacerLocked.performSurfacePlacement(true /* force */);
+
+            if (shouldRelayout) {
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "relayoutWindow: viewVisibility_1");
+
+                result = win.relayoutVisibleWindow(result);
+
                 if ((result & WindowManagerGlobal.RELAYOUT_RES_FIRST_TIME) != 0) {
                     focusMayChange = true;
                 }
diff --git a/services/tests/servicestests/src/com/android/server/powerstats/PowerStatsServiceTest.java b/services/tests/servicestests/src/com/android/server/powerstats/PowerStatsServiceTest.java
index ddbe81c..26b34fd 100644
--- a/services/tests/servicestests/src/com/android/server/powerstats/PowerStatsServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/powerstats/PowerStatsServiceTest.java
@@ -318,7 +318,8 @@
         assertTrue(pssProto.energyMeasurement.length == ENERGY_METER_COUNT);
         for (int i = 0; i < pssProto.energyMeasurement.length; i++) {
             assertTrue(pssProto.energyMeasurement[i].id == i);
-            assertTrue(pssProto.energyMeasurement[i].timestampMs == i);
+            assertTrue(pssProto.energyMeasurement[i].timestampMs ==
+                    i + mPowerStatsLogger.getStartWallTime());
             assertTrue(pssProto.energyMeasurement[i].durationMs == i);
             assertTrue(pssProto.energyMeasurement[i].energyUws == i);
         }
@@ -359,7 +360,8 @@
         assertTrue(pssProto.energyConsumerResult.length == ENERGY_CONSUMER_COUNT);
         for (int i = 0; i < pssProto.energyConsumerResult.length; i++) {
             assertTrue(pssProto.energyConsumerResult[i].id == i);
-            assertTrue(pssProto.energyConsumerResult[i].timestampMs == i);
+            assertTrue(pssProto.energyConsumerResult[i].timestampMs ==
+                    i + mPowerStatsLogger.getStartWallTime());
             assertTrue(pssProto.energyConsumerResult[i].energyUws == i);
             assertTrue(pssProto.energyConsumerResult[i].attribution.length
                     == ENERGY_CONSUMER_ATTRIBUTION_COUNT);
@@ -420,7 +422,8 @@
                 assertTrue(stateResidency.id == j);
                 assertTrue(stateResidency.totalTimeInStateMs == j);
                 assertTrue(stateResidency.totalStateEntryCount == j);
-                assertTrue(stateResidency.lastEntryTimestampMs == j);
+                assertTrue(stateResidency.lastEntryTimestampMs ==
+                        j + mPowerStatsLogger.getStartWallTime());
             }
         }
     }