Add view walking to notification memory monitor

This adds monitor for bitmap memory usage of Notification views - public and private versions for now. It counts both existing templates and custom views. The code is still flag guarded. Bug: 235451049 Test: atest com.android.systemui.statusbar.notification Change-Id: Ibdc8caaa002622ed51303e24b7d57e225465357c
2022-10-07 10:50:27 +02:00
parent b0d018239d
commit a52c75387d
7 changed files with 698 additions and 233 deletions
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
@@ -20,8 +20,9 @@ package com.android.systemui.statusbar.notification.logging
 /** Describes usage of a notification. */
 data class NotificationMemoryUsage(
    val packageName: String,
-    val notificationId: String,
+    val notificationKey: String,
    val objectUsage: NotificationObjectUsage,
+    val viewUsage: List<NotificationViewUsage>
 )

 /**
@@ -39,3 +40,26 @@ data class NotificationObjectUsage(
    val extender: Int,
    val hasCustomView: Boolean,
 )
+
+enum class ViewType {
+    PUBLIC_VIEW,
+    PRIVATE_CONTRACTED_VIEW,
+    PRIVATE_EXPANDED_VIEW,
+    PRIVATE_HEADS_UP_VIEW,
+    TOTAL
+}
+
+/**
+ * Describes current memory of a notification view hierarchy.
+ *
+ * The values are in bytes.
+ */
+data class NotificationViewUsage(
+    val viewType: ViewType,
+    val smallIcon: Int,
+    val largeIcon: Int,
+    val systemIcons: Int,
+    val style: Int,
+    val customViews: Int,
+    val softwareBitmapsPenalty: Int,
+)
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt
@@ -0,0 +1,212 @@
+package com.android.systemui.statusbar.notification.logging
+
+import android.app.Notification
+import android.app.Person
+import android.graphics.Bitmap
+import android.graphics.drawable.Icon
+import android.os.Bundle
+import android.os.Parcel
+import android.os.Parcelable
+import androidx.annotation.WorkerThread
+import com.android.systemui.statusbar.notification.NotificationUtils
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+
+/** Calculates estimated memory usage of [Notification] and [NotificationEntry] objects. */
+internal object NotificationMemoryMeter {
+
+    private const val CAR_EXTENSIONS = "android.car.EXTENSIONS"
+    private const val CAR_EXTENSIONS_LARGE_ICON = "large_icon"
+    private const val TV_EXTENSIONS = "android.tv.EXTENSIONS"
+    private const val WEARABLE_EXTENSIONS = "android.wearable.EXTENSIONS"
+    private const val WEARABLE_EXTENSIONS_BACKGROUND = "background"
+
+    /** Returns a list of memory use entries for currently shown notifications. */
+    @WorkerThread
+    fun notificationMemoryUse(
+        notifications: Collection<NotificationEntry>,
+    ): List<NotificationMemoryUsage> {
+        return notifications
+            .asSequence()
+            .map { entry ->
+                val packageName = entry.sbn.packageName
+                val notificationObjectUsage =
+                    notificationMemoryUse(entry.sbn.notification, hashSetOf())
+                val notificationViewUsage = NotificationMemoryViewWalker.getViewUsage(entry.row)
+                NotificationMemoryUsage(
+                    packageName,
+                    NotificationUtils.logKey(entry.sbn.key),
+                    notificationObjectUsage,
+                    notificationViewUsage
+                )
+            }
+            .toList()
+    }
+
+    @WorkerThread
+    fun notificationMemoryUse(
+        entry: NotificationEntry,
+        seenBitmaps: HashSet<Int> = hashSetOf(),
+    ): NotificationMemoryUsage {
+        return NotificationMemoryUsage(
+            entry.sbn.packageName,
+            NotificationUtils.logKey(entry.sbn.key),
+            notificationMemoryUse(entry.sbn.notification, seenBitmaps),
+            NotificationMemoryViewWalker.getViewUsage(entry.row)
+        )
+    }
+
+    /**
+     * Computes the estimated memory usage of a given [Notification] object. It'll attempt to
+     * inspect Bitmaps in the object and provide summary of memory usage.
+     */
+    @WorkerThread
+    fun notificationMemoryUse(
+        notification: Notification,
+        seenBitmaps: HashSet<Int> = hashSetOf(),
+    ): NotificationObjectUsage {
+        val extras = notification.extras
+        val smallIconUse = computeIconUse(notification.smallIcon, seenBitmaps)
+        val largeIconUse = computeIconUse(notification.getLargeIcon(), seenBitmaps)
+
+        // Collect memory usage of extra styles
+
+        // Big Picture
+        val bigPictureIconUse =
+            computeParcelableUse(extras, Notification.EXTRA_LARGE_ICON_BIG, seenBitmaps)
+        val bigPictureUse =
+            computeParcelableUse(extras, Notification.EXTRA_PICTURE, seenBitmaps) +
+                computeParcelableUse(extras, Notification.EXTRA_PICTURE_ICON, seenBitmaps)
+
+        // People
+        val peopleList = extras.getParcelableArrayList<Person>(Notification.EXTRA_PEOPLE_LIST)
+        val peopleUse =
+            peopleList?.sumOf { person -> computeIconUse(person.icon, seenBitmaps) } ?: 0
+
+        // Calling
+        val callingPersonUse =
+            computeParcelableUse(extras, Notification.EXTRA_CALL_PERSON, seenBitmaps)
+        val verificationIconUse =
+            computeParcelableUse(extras, Notification.EXTRA_VERIFICATION_ICON, seenBitmaps)
+
+        // Messages
+        val messages =
+            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
+                extras.getParcelableArray(Notification.EXTRA_MESSAGES)
+            )
+        val messagesUse =
+            messages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
+        val historicMessages =
+            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
+                extras.getParcelableArray(Notification.EXTRA_HISTORIC_MESSAGES)
+            )
+        val historyicMessagesUse =
+            historicMessages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
+
+        // Extenders
+        val carExtender = extras.getBundle(CAR_EXTENSIONS)
+        val carExtenderSize = carExtender?.let { computeBundleSize(it) } ?: 0
+        val carExtenderIcon =
+            computeParcelableUse(carExtender, CAR_EXTENSIONS_LARGE_ICON, seenBitmaps)
+
+        val tvExtender = extras.getBundle(TV_EXTENSIONS)
+        val tvExtenderSize = tvExtender?.let { computeBundleSize(it) } ?: 0
+
+        val wearExtender = extras.getBundle(WEARABLE_EXTENSIONS)
+        val wearExtenderSize = wearExtender?.let { computeBundleSize(it) } ?: 0
+        val wearExtenderBackground =
+            computeParcelableUse(wearExtender, WEARABLE_EXTENSIONS_BACKGROUND, seenBitmaps)
+
+        val style = notification.notificationStyle
+        val hasCustomView = notification.contentView != null || notification.bigContentView != null
+        val extrasSize = computeBundleSize(extras)
+
+        return NotificationObjectUsage(
+            smallIcon = smallIconUse,
+            largeIcon = largeIconUse,
+            extras = extrasSize,
+            style = style?.simpleName,
+            styleIcon =
+                bigPictureIconUse +
+                    peopleUse +
+                    callingPersonUse +
+                    verificationIconUse +
+                    messagesUse +
+                    historyicMessagesUse,
+            bigPicture = bigPictureUse,
+            extender =
+                carExtenderSize +
+                    carExtenderIcon +
+                    tvExtenderSize +
+                    wearExtenderSize +
+                    wearExtenderBackground,
+            hasCustomView = hasCustomView
+        )
+    }
+
+    /**
+     * Calculates size of the bundle data (excluding FDs and other shared objects like ashmem
+     * bitmaps). Can be slow.
+     */
+    private fun computeBundleSize(extras: Bundle): Int {
+        val parcel = Parcel.obtain()
+        try {
+            extras.writeToParcel(parcel, 0)
+            return parcel.dataSize()
+        } finally {
+            parcel.recycle()
+        }
+    }
+
+    /**
+     * Deserializes [Icon], [Bitmap] or [Person] from extras and computes its memory use. Returns 0
+     * if the key does not exist in extras.
+     */
+    private fun computeParcelableUse(extras: Bundle?, key: String, seenBitmaps: HashSet<Int>): Int {
+        return when (val parcelable = extras?.getParcelable<Parcelable>(key)) {
+            is Bitmap -> computeBitmapUse(parcelable, seenBitmaps)
+            is Icon -> computeIconUse(parcelable, seenBitmaps)
+            is Person -> computeIconUse(parcelable.icon, seenBitmaps)
+            else -> 0
+        }
+    }
+
+    /**
+     * Calculates the byte size of bitmaps or data in the Icon object. Returns 0 if the icon is
+     * defined via Uri or a resource.
+     *
+     * @return memory usage in bytes or 0 if the icon is Uri/Resource based
+     */
+    private fun computeIconUse(icon: Icon?, seenBitmaps: HashSet<Int>) =
+        when (icon?.type) {
+            Icon.TYPE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
+            Icon.TYPE_ADAPTIVE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
+            Icon.TYPE_DATA -> computeDataUse(icon, seenBitmaps)
+            else -> 0
+        }
+
+    /**
+     * Returns the amount of memory a given bitmap is using. If the bitmap reference is part of
+     * seenBitmaps set, this method returns 0 to avoid double counting.
+     *
+     * @return memory usage of the bitmap in bytes
+     */
+    private fun computeBitmapUse(bitmap: Bitmap, seenBitmaps: HashSet<Int>? = null): Int {
+        val refId = System.identityHashCode(bitmap)
+        if (seenBitmaps?.contains(refId) == true) {
+            return 0
+        }
+
+        seenBitmaps?.add(refId)
+        return bitmap.allocationByteCount
+    }
+
+    private fun computeDataUse(icon: Icon, seenBitmaps: HashSet<Int>): Int {
+        val refId = System.identityHashCode(icon.dataBytes)
+        if (seenBitmaps.contains(refId)) {
+            return 0
+        }
+
+        seenBitmaps.add(refId)
+        return icon.dataLength
+    }
+}
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
@@ -17,22 +17,11 @@

 package com.android.systemui.statusbar.notification.logging

