From 86ce1809aef3189254c17d4354a691fb28fbec3d Mon Sep 17 00:00:00 2001
From: Jernej Virag <jernej@google.com>
Date: Tue, 20 Dec 2022 16:15:47 +0100
Subject: [PATCH] Do not keep MediaSession Tokens for Notification accounting

In some cases we're tracking 50+ MediaSession.Tokens being retained. Since Tokens contain strong binder links, they can retain big resources on the other side of the link.

The code inside MediaSessionBasedFilter doesn't actually need to hold tokens at all - it just tests for equality between tokens. As such, we can only retain their identifiers and let Tokens themselves be deallocated if necessary.

Bug: 261882243
Test: atest MediaSessionBasedFilterTest, verified Media behaviour on device
Change-Id: I65fc3b47ad87cd6b90a814ef217b2ec1c90cac29
---
 .../pipeline/MediaSessionBasedFilter.kt       | 25 +++++++++++++------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaSessionBasedFilter.kt b/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaSessionBasedFilter.kt
index ab93b292308e4..d6f941de1b6d4 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaSessionBasedFilter.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaSessionBasedFilter.kt
@@ -58,10 +58,10 @@ constructor(
 
     // Keep track of the key used for the session tokens. This information is used to know when to
     // dispatch a removed event so that a media object for a local session will be removed.
-    private val keyedTokens: MutableMap<String, MutableSet<MediaSession.Token>> = mutableMapOf()
+    private val keyedTokens: MutableMap<String, MutableSet<TokenId>> = mutableMapOf()
 
     // Keep track of which media session tokens have associated notifications.
-    private val tokensWithNotifications: MutableSet<MediaSession.Token> = mutableSetOf()
+    private val tokensWithNotifications: MutableSet<TokenId> = mutableSetOf()
 
     private val sessionListener =
         object : MediaSessionManager.OnActiveSessionsChangedListener {
@@ -101,15 +101,15 @@ constructor(
         isSsReactivated: Boolean
     ) {
         backgroundExecutor.execute {
-            data.token?.let { tokensWithNotifications.add(it) }
+            data.token?.let { tokensWithNotifications.add(TokenId(it)) }
             val isMigration = oldKey != null && key != oldKey
             if (isMigration) {
                 keyedTokens.remove(oldKey)?.let { removed -> keyedTokens.put(key, removed) }
             }
             if (data.token != null) {
-                keyedTokens.get(key)?.let { tokens -> tokens.add(data.token) }
+                keyedTokens.get(key)?.let { tokens -> tokens.add(TokenId(data.token)) }
                     ?: run {
-                        val tokens = mutableSetOf(data.token)
+                        val tokens = mutableSetOf(TokenId(data.token))
                         keyedTokens.put(key, tokens)
                     }
             }
@@ -125,7 +125,7 @@ constructor(
                 isMigration ||
                     remote == null ||
                     remote.sessionToken == data.token ||
-                    !tokensWithNotifications.contains(remote.sessionToken)
+                    !tokensWithNotifications.contains(TokenId(remote.sessionToken))
             ) {
                 // Not filtering in this case. Passing the event along to listeners.
                 dispatchMediaDataLoaded(key, oldKey, data, immediately)
@@ -136,7 +136,7 @@ constructor(
                 // If the local session uses a different notification key, then lets go a step
                 // farther and dismiss the media data so that media controls for the local session
                 // don't hang around while casting.
-                if (!keyedTokens.get(key)!!.contains(remote.sessionToken)) {
+                if (!keyedTokens.get(key)!!.contains(TokenId(remote.sessionToken))) {
                     dispatchMediaDataRemoved(key)
                 }
             }
@@ -199,6 +199,15 @@ constructor(
                     packageControllers.put(controller.packageName, tokens)
                 }
         }
-        tokensWithNotifications.retainAll(controllers.map { it.sessionToken })
+        tokensWithNotifications.retainAll(controllers.map { TokenId(it.sessionToken) })
+    }
+
+    /**
+     * Represents a unique identifier for a [MediaSession.Token].
+     *
+     * It's used to avoid storing strong binders for media session tokens.
+     */
+    private data class TokenId(val id: Int) {
+        constructor(token: MediaSession.Token) : this(token.hashCode())
     }
 }