Wifi: Fix Wifi Multicast Wakelock Statistics

In current implementation, when the same UID holds multiple WiFi
wakelocks, the timer for the lock acquisition in batterystats is started
with the first wakelock acquire and stops with first release.
So the nesting of wakelocks of the same Uid is not correctly accounted for.
(so if Lock#1 acquired at T1, Lock# 2 acquired at T2, Lock#2 released at T3,
lock #1 released at T4 , the timer will only consider the interval T3 - T1 which
is not the correct value).

This commit fixes this issue by starting the timer on the first
acquisition, and stopping it on the last release

Bug: 70691435
Test: Unit Test
Test: $ atest BatteryStatsImplTest

Change-Id: Ie955b6ab0eb1e9837fdda2b8d81835038b8d71da
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index c65dd6f..07f74f0 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -6724,7 +6724,7 @@
         int mWifiBatchedScanBinStarted = NO_BATCHED_SCAN_STARTED;
         StopwatchTimer[] mWifiBatchedScanTimer;
 
-        boolean mWifiMulticastEnabled;
+        int mWifiMulticastWakelockCount;
         StopwatchTimer mWifiMulticastTimer;
 
         StopwatchTimer mAudioTurnedOnTimer;
@@ -7183,8 +7183,7 @@
 
         @Override
         public void noteWifiMulticastEnabledLocked(long elapsedRealtimeMs) {
-            if (!mWifiMulticastEnabled) {
-                mWifiMulticastEnabled = true;
+            if (mWifiMulticastWakelockCount == 0) {
                 if (mWifiMulticastTimer == null) {
                     mWifiMulticastTimer = new StopwatchTimer(mBsi.mClocks, Uid.this,
                             WIFI_MULTICAST_ENABLED, mBsi.mWifiMulticastTimers, mBsi.mOnBatteryTimeBase);
@@ -7194,12 +7193,17 @@
                         StatsLog.WIFI_MULTICAST_LOCK_STATE_CHANGED, getUid(), null,
                         StatsLog.WIFI_MULTICAST_LOCK_STATE_CHANGED__STATE__ON);
             }
+            mWifiMulticastWakelockCount++;
         }
 
         @Override
         public void noteWifiMulticastDisabledLocked(long elapsedRealtimeMs) {
-            if (mWifiMulticastEnabled) {
-                mWifiMulticastEnabled = false;
+            if (mWifiMulticastWakelockCount == 0) {
+                return;
+            }
+
+            mWifiMulticastWakelockCount--;
+            if (mWifiMulticastWakelockCount == 0) {
                 mWifiMulticastTimer.stopRunningLocked(elapsedRealtimeMs);
                 StatsLog.write_non_chained(
                         StatsLog.WIFI_MULTICAST_LOCK_STATE_CHANGED, getUid(), null,
@@ -7932,7 +7936,7 @@
             }
             if (mWifiMulticastTimer != null) {
                 active |= !mWifiMulticastTimer.reset(false);
-                active |= mWifiMulticastEnabled;
+                active |= (mWifiMulticastWakelockCount > 0);
             }
 
             active |= !resetIfNotNull(mAudioTurnedOnTimer, false);
@@ -8590,7 +8594,7 @@
                     mWifiBatchedScanTimer[i] = null;
                 }
             }
-            mWifiMulticastEnabled = false;
+            mWifiMulticastWakelockCount = 0;
             if (in.readInt() != 0) {
                 mWifiMulticastTimer = new StopwatchTimer(mBsi.mClocks, Uid.this, WIFI_MULTICAST_ENABLED,
                         mBsi.mWifiMulticastTimers, mBsi.mOnBatteryTimeBase, in);
@@ -14031,7 +14035,7 @@
                     u.mWifiBatchedScanTimer[i].readSummaryFromParcelLocked(in);
                 }
             }
-            u.mWifiMulticastEnabled = false;
+            u.mWifiMulticastWakelockCount = 0;
             if (in.readInt() != 0) {
                 u.mWifiMulticastTimer.readSummaryFromParcelLocked(in);
             }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
index a8094ea..613de45 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsImplTest.java
@@ -24,6 +24,7 @@
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
@@ -302,6 +303,113 @@
         assertArrayEquals(expected, mBatteryStatsImpl.addCpuTimes(timesA, timesB));
     }
 
