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 kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
+import com.android.settingslib.spa.framework.util.asyncMap
 
 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 {
-        getCycles().map {
-            async {
-                queryDetails(it)
-            }
-        }.awaitAll().filter { it.totalUsage > 0 }
-    }
+    override suspend fun queryDetailsForCycles(): List<NetworkUsageDetailsData> =
+        getCycles().asyncMap { queryDetails(it) }.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
-        var foregroundUsage = 0L
-        for (uid in withSdkSandboxUids) {
-            val usage = getUsage(range, uid, NetworkStats.Bucket.STATE_ALL)
-            if (usage > 0L) {
-                totalUsage += usage
-                foregroundUsage += getUsage(range, uid, NetworkStats.Bucket.STATE_FOREGROUND)
-            }
-        }
+        val buckets = networkStatsRepository.queryBuckets(range.lower, range.upper)
+            .filter { it.uid in withSdkSandboxUids }
+        val totalUsage = buckets.sumOf { it.bytes }
+        val foregroundUsage =
+            buckets.filter { it.state == NetworkStats.Bucket.STATE_FOREGROUND }.sumOf { it.bytes }
         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 kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
+import com.android.settingslib.spa.framework.util.asyncMap
 
 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 {
-        withSdkSandboxUids(listOf(uid)).map { uid ->
-            async {
-                networkStatsRepository.queryAggregateForUid(range = AllTimeRange, uid = uid)
-            }
-        }.awaitAll().filterNotNull().aggregate()
-    }
+    override suspend fun querySummary(uid: Int): NetworkUsageData? =
+        withSdkSandboxUids(listOf(uid)).asyncMap {
+            networkStatsRepository.queryAggregateForUid(range = AllTimeRange, uid = it)
+        }.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 {
-        const val DATE_FORMAT = DateUtils.FORMAT_SHOW_DATE or DateUtils.FORMAT_ABBREV_MONTH
+    companion object {
+        val AllZero = NetworkUsageData(
+            startTime = 0L,
+            endTime = 0L,
+            usage = 0L,
+        )
+
+        private const val DATE_FORMAT = DateUtils.FORMAT_SHOW_DATE or DateUtils.FORMAT_ABBREV_MONTH
     }
 }