-import android.app.Notification
-import android.app.Person
-import android.graphics.Bitmap
-import android.graphics.drawable.Icon
-import android.os.Bundle
-import android.os.Parcel
-import android.os.Parcelable
 import android.util.Log
-import androidx.annotation.WorkerThread
-import androidx.core.util.contains
 import com.android.systemui.Dumpable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.statusbar.notification.NotificationUtils
 import com.android.systemui.statusbar.notification.collection.NotifPipeline
-import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import java.io.PrintWriter
 import javax.inject.Inject

@@ -46,12 +35,7 @@ constructor(
 ) : Dumpable {

    companion object {
-        private const val TAG = "NotificationMemMonitor"
-        private const val CAR_EXTENSIONS = "android.car.EXTENSIONS"
-        private const val CAR_EXTENSIONS_LARGE_ICON = "large_icon"
-        private const val TV_EXTENSIONS = "android.tv.EXTENSIONS"
-        private const val WEARABLE_EXTENSIONS = "android.wearable.EXTENSIONS"
-        private const val WEARABLE_EXTENSIONS_BACKGROUND = "background"
+        private const val TAG = "NotificationMemory"
    }

    fun init() {
@@ -60,184 +44,123 @@ constructor(
    }

    override fun dump(pw: PrintWriter, args: Array<out String>) {
-        currentNotificationMemoryUse().forEach { use -> pw.println(use.toString()) }
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(notificationPipeline.allNotifs)
+                .sortedWith(compareBy({ it.packageName }, { it.notificationKey }))
+        dumpNotificationObjects(pw, memoryUse)
+        dumpNotificationViewUsage(pw, memoryUse)
    }

-    @WorkerThread
-    fun currentNotificationMemoryUse(): List<NotificationMemoryUsage> {
-        return notificationMemoryUse(notificationPipeline.allNotifs)
-    }
-
-    /** Returns a list of memory use entries for currently shown notifications. */
-    @WorkerThread
-    fun notificationMemoryUse(
-        notifications: Collection<NotificationEntry>
-    ): List<NotificationMemoryUsage> {
-        return notifications
-            .asSequence()
-            .map { entry ->
-                val packageName = entry.sbn.packageName
-                val notificationObjectUsage =
-                    computeNotificationObjectUse(entry.sbn.notification, hashSetOf())
-                NotificationMemoryUsage(
-                    packageName,
-                    NotificationUtils.logKey(entry.sbn.key),
-                    notificationObjectUsage
-                )
-            }
-            .toList()
-    }
-
-    /**
-     * Computes the estimated memory usage of a given [Notification] object. It'll attempt to
-     * inspect Bitmaps in the object and provide summary of memory usage.
-     */
-    private fun computeNotificationObjectUse(
-        notification: Notification,
-        seenBitmaps: HashSet<Int>
-    ): NotificationObjectUsage {
-        val extras = notification.extras
-        val smallIconUse = computeIconUse(notification.smallIcon, seenBitmaps)
-        val largeIconUse = computeIconUse(notification.getLargeIcon(), seenBitmaps)
-
-        // Collect memory usage of extra styles
-
-        // Big Picture
-        val bigPictureIconUse =
-            computeParcelableUse(extras, Notification.EXTRA_PICTURE_ICON, seenBitmaps) +
-                computeParcelableUse(extras, Notification.EXTRA_LARGE_ICON_BIG, seenBitmaps)
-        val bigPictureUse =
-            computeParcelableUse(extras, Notification.EXTRA_PICTURE, seenBitmaps) +
-                computeParcelableUse(extras, Notification.EXTRA_PICTURE_ICON, seenBitmaps)
-
-        // People
-        val peopleList = extras.getParcelableArrayList<Person>(Notification.EXTRA_PEOPLE_LIST)
-        val peopleUse =
-            peopleList?.sumOf { person -> computeIconUse(person.icon, seenBitmaps) } ?: 0
-
-        // Calling
-        val callingPersonUse =
-            computeParcelableUse(extras, Notification.EXTRA_CALL_PERSON, seenBitmaps)
-        val verificationIconUse =
-            computeParcelableUse(extras, Notification.EXTRA_VERIFICATION_ICON, seenBitmaps)
-
-        // Messages
-        val messages =
-            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
-                extras.getParcelableArray(Notification.EXTRA_MESSAGES)
-            )
-        val messagesUse =
-            messages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
-        val historicMessages =
-            Notification.MessagingStyle.Message.getMessagesFromBundleArray(
-                extras.getParcelableArray(Notification.EXTRA_HISTORIC_MESSAGES)
-            )
-        val historyicMessagesUse =
-            historicMessages.sumOf { msg -> computeIconUse(msg.senderPerson?.icon, seenBitmaps) }
-
-        // Extenders
-        val carExtender = extras.getBundle(CAR_EXTENSIONS)
-        val carExtenderSize = carExtender?.let { computeBundleSize(it) } ?: 0
-        val carExtenderIcon =
-            computeParcelableUse(carExtender, CAR_EXTENSIONS_LARGE_ICON, seenBitmaps)
-
-        val tvExtender = extras.getBundle(TV_EXTENSIONS)
-        val tvExtenderSize = tvExtender?.let { computeBundleSize(it) } ?: 0
-
-        val wearExtender = extras.getBundle(WEARABLE_EXTENSIONS)
-        val wearExtenderSize = wearExtender?.let { computeBundleSize(it) } ?: 0
-        val wearExtenderBackground =
-            computeParcelableUse(wearExtender, WEARABLE_EXTENSIONS_BACKGROUND, seenBitmaps)
-
-        val style = notification.notificationStyle
-        val hasCustomView = notification.contentView != null || notification.bigContentView != null
-        val extrasSize = computeBundleSize(extras)
-
-        return NotificationObjectUsage(
-            smallIconUse,
-            largeIconUse,
-            extrasSize,
-            style?.simpleName,
-            bigPictureIconUse +
-                peopleUse +
-                callingPersonUse +
-                verificationIconUse +
-                messagesUse +
-                historyicMessagesUse,
-            bigPictureUse,
-            carExtenderSize +
-                carExtenderIcon +
-                tvExtenderSize +
-                wearExtenderSize +
-                wearExtenderBackground,
-            hasCustomView
+    /** Renders a table of notification object usage into passed [PrintWriter]. */
+    private fun dumpNotificationObjects(pw: PrintWriter, memoryUse: List<NotificationMemoryUsage>) {
+        pw.println("Notification Object Usage")
+        pw.println("-----------")
+        pw.println(
+            "Package".padEnd(35) +
+                "\t\tSmall\tLarge\t${"Style".padEnd(15)}\t\tStyle\tBig\tExtend.\tExtras\tCustom"
        )
+        pw.println("".padEnd(35) + "\t\tIcon\tIcon\t${"".padEnd(15)}\t\tIcon\tPicture\t \t \tView")
+        pw.println()
+
+        memoryUse.forEach { use ->
+            pw.println(
+                use.packageName.padEnd(35) +
+                    "\t\t" +
+                    "${use.objectUsage.smallIcon}\t${use.objectUsage.largeIcon}\t" +
+                    (use.objectUsage.style?.take(15) ?: "").padEnd(15) +
+                    "\t\t${use.objectUsage.styleIcon}\t" +
+                    "${use.objectUsage.bigPicture}\t${use.objectUsage.extender}\t" +
+                    "${use.objectUsage.extras}\t${use.objectUsage.hasCustomView}\t" +
+                    use.notificationKey
+            )
+        }
+
+        // Calculate totals for easily glanceable summary.
+        data class Totals(
+            var smallIcon: Int = 0,
+            var largeIcon: Int = 0,
+            var styleIcon: Int = 0,
+            var bigPicture: Int = 0,
+            var extender: Int = 0,
+            var extras: Int = 0,
+        )
+
+        val totals =
+            memoryUse.fold(Totals()) { t, usage ->
+                t.smallIcon += usage.objectUsage.smallIcon
+                t.largeIcon += usage.objectUsage.largeIcon
+                t.styleIcon += usage.objectUsage.styleIcon
+                t.bigPicture += usage.objectUsage.bigPicture
+                t.extender += usage.objectUsage.extender
+                t.extras += usage.objectUsage.extras
+                t
+            }
+
+        pw.println()
+        pw.println("TOTALS")
+        pw.println(
+            "".padEnd(35) +
+                "\t\t" +
+                "${toKb(totals.smallIcon)}\t${toKb(totals.largeIcon)}\t" +
+                "".padEnd(15) +
+                "\t\t${toKb(totals.styleIcon)}\t" +
+                "${toKb(totals.bigPicture)}\t${toKb(totals.extender)}\t" +
+                toKb(totals.extras)
+        )
+        pw.println()
    }

-    /**
-     * Calculates size of the bundle data (excluding FDs and other shared objects like ashmem
-     * bitmaps). Can be slow.
-     */
-    private fun computeBundleSize(extras: Bundle): Int {
-        val parcel = Parcel.obtain()
-        try {
-            extras.writeToParcel(parcel, 0)
-            return parcel.dataSize()
-        } finally {
-            parcel.recycle()
-        }
+    /** Renders a table of notification view usage into passed [PrintWriter] */
+    private fun dumpNotificationViewUsage(
+        pw: PrintWriter,
+        memoryUse: List<NotificationMemoryUsage>,
+    ) {
+
+        data class Totals(
+            var smallIcon: Int = 0,
+            var largeIcon: Int = 0,
+            var style: Int = 0,
+            var customViews: Int = 0,
+            var softwareBitmapsPenalty: Int = 0,
+        )
+
+        val totals = Totals()
+        pw.println("Notification View Usage")
+        pw.println("-----------")
+        pw.println("View Type".padEnd(24) + "\tSmall\tLarge\tStyle\tCustom\tSoftware")
+        pw.println("".padEnd(24) + "\tIcon\tIcon\tUse\tView\tBitmaps")
+        pw.println()
+        memoryUse
+            .filter { it.viewUsage.isNotEmpty() }
+            .forEach { use ->
+                pw.println(use.packageName + " " + use.notificationKey)
+                use.viewUsage.forEach { view ->
+                    pw.println(
+                        "  ${view.viewType.toString().padEnd(24)}\t${view.smallIcon}" +
+                            "\t${view.largeIcon}\t${view.style}" +
+                            "\t${view.customViews}\t${view.softwareBitmapsPenalty}"
+                    )
+
+                    if (view.viewType == ViewType.TOTAL) {
+                        totals.smallIcon += view.smallIcon
+                        totals.largeIcon += view.largeIcon
+                        totals.style += view.style
+                        totals.customViews += view.customViews
+                        totals.softwareBitmapsPenalty += view.softwareBitmapsPenalty
+                    }
+                }
+            }
+        pw.println()
+        pw.println("TOTALS")
+        pw.println(
+            "  ${"".padEnd(24)}\t${toKb(totals.smallIcon)}" +
+                "\t${toKb(totals.largeIcon)}\t${toKb(totals.style)}" +
+                "\t${toKb(totals.customViews)}\t${toKb(totals.softwareBitmapsPenalty)}"
+        )
+        pw.println()
    }

-    /**
-     * Deserializes [Icon], [Bitmap] or [Person] from extras and computes its memory use. Returns 0
-     * if the key does not exist in extras.
-     */
-    private fun computeParcelableUse(extras: Bundle?, key: String, seenBitmaps: HashSet<Int>): Int {
-        return when (val parcelable = extras?.getParcelable<Parcelable>(key)) {
-            is Bitmap -> computeBitmapUse(parcelable, seenBitmaps)
-            is Icon -> computeIconUse(parcelable, seenBitmaps)
-            is Person -> computeIconUse(parcelable.icon, seenBitmaps)
-            else -> 0
-        }
-    }
-
-    /**
-     * Calculates the byte size of bitmaps or data in the Icon object. Returns 0 if the icon is
-     * defined via Uri or a resource.
-     *
-     * @return memory usage in bytes or 0 if the icon is Uri/Resource based
-     */
-    private fun computeIconUse(icon: Icon?, seenBitmaps: HashSet<Int>) =
-        when (icon?.type) {
-            Icon.TYPE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
-            Icon.TYPE_ADAPTIVE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
-            Icon.TYPE_DATA -> computeDataUse(icon, seenBitmaps)
-            else -> 0
-        }
-
-    /**
-     * Returns the amount of memory a given bitmap is using. If the bitmap reference is part of
-     * seenBitmaps set, this method returns 0 to avoid double counting.
-     *
-     * @return memory usage of the bitmap in bytes
-     */
-    private fun computeBitmapUse(bitmap: Bitmap, seenBitmaps: HashSet<Int>? = null): Int {
-        val refId = System.identityHashCode(bitmap)
-        if (seenBitmaps?.contains(refId) == true) {
-            return 0
-        }
-
-        seenBitmaps?.add(refId)
-        return bitmap.allocationByteCount
-    }
-
-    private fun computeDataUse(icon: Icon, seenBitmaps: HashSet<Int>): Int {
-        val refId = System.identityHashCode(icon.dataBytes)
-        if (seenBitmaps.contains(refId)) {
-            return 0
-        }
-
-        seenBitmaps.add(refId)
-        return icon.dataLength
+    private fun toKb(bytes: Int): String {
+        return (bytes / 1024).toString() + " KB"
    }
 }
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
@@ -0,0 +1,173 @@
+package com.android.systemui.statusbar.notification.logging
+
+import android.graphics.Bitmap
+import android.graphics.drawable.BitmapDrawable
+import android.graphics.drawable.Drawable
+import android.util.Log
+import android.view.View
+import android.view.ViewGroup
+import android.widget.ImageView
+import com.android.internal.R
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
+import com.android.systemui.util.children
+
+/** Walks view hiearchy of a given notification to estimate its memory use. */
+internal object NotificationMemoryViewWalker {
+
+    private const val TAG = "NotificationMemory"
+
+    /** Builder for [NotificationViewUsage] objects. */
+    private class UsageBuilder {
+        private var smallIcon: Int = 0
+        private var largeIcon: Int = 0
+        private var systemIcons: Int = 0
+        private var style: Int = 0
+        private var customViews: Int = 0
+        private var softwareBitmaps = 0
+
+        fun addSmallIcon(smallIconUse: Int) = apply { smallIcon += smallIconUse }
+        fun addLargeIcon(largeIconUse: Int) = apply { largeIcon += largeIconUse }
+        fun addSystem(systemIconUse: Int) = apply { systemIcons += systemIconUse }
+        fun addStyle(styleUse: Int) = apply { style += styleUse }
+        fun addSoftwareBitmapPenalty(softwareBitmapUse: Int) = apply {
+            softwareBitmaps += softwareBitmapUse
+        }
+
+        fun addCustomViews(customViewsUse: Int) = apply { customViews += customViewsUse }
+
+        fun build(viewType: ViewType): NotificationViewUsage {
+            return NotificationViewUsage(
+                viewType = viewType,
+                smallIcon = smallIcon,
+                largeIcon = largeIcon,
+                systemIcons = systemIcons,
+                style = style,
+                customViews = customViews,
+                softwareBitmapsPenalty = softwareBitmaps,
+            )
+        }
+    }
+
+    /**
+     * Returns memory usage of public and private views contained in passed
+     * [ExpandableNotificationRow]
+     */
+    fun getViewUsage(row: ExpandableNotificationRow?): List<NotificationViewUsage> {
+        if (row == null) {
+            return listOf()
+        }
+
+        // The ordering here is significant since it determines deduplication of seen drawables.
+        return listOf(
+            getViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, row.privateLayout?.expandedChild),
+            getViewUsage(ViewType.PRIVATE_CONTRACTED_VIEW, row.privateLayout?.contractedChild),
+            getViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, row.privateLayout?.headsUpChild),
+            getViewUsage(ViewType.PUBLIC_VIEW, row.publicLayout),
+            getTotalUsage(row)
+        )
+    }
+
+    /**
+     * Calculate total usage of all views - we need to do a separate traversal to make sure we don't
+     * double count fields.
+     */
+    private fun getTotalUsage(row: ExpandableNotificationRow): NotificationViewUsage {
+        val totalUsage = UsageBuilder()
+        val seenObjects = hashSetOf<Int>()
+
+        row.publicLayout?.let { computeViewHierarchyUse(it, totalUsage, seenObjects) }
+        row.privateLayout?.let { child ->
+            for (view in listOf(child.expandedChild, child.contractedChild, child.headsUpChild)) {
+                (view as? ViewGroup)?.let { v ->
+                    computeViewHierarchyUse(v, totalUsage, seenObjects)
+                }
+            }
+        }
+        return totalUsage.build(ViewType.TOTAL)
+    }
+
+    private fun getViewUsage(
+        type: ViewType,
+        rootView: View?,
+        seenObjects: HashSet<Int> = hashSetOf()
+    ): NotificationViewUsage {
+        val usageBuilder = UsageBuilder()
+        (rootView as? ViewGroup)?.let { computeViewHierarchyUse(it, usageBuilder, seenObjects) }
+        return usageBuilder.build(type)
+    }
+
+    private fun computeViewHierarchyUse(
+        rootView: ViewGroup,
+        builder: UsageBuilder,
+        seenObjects: HashSet<Int> = hashSetOf(),
+    ) {
+        for (child in rootView.children) {
+            if (child is ViewGroup) {
+                computeViewHierarchyUse(child, builder, seenObjects)
+            } else {
+                computeViewUse(child, builder, seenObjects)
+            }
+        }
+    }
+
+    private fun computeViewUse(view: View, builder: UsageBuilder, seenObjects: HashSet<Int>) {
+        if (view !is ImageView) return
+        val drawable = view.drawable ?: return
+        val drawableRef = System.identityHashCode(drawable)
+        if (seenObjects.contains(drawableRef)) return
+        val drawableUse = computeDrawableUse(drawable, seenObjects)
+        // TODO(b/235451049): We need to make sure we traverse large icon before small icon -
+        // sometimes the large icons are assigned to small icon views and we want to
+        // attribute them to large view in those cases.
+        when (view.id) {
+            R.id.left_icon,
+            R.id.icon,
+            R.id.conversation_icon -> builder.addSmallIcon(drawableUse)
+            R.id.right_icon -> builder.addLargeIcon(drawableUse)
+            R.id.big_picture -> builder.addStyle(drawableUse)
+            // Elements that are part of platform with resources
+            R.id.phishing_alert,
+            R.id.feedback,
+            R.id.alerted_icon,
+            R.id.expand_button_icon,
+            R.id.remote_input_send -> builder.addSystem(drawableUse)
+            // Custom view ImageViews
+            else -> {
+                if (Log.isLoggable(TAG, Log.DEBUG)) {
+                    Log.d(TAG, "Custom view: ${identifierForView(view)}")
+                }
+                builder.addCustomViews(drawableUse)
+            }
+        }
+
+        if (isDrawableSoftwareBitmap(drawable)) {
+            builder.addSoftwareBitmapPenalty(drawableUse)
+        }
+
+        seenObjects.add(drawableRef)
+    }
+
+    private fun computeDrawableUse(drawable: Drawable, seenObjects: HashSet<Int>): Int =
+        when (drawable) {
+            is BitmapDrawable -> {
+                val ref = System.identityHashCode(drawable.bitmap)
+                if (seenObjects.contains(ref)) {
+                    0
+                } else {
+                    seenObjects.add(ref)
+                    drawable.bitmap.allocationByteCount
+                }
+            }
+            else -> 0
+        }
+
+    private fun isDrawableSoftwareBitmap(drawable: Drawable) =
+        drawable is BitmapDrawable && drawable.bitmap.config != Bitmap.Config.HARDWARE
+
+    private fun identifierForView(view: View) =
+        if (view.id == View.NO_ID) {
+            "no-id"
+        } else {
+            view.resources.getResourceName(view.id)
+        }
+}
--- a/packages/SystemUI/tests/res/layout/custom_view_dark.xml
+++ b/packages/SystemUI/tests/res/layout/custom_view_dark.xml
@@ -14,6 +14,7 @@
    limitations under the License.
 -->
 <ImageView xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/custom_view_dark_image"
    android:layout_width="match_parent"
    android:layout_height="match_parent"
    android:background="#ff000000"
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitorTest.kt
@@ -28,30 +28,21 @@ import android.widget.RemoteViews
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.statusbar.notification.NotificationUtils
-import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.MockitoAnnotations

@SmallTest
@RunWith(AndroidTestingRunner::class)
-class NotificationMemoryMonitorTest : SysuiTestCase() {
-
-    @Before
-    fun setUp() {
-        MockitoAnnotations.initMocks(this)
-    }
+class NotificationMemoryMeterTest : SysuiTestCase() {

    @Test
    fun currentNotificationMemoryUse_plainNotification() {
        val notification = createBasicNotification().build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -69,8 +60,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
    fun currentNotificationMemoryUse_plainNotification_dontDoubleCountSameBitmap() {
        val icon = Icon.createWithBitmap(Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888))
        val notification = createBasicNotification().setLargeIcon(icon).setSmallIcon(icon).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -92,8 +83,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
                    RemoteViews(context.packageName, android.R.layout.list_content)
                )
                .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -112,8 +103,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
        val dataIcon = Icon.createWithData(ByteArray(444444), 0, 444444)
        val notification =
            createBasicNotification().setLargeIcon(dataIcon).setSmallIcon(dataIcon).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = 444444,
@@ -141,8 +132,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
                        .bigLargeIcon(bigPictureIcon)
                )
                .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -167,8 +158,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
            createBasicNotification()
                .setStyle(Notification.CallStyle.forIncomingCall(person, fakeIntent, fakeIntent))
                .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -203,8 +194,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
                        .addHistoricMessage(historicMessage)
                )
                .build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -225,8 +216,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
        val carIcon = Bitmap.createBitmap(432, 322, Bitmap.Config.ARGB_8888)
        val extender = Notification.CarExtender().setLargeIcon(carIcon)
        val notification = createBasicNotification().extend(extender).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -246,8 +237,8 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
        val wearBackground = Bitmap.createBitmap(443, 433, Bitmap.Config.ARGB_8888)
        val wearExtender = Notification.WearableExtender().setBackground(wearBackground)
        val notification = createBasicNotification().extend(tvExtender).extend(wearExtender).build()
-        val nmm = createNMMWithNotifications(listOf(notification))
-        val memoryUse = getUseObject(nmm.currentNotificationMemoryUse())
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(createNotificationEntry(notification))
        assertNotificationObjectSizes(
            memoryUse = memoryUse,
            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
@@ -283,10 +274,10 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
        extender: Int,
        style: String?,
        styleIcon: Int,
-        hasCustomView: Boolean
+        hasCustomView: Boolean,
    ) {
        assertThat(memoryUse.packageName).isEqualTo("test_pkg")
-        assertThat(memoryUse.notificationId)
+        assertThat(memoryUse.notificationKey)
            .isEqualTo(NotificationUtils.logKey("0|test_pkg|0|test|0"))
        assertThat(memoryUse.objectUsage.smallIcon).isEqualTo(smallIcon)
        assertThat(memoryUse.objectUsage.largeIcon).isEqualTo(largeIcon)
@@ -301,21 +292,14 @@ class NotificationMemoryMonitorTest : SysuiTestCase() {
    }

    private fun getUseObject(
-        singleItemUseList: List<NotificationMemoryUsage>
+        singleItemUseList: List<NotificationMemoryUsage>,
    ): NotificationMemoryUsage {
        assertThat(singleItemUseList).hasSize(1)
        return singleItemUseList[0]
    }

-    private fun createNMMWithNotifications(
-        notifications: List<Notification>
-    ): NotificationMemoryMonitor {
-        val notifPipeline: NotifPipeline = mock()
-        val notificationEntries =
-            notifications.map { n ->
-                NotificationEntryBuilder().setTag("test").setNotification(n).build()
-            }
-        whenever(notifPipeline.allNotifs).thenReturn(notificationEntries)
-        return NotificationMemoryMonitor(notifPipeline, mock())
-    }
+    private fun createNotificationEntry(
+        notification: Notification,
+    ): NotificationEntry =
+        NotificationEntryBuilder().setTag("test").setNotification(notification).build()
 }
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt
@@ -0,0 +1,148 @@
+package com.android.systemui.statusbar.notification.logging
+
+import android.app.Notification
+import android.graphics.Bitmap
+import android.graphics.drawable.Icon
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import android.widget.RemoteViews
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.notification.row.NotificationTestHelper
+import com.android.systemui.tests.R
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@TestableLooper.RunWithLooper
+class NotificationMemoryViewWalkerTest : SysuiTestCase() {
+
+    private lateinit var testHelper: NotificationTestHelper
+
+    @Before
+    fun setUp() {
+        allowTestableLooperAsMainThread()
+        testHelper = NotificationTestHelper(mContext, mDependency, TestableLooper.get(this))
+    }
+
+    @Test
+    fun testViewWalker_nullRow_returnsEmptyView() {
+        val result = NotificationMemoryViewWalker.getViewUsage(null)
+        assertThat(result).isNotNull()
+        assertThat(result).isEmpty()
+    }
+
+    @Test
+    fun testViewWalker_plainNotification() {
+        val row = testHelper.createRow()
+        val result = NotificationMemoryViewWalker.getViewUsage(row)
+        assertThat(result).hasSize(5)
+        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_CONTRACTED_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(NotificationViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, 0, 0, 0, 0, 0, 0))
+    }
+
+    @Test
+    fun testViewWalker_bigPictureNotification() {
+        val bigPicture = Bitmap.createBitmap(300, 300, Bitmap.Config.ARGB_8888)
+        val icon = Icon.createWithBitmap(Bitmap.createBitmap(20, 20, Bitmap.Config.ARGB_8888))
+        val largeIcon = Icon.createWithBitmap(Bitmap.createBitmap(60, 60, Bitmap.Config.ARGB_8888))
+        val row =
+            testHelper.createRow(
+                Notification.Builder(mContext)
+                    .setContentText("Test")
+                    .setContentTitle("title")
+                    .setSmallIcon(icon)
+                    .setLargeIcon(largeIcon)
+                    .setStyle(Notification.BigPictureStyle().bigPicture(bigPicture))
+                    .build()
+            )
+        val result = NotificationMemoryViewWalker.getViewUsage(row)
+        assertThat(result).hasSize(5)
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_EXPANDED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    largeIcon.bitmap.allocationByteCount,
+                    0,
+                    bigPicture.allocationByteCount,
+                    0,
+                    bigPicture.allocationByteCount +
+                        icon.bitmap.allocationByteCount +
+                        largeIcon.bitmap.allocationByteCount
+                )
+            )
+
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_CONTRACTED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    largeIcon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    icon.bitmap.allocationByteCount + largeIcon.bitmap.allocationByteCount
+                )
+            )
+        // Due to deduplication, this should all be 0.
+        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+    }
+
+    @Test
+    fun testViewWalker_customView() {
+        val icon = Icon.createWithBitmap(Bitmap.createBitmap(20, 20, Bitmap.Config.ARGB_8888))
+        val bitmap = Bitmap.createBitmap(300, 300, Bitmap.Config.ARGB_8888)
+
+        val views = RemoteViews(mContext.packageName, R.layout.custom_view_dark)
+        views.setImageViewBitmap(R.id.custom_view_dark_image, bitmap)
+        val row =
+            testHelper.createRow(
+                Notification.Builder(mContext)
+                    .setContentText("Test")
+                    .setContentTitle("title")
+                    .setSmallIcon(icon)
+                    .setCustomContentView(views)
+                    .setCustomBigContentView(views)
+                    .build()
+            )
+        val result = NotificationMemoryViewWalker.getViewUsage(row)
+        assertThat(result).hasSize(5)
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_CONTRACTED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    bitmap.allocationByteCount,
+                    bitmap.allocationByteCount + icon.bitmap.allocationByteCount
+                )
+            )
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_EXPANDED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    bitmap.allocationByteCount,
+                    bitmap.allocationByteCount + icon.bitmap.allocationByteCount
+                )
+            )
+        // Due to deduplication, this should all be 0.
+        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+    }
+}