+    @Test
+    public void testMulticastWakelockAcqRel() {
+        final int testUid = 10032;
+        final int acquireTimeMs = 1000;
+        final int releaseTimeMs = 1005;
+        final int currentTimeMs = 1011;
+
+        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+
+        // Create a Uid Object
+        final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
+        assertNotNull(u);
+
+        // Acquire and release the lock
+        u.noteWifiMulticastEnabledLocked(acquireTimeMs);
+        u.noteWifiMulticastDisabledLocked(releaseTimeMs);
+
+        // Get the total acquisition time
+        long totalTime = u.getWifiMulticastTime(currentTimeMs*1000,
+                BatteryStats.STATS_SINCE_UNPLUGGED);
+        assertEquals("Miscalculations of Multicast wakelock acquisition time",
+                (releaseTimeMs - acquireTimeMs) * 1000, totalTime);
+    }
+
+    @Test
+    public void testMulticastWakelockAcqNoRel() {
+        final int testUid = 10032;
+        final int acquireTimeMs = 1000;
+        final int currentTimeMs = 1011;
+
+        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+
+        // Create a Uid Object
+        final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
+        assertNotNull(u);
+
+        // Acquire the lock
+        u.noteWifiMulticastEnabledLocked(acquireTimeMs);
+
+        // Get the total acquisition time
+        long totalTime =  u.getWifiMulticastTime(currentTimeMs*1000,
+                BatteryStats.STATS_SINCE_UNPLUGGED);
+        assertEquals("Miscalculations of Multicast wakelock acquisition time",
+                (currentTimeMs - acquireTimeMs) * 1000, totalTime);
+    }
+
+    @Test
+    public void testMulticastWakelockAcqAcqRelRel() {
+        final int testUid = 10032;
+        final int acquireTimeMs_1 = 1000;
+        final int acquireTimeMs_2 = 1002;
+
+        final int releaseTimeMs_1 = 1005;
+        final int releaseTimeMs_2 = 1009;
+        final int currentTimeMs = 1011;
+
+        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+
+        // Create a Uid Object
+        final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
+        assertNotNull(u);
+
+        // Acquire and release the lock (twice in nested way)
+        u.noteWifiMulticastEnabledLocked(acquireTimeMs_1);
+        u.noteWifiMulticastEnabledLocked(acquireTimeMs_2);
+
+        u.noteWifiMulticastDisabledLocked(releaseTimeMs_1);
+        u.noteWifiMulticastDisabledLocked(releaseTimeMs_2);
+
+        // Get the total acquisition time
+        long totalTime =  u.getWifiMulticastTime(currentTimeMs*1000,
+                BatteryStats.STATS_SINCE_UNPLUGGED);
+        assertEquals("Miscalculations of Multicast wakelock acquisition time",
+                (releaseTimeMs_2 - acquireTimeMs_1) * 1000, totalTime);
+    }
+
+    @Test
+    public void testMulticastWakelockAcqRelAcqRel() {
+        final int testUid = 10032;
+        final int acquireTimeMs_1 = 1000;
+        final int acquireTimeMs_2 = 1005;
+
+        final int releaseTimeMs_1 = 1002;
+        final int releaseTimeMs_2 = 1009;
+        final int currentTimeMs = 1011;
+
+        mBatteryStatsImpl.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+
+        // Create a Uid Object
+        final BatteryStats.Uid u = mBatteryStatsImpl.getUidStatsLocked(testUid);
+        assertNotNull(u);
+
+        // Acquire and release the lock (twice)
+        u.noteWifiMulticastEnabledLocked(acquireTimeMs_1);
+        u.noteWifiMulticastDisabledLocked(releaseTimeMs_1);
+
+        u.noteWifiMulticastEnabledLocked(acquireTimeMs_2);
+        u.noteWifiMulticastDisabledLocked(releaseTimeMs_2);
+
+        // Get the total acquisition time
+        long totalTime =  u.getWifiMulticastTime(currentTimeMs*1000,
+                BatteryStats.STATS_SINCE_UNPLUGGED);
+        assertEquals("Miscalculations of Multicast wakelock acquisition time",
+                ((releaseTimeMs_1 - acquireTimeMs_1) + (releaseTimeMs_2 - acquireTimeMs_2))
+                * 1000, totalTime);
+    }
+
     private void addIsolatedUid(int parentUid, int childUid) {
         final BatteryStatsImpl.Uid u = mBatteryStatsImpl.getUidStatsLocked(parentUid);
         u.addIsolatedUid(childUid);