Make ConnectivityService update BatteryStats radio power state

Currently, BatteryStatsService register netd callback and track radio
power state only for the default network.
So, battery stats are not accurate if cellular and Wi-Fi networks are
active at the same time.

aosp/2561090 updates BatteryStatsService not to register netd callback
on V+ devices.
This CL updates LegacyNetworkActivityTracker to call BatteryStats API
and update radio power state for networks including non default networks
on V+ devices.

Bug: 267870186
Bug: 279380356
Test: atest FrameworksNetTests
Change-Id: I826b3b4046cc7b8aef46c13300854eaf14b1b777
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index a9af7b4..8d8117b 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -11802,6 +11802,10 @@
         // Key is netId. Value is configured idle timer information.
         private final SparseArray<IdleTimerParams> mActiveIdleTimers = new SparseArray<>();
         private final boolean mTrackMultiNetworkActivities;
+        // Store netIds of Wi-Fi networks whose idletimers report that they are active
+        private final Set<Integer> mActiveWifiNetworks = new ArraySet<>();
+        // Store netIds of cellular networks whose idletimers report that they are active
+        private final Set<Integer> mActiveCellularNetworks = new ArraySet<>();
 
         private static class IdleTimerParams {
             public final int timeout;
@@ -11828,6 +11832,40 @@
             }
         }
 
