Improve AppDataUsageDetailsRepository

Reduce latency by reduce the number of system calls.

Bug: 295260929
Test: manual - on AppDataUsage page
Test: unit test
Change-Id: If233e223db744cd15d3a769416fd4c5957085417

src/com/android/settings/datausage/ChartDataUsagePreferenceController.kt

@@ -26,6 +26,7 @@ import androidx.lifecycle.lifecycleScope
 import androidx.preference.PreferenceScreen
 import com.android.settings.core.BasePreferenceController
 import com.android.settings.datausage.lib.INetworkCycleDataRepository
+import com.android.settings.datausage.lib.NetworkCycleChartData
 import com.android.settings.datausage.lib.NetworkCycleDataRepository
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -72,6 +73,7 @@ open class ChartDataUsagePreferenceController(context: Context, preferenceKey: S
 
     fun update(startTime: Long, endTime: Long) {
         preference.setTime(startTime, endTime)
+        preference.setNetworkCycleData(NetworkCycleChartData.AllZero)
         lifecycleScope.launch {
             val chartData = withContext(Dispatchers.Default) {
                 repository.queryChartData(startTime, endTime)

src/com/android/settings/datausage/lib/AppDataUsageDetailsRepository.kt

@@ -20,11 +20,8 @@ import android.app.usage.NetworkStats
 import android.content.Context
 import android.net.NetworkTemplate
 import android.util.Range
-import androidx.annotation.VisibleForTesting
 import com.android.settings.datausage.lib.AppDataUsageRepository.Companion.withSdkSandboxUids
-import kotlinx.coroutines.async
+import com.android.settingslib.spa.framework.util.asyncMap
-import kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
 
 interface IAppDataUsageDetailsRepository {
     suspend fun queryDetailsForCycles(): List<NetworkUsageDetailsData>
@@ -37,32 +34,24 @@ class AppDataUsageDetailsRepository @JvmOverloads constructor(
     uids: List<Int>,
     private val networkCycleDataRepository: INetworkCycleDataRepository =
         NetworkCycleDataRepository(context, template),
+    private val networkStatsRepository: NetworkStatsRepository =
+        NetworkStatsRepository(context, template),
 ) : IAppDataUsageDetailsRepository {
     private val withSdkSandboxUids = withSdkSandboxUids(uids)
-    private val networkStatsRepository = NetworkStatsRepository(context, template)
 
-    override suspend fun queryDetailsForCycles(): List<NetworkUsageDetailsData> = coroutineScope {
+    override suspend fun queryDetailsForCycles(): List<NetworkUsageDetailsData> =
-        getCycles().map {
+        getCycles().asyncMap { queryDetails(it) }.filter { it.totalUsage > 0 }
-            async {
-                queryDetails(it)
-            }
-        }.awaitAll().filter { it.totalUsage > 0 }
-    }
 
     private fun getCycles(): List<Range<Long>> =
         cycles?.zipWithNext { endTime, startTime -> Range(startTime, endTime) }
             ?: networkCycleDataRepository.getCycles()
 
     private fun queryDetails(range: Range<Long>): NetworkUsageDetailsData {
-        var totalUsage = 0L
+        val buckets = networkStatsRepository.queryBuckets(range.lower, range.upper)
-        var foregroundUsage = 0L
+            .filter { it.uid in withSdkSandboxUids }
-        for (uid in withSdkSandboxUids) {
+        val totalUsage = buckets.sumOf { it.bytes }
-            val usage = getUsage(range, uid, NetworkStats.Bucket.STATE_ALL)
+        val foregroundUsage =
-            if (usage > 0L) {
+            buckets.filter { it.state == NetworkStats.Bucket.STATE_FOREGROUND }.sumOf { it.bytes }
-                totalUsage += usage
-                foregroundUsage += getUsage(range, uid, NetworkStats.Bucket.STATE_FOREGROUND)
-            }
-        }
         return NetworkUsageDetailsData(
             range = range,
             totalUsage = totalUsage,
@@ -70,8 +59,4 @@ class AppDataUsageDetailsRepository @JvmOverloads constructor(
             backgroundUsage = totalUsage - foregroundUsage,
         )
     }
-
-    @VisibleForTesting
-    fun getUsage(range: Range<Long>, uid: Int, state: Int): Long =
-        networkStatsRepository.queryAggregateForUid(range, uid, state)?.usage ?: 0
 }

src/com/android/settings/datausage/lib/AppDataUsageRepository.kt

@@ -41,7 +41,7 @@ class AppDataUsageRepository(
     private val networkStatsRepository = NetworkStatsRepository(context, template)
 
     fun getAppPercent(carrierId: Int?, startTime: Long, endTime: Long): List<Pair<AppItem, Int>> {
-        val buckets = networkStatsRepository.querySummary(startTime, endTime)
+        val buckets = networkStatsRepository.queryBuckets(startTime, endTime)
         return getAppPercent(carrierId, buckets)
     }
 

src/com/android/settings/datausage/lib/AppDataUsageSummaryRepository.kt

@@ -20,9 +20,7 @@ import android.content.Context
 import android.net.NetworkTemplate
 import com.android.settings.datausage.lib.AppDataUsageRepository.Companion.withSdkSandboxUids
 import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.AllTimeRange
-import kotlinx.coroutines.async
+import com.android.settingslib.spa.framework.util.asyncMap
-import kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
 
 interface IAppDataUsageSummaryRepository {
     suspend fun querySummary(uid: Int): NetworkUsageData?
@@ -35,11 +33,8 @@ class AppDataUsageSummaryRepository(
         NetworkStatsRepository(context, template),
 ) : IAppDataUsageSummaryRepository {
 
-    override suspend fun querySummary(uid: Int): NetworkUsageData? = coroutineScope {
+    override suspend fun querySummary(uid: Int): NetworkUsageData? =
-        withSdkSandboxUids(listOf(uid)).map { uid ->
+        withSdkSandboxUids(listOf(uid)).asyncMap {
-            async {
+            networkStatsRepository.queryAggregateForUid(range = AllTimeRange, uid = it)
-                networkStatsRepository.queryAggregateForUid(range = AllTimeRange, uid = uid)
+        }.filterNotNull().aggregate()
-            }
-        }.awaitAll().filterNotNull().aggregate()
-    }
 }

src/com/android/settings/datausage/lib/NetworkCycleChartData.kt

@@ -26,6 +26,11 @@ data class NetworkCycleChartData(
     val dailyUsage: List<NetworkUsageData>,
 ) {
     companion object {
+        val AllZero = NetworkCycleChartData(
+            total = NetworkUsageData.AllZero,
+            dailyUsage = emptyList(),
+        )
+
         val BUCKET_DURATION = 1.days
     }
 }

src/com/android/settings/datausage/lib/NetworkStatsRepository.kt

@@ -54,7 +54,7 @@ class NetworkStatsRepository(context: Context, private val template: NetworkTemp
         0
     }
 
-    fun querySummary(startTime: Long, endTime: Long): List<Bucket> = try {
+    fun queryBuckets(startTime: Long, endTime: Long): List<Bucket> = try {
         networkStatsManager.querySummary(template, startTime, endTime).convertToBuckets()
     } catch (e: Exception) {
         Log.e(TAG, "Exception querySummary", e)
@@ -69,13 +69,14 @@ class NetworkStatsRepository(context: Context, private val template: NetworkTemp
         data class Bucket(
             val uid: Int,
             val bytes: Long,
+            val state: Int = NetworkStats.Bucket.STATE_ALL,
         )
 
         private fun NetworkStats.convertToBuckets(): List<Bucket> = use {
             val buckets = mutableListOf<Bucket>()
             val bucket = NetworkStats.Bucket()
             while (getNextBucket(bucket)) {
-                buckets += Bucket(uid = bucket.uid, bytes = bucket.bytes)
+                buckets += Bucket(uid = bucket.uid, bytes = bucket.bytes, state = bucket.state)
             }
             buckets
         }

src/com/android/settings/datausage/lib/NetworkUsageData.kt

@@ -43,8 +43,14 @@ data class NetworkUsageData(
     fun getDataUsedString(context: Context): String =
         context.getString(R.string.data_used_template, formatUsage(context))
 
-    private companion object {
+    companion object {
-        const val DATE_FORMAT = DateUtils.FORMAT_SHOW_DATE or DateUtils.FORMAT_ABBREV_MONTH
+        val AllZero = NetworkUsageData(
+            startTime = 0L,
+            endTime = 0L,
+            usage = 0L,
+        )
+
+        private const val DATE_FORMAT = DateUtils.FORMAT_SHOW_DATE or DateUtils.FORMAT_ABBREV_MONTH
     }
 }
 

tests/spa_unit/src/com/android/settings/datausage/lib/AppDataUsageDetailsRepositoryTest.kt

@@ -16,7 +16,7 @@
 
 package com.android.settings.datausage.lib
 
-import android.app.usage.NetworkStats.Bucket
+import android.app.usage.NetworkStats
 import android.content.Context
 import android.net.NetworkTemplate
 import android.util.Range
@@ -28,8 +28,8 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
-import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
-import org.mockito.kotlin.whenever
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.Bucket
 
 @RunWith(AndroidJUnit4::class)
 class AppDataUsageDetailsRepositoryTest {
@@ -41,58 +41,78 @@ class AppDataUsageDetailsRepositoryTest {
         on { getCycles() } doReturn listOf(Range(CYCLE1_END_TIME, CYCLE2_END_TIME))
     }
 
+    private val networkStatsRepository = mock<NetworkStatsRepository>()
+
     @Test
     fun queryDetailsForCycles_hasCycles(): Unit = runBlocking {
-        val range = Range(CYCLE1_START_TIME, CYCLE1_END_TIME)
+        networkStatsRepository.stub {
-        val repository = spy(
+            on { queryBuckets(CYCLE1_START_TIME, CYCLE1_END_TIME) } doReturn listOf(
-            AppDataUsageDetailsRepository(
+                Bucket(
+                    uid = UID,
+                    state = NetworkStats.Bucket.STATE_DEFAULT,
+                    bytes = BACKGROUND_USAGE,
+                ),
+                Bucket(
+                    uid = UID,
+                    state = NetworkStats.Bucket.STATE_FOREGROUND,
+                    bytes = FOREGROUND_USAGE,
+                ),
+            )
+        }
+        val repository = AppDataUsageDetailsRepository(
             context = context,
             cycles = listOf(CYCLE1_END_TIME, CYCLE1_START_TIME),
             template = template,
             uids = listOf(UID),
             networkCycleDataRepository = networkCycleDataRepository,
+            networkStatsRepository = networkStatsRepository,
         )
-        ) {
-            doReturn(ALL_USAGE).whenever(mock).getUsage(range, UID, Bucket.STATE_ALL)
-            doReturn(FOREGROUND_USAGE).whenever(mock).getUsage(range, UID, Bucket.STATE_FOREGROUND)
-        }
 
         val detailsForCycles = repository.queryDetailsForCycles()
 
         assertThat(detailsForCycles).containsExactly(
             NetworkUsageDetailsData(
-                range = range,
+                range = Range(CYCLE1_START_TIME, CYCLE1_END_TIME),
-                totalUsage = ALL_USAGE,
+                totalUsage = BACKGROUND_USAGE + FOREGROUND_USAGE,
                 foregroundUsage = FOREGROUND_USAGE,
-                backgroundUsage = ALL_USAGE - FOREGROUND_USAGE,
+                backgroundUsage = BACKGROUND_USAGE,
             )
         )
     }
 
     @Test
     fun queryDetailsForCycles_defaultCycles(): Unit = runBlocking {
-        val range = Range(CYCLE1_END_TIME, CYCLE2_END_TIME)
+        networkStatsRepository.stub {
-        val repository = spy(
+            on { queryBuckets(CYCLE1_END_TIME, CYCLE2_END_TIME) } doReturn listOf(
-            AppDataUsageDetailsRepository(
+                Bucket(
+                    uid = UID,
+                    state = NetworkStats.Bucket.STATE_DEFAULT,
+                    bytes = BACKGROUND_USAGE,
+                ),
+                Bucket(
+                    uid = UID,
+                    state = NetworkStats.Bucket.STATE_FOREGROUND,
+                    bytes = FOREGROUND_USAGE,
+                ),
+            )
+        }
+        val repository = AppDataUsageDetailsRepository(
             context = context,
             cycles = null,
             template = template,
             uids = listOf(UID),
             networkCycleDataRepository = networkCycleDataRepository,
+            networkStatsRepository = networkStatsRepository,
         )
-        ) {
-            doReturn(ALL_USAGE).whenever(mock).getUsage(range, UID, Bucket.STATE_ALL)
-            doReturn(FOREGROUND_USAGE).whenever(mock).getUsage(range, UID, Bucket.STATE_FOREGROUND)
-        }
 
         val detailsForCycles = repository.queryDetailsForCycles()
 
         assertThat(detailsForCycles).containsExactly(
             NetworkUsageDetailsData(
-                range = range,
+                range = Range(CYCLE1_END_TIME, CYCLE2_END_TIME),
-                totalUsage = ALL_USAGE,
+                totalUsage = BACKGROUND_USAGE + FOREGROUND_USAGE,
                 foregroundUsage = FOREGROUND_USAGE,
-                backgroundUsage = ALL_USAGE - FOREGROUND_USAGE,
+                backgroundUsage = BACKGROUND_USAGE,
             )
         )
     }
@@ -103,7 +123,7 @@ class AppDataUsageDetailsRepositoryTest {
         const val CYCLE2_END_TIME = 1695566666000L
         const val UID = 10000
 
-        const val ALL_USAGE = 10L
+        const val BACKGROUND_USAGE = 8L
         const val FOREGROUND_USAGE = 2L
     }
 }