From 8c262d6c072304ed9f16feca64b70a18645cc908 Mon Sep 17 00:00:00 2001
From: tonihei <tonihei@google.com>
Date: Thu, 27 Apr 2023 12:54:17 +0100
Subject: [PATCH] Fix leaks of media session service.

References to the service are kept from MediaSessionStub
and from a long-delayed Handler messages in ConnectionTimeoutHandler.

Remove strong references from these places by making the timeout
handler static and ensuring ConnectedControllersManager only keeps
a weak reference to the service (as it's part of MediaSessionStub).

Issue: androidx/media#346
PiperOrigin-RevId: 527543396
---
 RELEASENOTES.md                                   |  3 +++
 .../session/ConnectedControllersManager.java      | 15 +++++++++++++--
 .../media3/session/MediaSessionLegacyStub.java    | 13 +++++++++----
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/RELEASENOTES.md b/RELEASENOTES.md
index f383948cc35..d0ea6966c1b 100644
--- a/RELEASENOTES.md
+++ b/RELEASENOTES.md
@@ -66,6 +66,9 @@
         `handlePlayPauseButtonAction` to write custom UI elements with a
         play/pause button.
 
+    *   Fix memory leak of `MediaSessionService` or `MediaLibraryService`
+        ([#346](https://github.com/androidx/media/issues/346)).
+
 *   Audio:
 
     *   Fix bug where some playbacks fail when tunneling is enabled and
diff --git a/libraries/session/src/main/java/androidx/media3/session/ConnectedControllersManager.java b/libraries/session/src/main/java/androidx/media3/session/ConnectedControllersManager.java
index 40e3d9619c4..71efbf6d8f0 100644
--- a/libraries/session/src/main/java/androidx/media3/session/ConnectedControllersManager.java
+++ b/libraries/session/src/main/java/androidx/media3/session/ConnectedControllersManager.java
@@ -26,6 +26,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
+import java.lang.ref.WeakReference;
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -62,14 +63,14 @@ public interface AsyncCommand {
   private final ArrayMap<ControllerInfo, ConnectedControllerRecord<T>> controllerRecords =
       new ArrayMap<>();
 
-  private final MediaSessionImpl sessionImpl;
+  private final WeakReference<MediaSessionImpl> sessionImpl;
 
   public ConnectedControllersManager(MediaSessionImpl session) {
     // Initialize default values.
     lock = new Object();
 
     // Initialize members with params.
-    sessionImpl = session;
+    sessionImpl = new WeakReference<>(session);
   }
 
   public void addController(
@@ -136,6 +137,10 @@ record = controllerRecords.remove(controllerInfo);
     }
 
     record.sequencedFutureManager.release();
+    @Nullable MediaSessionImpl sessionImpl = this.sessionImpl.get();
+    if (sessionImpl == null || sessionImpl.isReleased()) {
+      return;
+    }
     postOrRun(
         sessionImpl.getApplicationHandler(),
         () -> {
@@ -214,8 +219,10 @@ public boolean isPlayerCommandAvailable(
     synchronized (lock) {
       info = controllerRecords.get(controllerInfo);
     }
+    @Nullable MediaSessionImpl sessionImpl = this.sessionImpl.get();
     return info != null
         && info.playerCommands.contains(commandCode)
+        && sessionImpl != null
         && sessionImpl.getPlayerWrapper().getAvailableCommands().contains(commandCode);
   }
 
@@ -248,6 +255,10 @@ public void flushCommandQueue(ControllerInfo controllerInfo) {
 
   @GuardedBy("lock")
   private void flushCommandQueue(ConnectedControllerRecord<T> info) {
+    @Nullable MediaSessionImpl sessionImpl = this.sessionImpl.get();
+    if (sessionImpl == null) {
+      return;
+    }
     AtomicBoolean continueRunning = new AtomicBoolean(true);
     while (continueRunning.get()) {
       continueRunning.set(false);
diff --git a/libraries/session/src/main/java/androidx/media3/session/MediaSessionLegacyStub.java b/libraries/session/src/main/java/androidx/media3/session/MediaSessionLegacyStub.java
index cefbcbb6202..1699f29c7c2 100644
--- a/libraries/session/src/main/java/androidx/media3/session/MediaSessionLegacyStub.java
+++ b/libraries/session/src/main/java/androidx/media3/session/MediaSessionLegacyStub.java
@@ -145,12 +145,13 @@ public MediaSessionLegacyStub(
     appPackageName = context.getPackageName();
     sessionManager = MediaSessionManager.getSessionManager(context);
     controllerLegacyCbForBroadcast = new ControllerLegacyCbForBroadcast();
-    connectionTimeoutHandler =
-        new ConnectionTimeoutHandler(session.getApplicationHandler().getLooper());
     mediaPlayPauseKeyHandler =
         new MediaPlayPauseKeyHandler(session.getApplicationHandler().getLooper());
     connectedControllersManager = new ConnectedControllersManager<>(session);
     connectionTimeoutMs = DEFAULT_CONNECTION_TIMEOUT_MS;
+    connectionTimeoutHandler =
+        new ConnectionTimeoutHandler(
+            session.getApplicationHandler().getLooper(), connectedControllersManager);
 
     // Select a media button receiver component.
     ComponentName receiverComponentName = queryPackageManagerForMediaButtonReceiver(context);
@@ -1359,12 +1360,16 @@ public void onFailure(Throwable t) {
     }
   }
 
-  private class ConnectionTimeoutHandler extends Handler {
+  private static class ConnectionTimeoutHandler extends Handler {
 
     private static final int MSG_CONNECTION_TIMED_OUT = 1001;
 
-    public ConnectionTimeoutHandler(Looper looper) {
+    private final ConnectedControllersManager<RemoteUserInfo> connectedControllersManager;
+
+    public ConnectionTimeoutHandler(
+        Looper looper, ConnectedControllersManager<RemoteUserInfo> connectedControllersManager) {
       super(looper);
+      this.connectedControllersManager = connectedControllersManager;
     }
 
     @Override