+        /**
+         * Update network activity and call BatteryStats to update radio power state if the
+         * mobile or Wi-Fi activity is changed.
+         * LegacyNetworkActivityTracker considers the mobile network is active if at least one
+         * mobile network is active since BatteryStatsService only maintains a single power state
+         * for the mobile network.
+         * The Wi-Fi network is also the same.
+         *
+         * {@link #setupDataActivityTracking} and {@link #removeDataActivityTracking} use
+         * TRANSPORT_CELLULAR as the transportType argument if the network has both cell and Wi-Fi
+         * transports.
+         */
+        private void maybeUpdateRadioPowerState(final int netId, final int transportType,
+                final boolean isActive, final int uid) {
+            if (transportType != TRANSPORT_WIFI && transportType != TRANSPORT_CELLULAR) {
+                Log.e(TAG, "Unexpected transportType in maybeUpdateRadioPowerState: "
+                        + transportType);
+                return;
+            }
+            final Set<Integer> activeNetworks = transportType == TRANSPORT_WIFI
+                    ? mActiveWifiNetworks : mActiveCellularNetworks;
+
+            final boolean wasEmpty = activeNetworks.isEmpty();
+            if (isActive) {
+                activeNetworks.add(netId);
+            } else {
+                activeNetworks.remove(netId);
+            }
+
+            if (wasEmpty != activeNetworks.isEmpty()) {
+                updateRadioPowerState(isActive, transportType, uid);
+            }
+        }
+
         private void handleDefaultNetworkActivity(final int transportType,
                 final boolean isActive, final long timestampNs) {
             mIsDefaultNetworkActive = isActive;
@@ -11840,12 +11878,8 @@
 
         private void handleReportNetworkActivityWithNetIdLabel(
                 NetworkActivityParams activityParams) {
-            if (mDefaultNetwork == null || mDefaultNetwork.netId != activityParams.label) {
-                // This activity change is not for the default network.
-                return;
-            }
-
-            final IdleTimerParams idleTimerParams = mActiveIdleTimers.get(activityParams.label);
+            final int netId = activityParams.label;
+            final IdleTimerParams idleTimerParams = mActiveIdleTimers.get(netId);
             if (idleTimerParams == null) {
                 // This network activity change is not tracked anymore
                 // This can happen if netd callback post activity change event message but idle
@@ -11855,7 +11889,16 @@
             // TODO: if a network changes transports, storing the transport type in the
             // IdleTimerParams is not correct. Consider getting it from the network's
             // NetworkCapabilities instead.
-            handleDefaultNetworkActivity(idleTimerParams.transportType, activityParams.isActive,
+            final int transportType = idleTimerParams.transportType;
+            maybeUpdateRadioPowerState(netId, transportType,
+                    activityParams.isActive, activityParams.uid);
+
+            if (mDefaultNetwork == null || mDefaultNetwork.netId != netId) {
+                // This activity change is not for the default network.
+                return;
+            }
+
+            handleDefaultNetworkActivity(transportType, activityParams.isActive,
                     activityParams.timestampNs);
         }
 
@@ -11944,6 +11987,21 @@
             return trackMultiNetworkActivities ? netId : transportType;
         }
 
+        private boolean maybeCreateIdleTimer(
+                String iface, int netId, int timeout, int transportType) {
+            if (timeout <= 0 || iface == null) return false;
+            try {
+                final String label = Integer.toString(getIdleTimerLabel(
+                        mTrackMultiNetworkActivities, netId, transportType));
+                mNetd.idletimerAddInterface(iface, timeout, label);
+                mActiveIdleTimers.put(netId, new IdleTimerParams(timeout, transportType));
+                return true;
+            } catch (Exception e) {
+                loge("Exception in createIdleTimer", e);
+                return false;
+            }
+        }
+
         /**
          * Setup data activity tracking for the given network.
          *
@@ -11979,21 +12037,14 @@
                 return false; // do not track any other networks
             }
 
-            updateRadioPowerState(true /* isActive */, type);
-
-            if (timeout > 0 && iface != null) {
-                try {
-                    mActiveIdleTimers.put(netId, new IdleTimerParams(timeout, type));
-                    final String label = Integer.toString(getIdleTimerLabel(
-                            mTrackMultiNetworkActivities, netId, type));
-                    mNetd.idletimerAddInterface(iface, timeout, label);
-                    return true;
-                } catch (Exception e) {
-                    // You shall not crash!
-                    loge("Exception in setupDataActivityTracking " + e);
-                }
+            final boolean hasIdleTimer = maybeCreateIdleTimer(iface, netId, timeout, type);
+            if (hasIdleTimer || !mTrackMultiNetworkActivities) {
+                // If trackMultiNetwork is disabled, NetworkActivityTracker updates radio power
+                // state in all cases. If trackMultiNetwork is enabled, it updates radio power
+                // state only about a network that has an idletimer.
+                maybeUpdateRadioPowerState(netId, type, true /* isActive */, NO_UID);
             }
-            return false;
+            return hasIdleTimer;
         }
 
         /**
@@ -12020,7 +12071,7 @@
             }
 
             try {
-                updateRadioPowerState(false /* isActive */, type);
+                maybeUpdateRadioPowerState(netId, type, false /* isActive */, NO_UID);
                 final IdleTimerParams params = mActiveIdleTimers.get(netId);
                 if (params == null) {
                     // IdleTimer is not added if the configured timeout is 0 or negative value
@@ -12076,14 +12127,14 @@
             }
         }
 
-        private void updateRadioPowerState(boolean isActive, int transportType) {
+        private void updateRadioPowerState(boolean isActive, int transportType, int uid) {
             final BatteryStatsManager bs = mContext.getSystemService(BatteryStatsManager.class);
             switch (transportType) {
                 case NetworkCapabilities.TRANSPORT_CELLULAR:
-                    bs.reportMobileRadioPowerState(isActive, NO_UID);
+                    bs.reportMobileRadioPowerState(isActive, uid);
                     break;
                 case NetworkCapabilities.TRANSPORT_WIFI:
-                    bs.reportWifiRadioPowerState(isActive, NO_UID);
+                    bs.reportWifiRadioPowerState(isActive, uid);
                     break;
                 default:
                     logw("Untracked transport type:" + transportType);
@@ -12114,11 +12165,13 @@
                     pw.print("    timeout="); pw.print(params.timeout);
                     pw.print(" type="); pw.println(params.transportType);
                 }
+                pw.println("WiFi active networks: " + mActiveWifiNetworks);
+                pw.println("Cellular active networks: " + mActiveCellularNetworks);
             } catch (Exception e) {
-                // mActiveIdleTimers should only be accessed from handler thread, except dump().
-                // As dump() is never called in normal usage, it would be needlessly expensive
-                // to lock the collection only for its benefit.
-                // Also, mActiveIdleTimers is not expected to be updated frequently.
+                // mActiveIdleTimers, mActiveWifiNetworks, and mActiveCellularNetworks should only
+                // be accessed from handler thread, except dump(). As dump() is never called in
+                // normal usage, it would be needlessly expensive to lock the collection only for
+                // its benefit. Also, they are not expected to be updated frequently.
                 // So catching the exception and logging.
                 pw.println("Failed to dump NetworkActivityTracker: " + e);
             }
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 67aac9f..fbaa046 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -154,6 +154,8 @@
 import static android.net.resolv.aidl.IDnsResolverUnsolicitedEventListener.VALIDATION_RESULT_SUCCESS;
 import static android.os.Process.INVALID_UID;
 import static android.system.OsConstants.IPPROTO_TCP;
+import static android.telephony.DataConnectionRealTimeInfo.DC_POWER_STATE_HIGH;
+import static android.telephony.DataConnectionRealTimeInfo.DC_POWER_STATE_LOW;
 
 import static com.android.server.ConnectivityService.DELAY_DESTROY_FROZEN_SOCKETS_VERSION;
 import static com.android.net.module.util.DeviceConfigUtils.TETHERING_MODULE_NAME;
@@ -638,8 +640,8 @@
 
     // BatteryStatsManager is final and cannot be mocked with regular mockito, so just mock the
     // underlying binder calls.
-    final BatteryStatsManager mBatteryStatsManager =
-            new BatteryStatsManager(mock(IBatteryStats.class));
+    final IBatteryStats mIBatteryStats = mock(IBatteryStats.class);
+    final BatteryStatsManager mBatteryStatsManager = new BatteryStatsManager(mIBatteryStats);
 
     private ArgumentCaptor<ResolverParamsParcel> mResolverParamsParcelCaptor =
             ArgumentCaptor.forClass(ResolverParamsParcel.class);
@@ -11355,8 +11357,20 @@
         final ConditionVariable onNetworkActiveCv = new ConditionVariable();
         final ConnectivityManager.OnNetworkActiveListener listener = onNetworkActiveCv::open;
 
+        TestNetworkCallback defaultCallback = new TestNetworkCallback();
+
         testAndCleanup(() -> {
+            mCm.registerDefaultNetworkCallback(defaultCallback);
             agent.connect(true);
+            defaultCallback.expectAvailableThenValidatedCallbacks(agent);
+            if (transportType == TRANSPORT_CELLULAR) {
+                verify(mIBatteryStats).noteMobileRadioPowerState(eq(DC_POWER_STATE_HIGH),
+                        anyLong() /* timestampNs */, eq(NETWORK_ACTIVITY_NO_UID));
+            } else if (transportType == TRANSPORT_WIFI) {
+                verify(mIBatteryStats).noteWifiRadioPowerState(eq(DC_POWER_STATE_HIGH),
+                        anyLong() /* timestampNs */, eq(NETWORK_ACTIVITY_NO_UID));
+            }
+            clearInvocations(mIBatteryStats);
             final int idleTimerLabel = getIdleTimerLabel(agent.getNetwork().netId, transportType);
 
             // Network is considered active when the network becomes the default network.
@@ -11371,6 +11385,24 @@
                     TIMESTAMP);
             assertFalse(onNetworkActiveCv.block(TEST_CALLBACK_TIMEOUT_MS));
             assertFalse(mCm.isDefaultNetworkActive());
+            if (mDeps.isAtLeastV()) {
+                if (transportType == TRANSPORT_CELLULAR) {
+                    verify(mIBatteryStats).noteMobileRadioPowerState(eq(DC_POWER_STATE_LOW),
+                            anyLong() /* timestampNs */, eq(NETWORK_ACTIVITY_NO_UID));
+                } else if (transportType == TRANSPORT_WIFI) {
+                    verify(mIBatteryStats).noteWifiRadioPowerState(eq(DC_POWER_STATE_LOW),
+                            anyLong() /* timestampNs */, eq(NETWORK_ACTIVITY_NO_UID));
+                }
+            } else {
+                // If TrackMultiNetworks is disabled, LegacyNetworkActivityTracker does not call
+                // BatteryStats API by the netd activity change callback since BatteryStatsService
+                // listen to netd callback via NetworkManagementService and update battery stats by
+                // itself.
+                verify(mIBatteryStats, never())
+                        .noteMobileRadioPowerState(anyInt(), anyLong(), anyInt());
+                verify(mIBatteryStats, never())
+                        .noteWifiRadioPowerState(anyInt(), anyLong(), anyInt());
+            }
 
             // Interface goes to active state
             netdUnsolicitedEventListener.onInterfaceClassActivityChanged(true /* isActive */,
@@ -11378,7 +11410,27 @@
             mServiceContext.expectDataActivityBroadcast(legacyType, true /* isActive */, TIMESTAMP);
             assertTrue(onNetworkActiveCv.block(TEST_CALLBACK_TIMEOUT_MS));
             assertTrue(mCm.isDefaultNetworkActive());
+            if (mDeps.isAtLeastV()) {
+                if (transportType == TRANSPORT_CELLULAR) {
+                    verify(mIBatteryStats).noteMobileRadioPowerState(eq(DC_POWER_STATE_HIGH),
+                            anyLong() /* timestampNs */, eq(TEST_PACKAGE_UID));
+                } else if (transportType == TRANSPORT_WIFI) {
+                    verify(mIBatteryStats).noteWifiRadioPowerState(eq(DC_POWER_STATE_HIGH),
+                            anyLong() /* timestampNs */, eq(TEST_PACKAGE_UID));
+                }
+            } else {
+                // If TrackMultiNetworks is disabled, LegacyNetworkActivityTracker does not call
+                // BatteryStats API by the netd activity change callback since BatteryStatsService
+                // listen to netd callback via NetworkManagementService and update battery stats by
+                // itself.
+                verify(mIBatteryStats, never())
+                        .noteMobileRadioPowerState(anyInt(), anyLong(), anyInt());
+                verify(mIBatteryStats, never())
+                        .noteWifiRadioPowerState(anyInt(), anyLong(), anyInt());
+            }
         }, () -> { // Cleanup
+                mCm.unregisterNetworkCallback(defaultCallback);
+            }, () -> { // Cleanup
                 mCm.removeDefaultNetworkActiveListener(listener);
             }, () -> { // Cleanup
                 agent.disconnect();
diff --git a/tests/unit/java/com/android/server/connectivityservice/CSNetworkActivityTest.kt b/tests/unit/java/com/android/server/connectivityservice/CSNetworkActivityTest.kt
index 1f3ba44..526ec9d 100644
--- a/tests/unit/java/com/android/server/connectivityservice/CSNetworkActivityTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/CSNetworkActivityTest.kt
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 package com.android.server
 
 import android.net.ConnectivityManager
@@ -31,6 +32,8 @@
 import android.net.NetworkRequest
 import android.os.Build
 import android.os.ConditionVariable
+import android.telephony.DataConnectionRealTimeInfo.DC_POWER_STATE_HIGH
+import android.telephony.DataConnectionRealTimeInfo.DC_POWER_STATE_LOW
 import androidx.test.filters.SmallTest
 import com.android.net.module.util.BaseNetdUnsolicitedEventListener
 import com.android.server.CSTest.CSContext
@@ -46,7 +49,10 @@
 import org.mockito.ArgumentCaptor
 import org.mockito.ArgumentMatchers.eq
 import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.anyLong
+import org.mockito.Mockito.inOrder
 import org.mockito.Mockito.never
+import org.mockito.Mockito.timeout
 import org.mockito.Mockito.verify
 
 private const val DATA_CELL_IFNAME = "rmnet_data"
@@ -55,7 +61,7 @@
 private const val TIMESTAMP = 1234L
 private const val NETWORK_ACTIVITY_NO_UID = -1
 private const val PACKAGE_UID = 123
-private const val CALLBACK_TIMEOUT_MS = 250L
+private const val TIMEOUT_MS = 250L
 
 @RunWith(DevSdkIgnoreRunner::class)
 @SmallTest
@@ -71,6 +77,7 @@
     @Test
     fun testInterfaceClassActivityChanged_NonDefaultNetwork() {
         val netdUnsolicitedEventListener = getRegisteredNetdUnsolicitedEventListener()
+        val batteryStatsInorder = inOrder(batteryStats)
 
         val cellNr = NetworkRequest.Builder()
                 .clearCapabilities()
@@ -110,6 +117,8 @@
         val wifiAgent = Agent(nc = wifiNc, lp = wifiLp)
         wifiAgent.connect()
         defaultCb.expectAvailableCallbacks(wifiAgent.network, validated = false)
+        batteryStatsInorder.verify(batteryStats).noteWifiRadioPowerState(eq(DC_POWER_STATE_HIGH),
+                anyLong() /* timestampNs */, eq(NETWORK_ACTIVITY_NO_UID))
 
         val onNetworkActiveCv = ConditionVariable()
         val listener = ConnectivityManager.OnNetworkActiveListener { onNetworkActiveCv::open }
@@ -119,17 +128,23 @@
         netdUnsolicitedEventListener.onInterfaceClassActivityChanged(false /* isActive */,
                 cellAgent.network.netId, TIMESTAMP, NETWORK_ACTIVITY_NO_UID)
         // Non-default network activity change does not change default network activity
-        assertFalse(onNetworkActiveCv.block(CALLBACK_TIMEOUT_MS))
+        // But cellular radio power state is updated
+        assertFalse(onNetworkActiveCv.block(TIMEOUT_MS))
         context.expectNoDataActivityBroadcast(0 /* timeoutMs */)
         assertTrue(cm.isDefaultNetworkActive)
+        batteryStatsInorder.verify(batteryStats).noteMobileRadioPowerState(eq(DC_POWER_STATE_LOW),
+                anyLong() /* timestampNs */, eq(NETWORK_ACTIVITY_NO_UID))
 
         // Cellular network (non default network) goes to active state.
         netdUnsolicitedEventListener.onInterfaceClassActivityChanged(true /* isActive */,
                 cellAgent.network.netId, TIMESTAMP, PACKAGE_UID)
         // Non-default network activity change does not change default network activity
-        assertFalse(onNetworkActiveCv.block(CALLBACK_TIMEOUT_MS))
+        // But cellular radio power state is updated
+        assertFalse(onNetworkActiveCv.block(TIMEOUT_MS))
         context.expectNoDataActivityBroadcast(0 /* timeoutMs */)
         assertTrue(cm.isDefaultNetworkActive)
+        batteryStatsInorder.verify(batteryStats).noteMobileRadioPowerState(eq(DC_POWER_STATE_HIGH),
+                anyLong() /* timestampNs */, eq(PACKAGE_UID))
 
         cm.unregisterNetworkCallback(cellCb)
         cm.unregisterNetworkCallback(defaultCb)
@@ -138,6 +153,9 @@
 
     @Test
     fun testDataActivityTracking_MultiCellNetwork() {
+        val netdUnsolicitedEventListener = getRegisteredNetdUnsolicitedEventListener()
+        val batteryStatsInorder = inOrder(batteryStats)
+
         val dataNetworkNc = NetworkCapabilities.Builder()
                 .addTransportType(TRANSPORT_CELLULAR)
                 .addCapability(NET_CAPABILITY_INTERNET)
@@ -188,6 +206,40 @@
         verify(netd, never()).idletimerRemoveInterface(eq(IMS_CELL_IFNAME), anyInt(),
                 eq(imsNetworkNetId))
 
+        // Both cell networks go to inactive state
+        netdUnsolicitedEventListener.onInterfaceClassActivityChanged(false /* isActive */,
+                imsNetworkAgent.network.netId, TIMESTAMP, NETWORK_ACTIVITY_NO_UID)
+        netdUnsolicitedEventListener.onInterfaceClassActivityChanged(false /* isActive */,
+                dataNetworkAgent.network.netId, TIMESTAMP, NETWORK_ACTIVITY_NO_UID)
+
+        // Data cell network goes to active state. This should update the cellular radio power state
+        netdUnsolicitedEventListener.onInterfaceClassActivityChanged(true /* isActive */,
+                dataNetworkAgent.network.netId, TIMESTAMP, PACKAGE_UID)
+        batteryStatsInorder.verify(batteryStats, timeout(TIMEOUT_MS)).noteMobileRadioPowerState(
+                eq(DC_POWER_STATE_HIGH), anyLong() /* timestampNs */, eq(PACKAGE_UID))
+        // Ims cell network goes to active state. But this should not update the cellular radio
+        // power state since cellular radio power state is already high
+        netdUnsolicitedEventListener.onInterfaceClassActivityChanged(true /* isActive */,
+                imsNetworkAgent.network.netId, TIMESTAMP, PACKAGE_UID)
+        waitForIdle()
+        batteryStatsInorder.verify(batteryStats, never()).noteMobileRadioPowerState(anyInt(),
+                anyLong() /* timestampNs */, anyInt())
+
+        // Data cell network goes to inactive state. But this should not update the cellular radio
+        // power state ims cell network is still active state
+        netdUnsolicitedEventListener.onInterfaceClassActivityChanged(false /* isActive */,
+                dataNetworkAgent.network.netId, TIMESTAMP, NETWORK_ACTIVITY_NO_UID)
+        waitForIdle()
+        batteryStatsInorder.verify(batteryStats, never()).noteMobileRadioPowerState(anyInt(),
+                anyLong() /* timestampNs */, anyInt())
+
+        // Ims cell network goes to inactive state.
+        // This should update the cellular radio power state
+        netdUnsolicitedEventListener.onInterfaceClassActivityChanged(false /* isActive */,
+                imsNetworkAgent.network.netId, TIMESTAMP, NETWORK_ACTIVITY_NO_UID)
+        batteryStatsInorder.verify(batteryStats, timeout(TIMEOUT_MS)).noteMobileRadioPowerState(
+                eq(DC_POWER_STATE_LOW), anyLong() /* timestampNs */, eq(NETWORK_ACTIVITY_NO_UID))
+
         dataNetworkAgent.disconnect()
         dataNetworkCb.expect<Lost>(dataNetworkAgent.network)
         verify(netd).idletimerRemoveInterface(eq(DATA_CELL_IFNAME), anyInt(), eq(dataNetworkNetId))
diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
index a28fef7..8881c5c 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
@@ -158,7 +158,8 @@
     val proxyTracker = ProxyTracker(context, mock<Handler>(), 16 /* EVENT_PROXY_HAS_CHANGED */)
     val alarmManager = makeMockAlarmManager()
     val systemConfigManager = makeMockSystemConfigManager()
-    val batteryManager = BatteryStatsManager(mock<IBatteryStats>())
+    val batteryStats = mock<IBatteryStats>()
+    val batteryManager = BatteryStatsManager(batteryStats)
     val telephonyManager = mock<TelephonyManager>().also {
         doReturn(true).`when`(it).isDataCapable()
     }