Clean up mediabrowser callbacks

The test connection was not getting cleaned up properly, which is why
VLC's foreground notification stayed around. It also should have been
replaced with the actual resumption controls once those were ready.
At the same time Spotify was not restarting on the first tap because the
browser was being disconnected prematurely. So, if we have successfully
connected to restart, wait until the new session info has come in to
disconnect the browser.

Fixes: 154825715
Fixes: 155098687
Test: manual

Change-Id: I207ae4790e5ae2b073def1ceb654808907d5a5c4
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaControlPanel.java b/packages/SystemUI/src/com/android/systemui/media/MediaControlPanel.java
index f876053..557132b 100644
--- a/packages/SystemUI/src/com/android/systemui/media/MediaControlPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaControlPanel.java
@@ -98,6 +98,7 @@
     private SharedPreferences mSharedPrefs;
     private boolean mCheckedForResumption = false;
     private boolean mIsRemotePlayback;
+    private QSMediaBrowser mQSMediaBrowser;
 
     // Button IDs used in notifications
     protected static final int[] NOTIF_ACTION_IDS = {
@@ -232,6 +233,11 @@
             String key) {
         // Ensure that component names are updated if token has changed
         if (mToken == null || !mToken.equals(token)) {
+            if (mQSMediaBrowser != null) {
+                Log.d(TAG, "Disconnecting old media browser");
+                mQSMediaBrowser.disconnect();
+                mQSMediaBrowser = null;
+            }
             mToken = token;
             mServiceComponent = null;
             mCheckedForResumption = false;
@@ -618,8 +624,22 @@
         ImageButton btn = mMediaNotifView.findViewById(mActionIds[0]);
         btn.setOnClickListener(v -> {
             Log.d(TAG, "Attempting to restart session");
-            QSMediaBrowser browser = new QSMediaBrowser(mContext, null, mServiceComponent);
-            browser.restart();
+            if (mQSMediaBrowser != null) {
+                mQSMediaBrowser.disconnect();
+            }
+            mQSMediaBrowser = new QSMediaBrowser(mContext, new QSMediaBrowser.Callback(){
+                @Override
+                public void onConnected() {
+                    Log.d(TAG, "Successfully restarted");
+                }
+                @Override
+                public void onError() {
+                    Log.e(TAG, "Error restarting");
+                    mQSMediaBrowser.disconnect();
+                    mQSMediaBrowser = null;
+                }
+            }, mServiceComponent);
+            mQSMediaBrowser.restart();
         });
         btn.setImageDrawable(mContext.getResources().getDrawable(R.drawable.lb_ic_play));
         btn.setImageTintList(ColorStateList.valueOf(mForegroundColor));
@@ -654,13 +674,18 @@
      */
     private void tryUpdateResumptionList(ComponentName componentName) {
         Log.d(TAG, "Testing if we can connect to " + componentName);
-        QSMediaBrowser.testConnection(mContext,
+        if (mQSMediaBrowser != null) {
+            mQSMediaBrowser.disconnect();
+        }
+        mQSMediaBrowser = new QSMediaBrowser(mContext,
                 new QSMediaBrowser.Callback() {
                     @Override
                     public void onConnected() {
                         Log.d(TAG, "yes we can resume with " + componentName);
                         mServiceComponent = componentName;
                         updateResumptionList(componentName);
+                        mQSMediaBrowser.disconnect();
+                        mQSMediaBrowser = null;
                     }
 
                     @Override
@@ -671,9 +696,12 @@
                             // If it's not active and we can't resume, remove
                             removePlayer();
                         }
+                        mQSMediaBrowser.disconnect();
+                        mQSMediaBrowser = null;
                     }
                 },
                 componentName);
+        mQSMediaBrowser.testConnection();
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSMediaBrowser.java b/packages/SystemUI/src/com/android/systemui/qs/QSMediaBrowser.java
index 9e53286..a5b73dc 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSMediaBrowser.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSMediaBrowser.java
@@ -58,21 +58,25 @@
         mContext = context;
         mCallback = callback;
         mComponentName = componentName;
+    }
 
+    /**
+     * Connects to the MediaBrowserService and looks for valid media. If a media item is returned
+     * by the service, QSMediaBrowser.Callback#addTrack will be called with its MediaDescription.
+     * QSMediaBrowser.Callback#onConnected and QSMediaBrowser.Callback#onError will also be called
+     * when the initial connection is successful, or an error occurs. Note that it is possible for
+     * the service to connect but for no playable tracks to be found later.
+     * QSMediaBrowser#disconnect will be called automatically with this function.
+     */
+    public void findRecentMedia() {
+        Log.d(TAG, "Connecting to " + mComponentName);
+        disconnect();
         Bundle rootHints = new Bundle();
         rootHints.putBoolean(MediaBrowserService.BrowserRoot.EXTRA_RECENT, true);
         mMediaBrowser = new MediaBrowser(mContext,
                 mComponentName,
                 mConnectionCallback,
                 rootHints);
-    }
-
-    /**
-     * Connects to the MediaBrowserService and looks for valid media. If a media item is returned
-     * by the service, QSMediaBrowser.Callback#addTrack will be called with its MediaDescription
-     */
-    public void findRecentMedia() {
-        Log.d(TAG, "Connecting to " + mComponentName);
         mMediaBrowser.connect();
     }
 
@@ -94,20 +98,22 @@
             } else {
                 Log.e(TAG, "Child found but not playable for " + mComponentName);
             }
-            mMediaBrowser.disconnect();
+            disconnect();
         }
 
         @Override
         public void onError(String parentId) {
             Log.e(TAG, "Subscribe error for " + mComponentName + ": " + parentId);
-            mMediaBrowser.disconnect();
+            mCallback.onError();
+            disconnect();
         }
 
         @Override
         public void onError(String parentId, Bundle options) {
             Log.e(TAG, "Subscribe error for " + mComponentName + ": " + parentId
                     + ", options: " + options);
-            mMediaBrowser.disconnect();
+            mCallback.onError();
+            disconnect();
         }
     };
 
@@ -134,6 +140,8 @@
         @Override
         public void onConnectionSuspended() {
             Log.d(TAG, "Connection suspended for " + mComponentName);
+            mCallback.onError();
+            disconnect();
         }
 
         /**
@@ -143,16 +151,28 @@
         public void onConnectionFailed() {
             Log.e(TAG, "Connection failed for " + mComponentName);
             mCallback.onError();
+            disconnect();
         }
     };
 
     /**
-     * Connects to the MediaBrowserService and starts playback
+     * Disconnect the media browser. This should be called after restart or testConnection have
+     * completed to close the connection.
      */
-    public void restart() {
-        if (mMediaBrowser.isConnected()) {
+    public void disconnect() {
+        if (mMediaBrowser != null) {
             mMediaBrowser.disconnect();
         }
+        mMediaBrowser = null;
+    }
+
+    /**
+     * Connects to the MediaBrowserService and starts playback. QSMediaBrowser.Callback#onError or
+     * QSMediaBrowser.Callback#onConnected will be called depending on whether it was successful.
+     * QSMediaBrowser#disconnect should be called after this to ensure the connection is closed.
+     */
+    public void restart() {
+        disconnect();
         Bundle rootHints = new Bundle();
         rootHints.putBoolean(MediaBrowserService.BrowserRoot.EXTRA_RECENT, true);
         mMediaBrowser = new MediaBrowser(mContext, mComponentName,
@@ -165,6 +185,17 @@
                         controller.getTransportControls();
                         controller.getTransportControls().prepare();
                         controller.getTransportControls().play();
+                        mCallback.onConnected();
+                    }
+
+                    @Override
+                    public void onConnectionFailed() {
+                        mCallback.onError();
+                    }
+
+                    @Override
+                    public void onConnectionSuspended() {
+                        mCallback.onError();
                     }
                 }, rootHints);
         mMediaBrowser.connect();
@@ -192,42 +223,44 @@
     }
 
     /**
-     * Used to test if SystemUI is allowed to connect to the given component as a MediaBrowser
-     * @param mContext the context
-     * @param callback methods onConnected or onError will be called to indicate whether the
-     *                 connection was successful or not
-     * @param mComponentName Component name of the MediaBrowserService this browser will connect to
+     * Used to test if SystemUI is allowed to connect to the given component as a MediaBrowser.
+     * QSMediaBrowser.Callback#onError or QSMediaBrowser.Callback#onConnected will be called
+     * depending on whether it was successful.
+     * QSMediaBrowser#disconnect should be called after this to ensure the connection is closed.
      */
-    public static MediaBrowser testConnection(Context mContext, Callback callback,
-            ComponentName mComponentName) {
-        final MediaBrowser.ConnectionCallback mConnectionCallback =
+    public void testConnection() {
+        disconnect();
+        final MediaBrowser.ConnectionCallback connectionCallback =
                 new MediaBrowser.ConnectionCallback() {
                     @Override
                     public void onConnected() {
                         Log.d(TAG, "connected");
-                        callback.onConnected();
+                        if (mMediaBrowser.getRoot() == null) {
+                            mCallback.onError();
+                        } else {
+                            mCallback.onConnected();
+                        }
                     }
 
                     @Override
                     public void onConnectionSuspended() {
                         Log.d(TAG, "suspended");
-                        callback.onError();
+                        mCallback.onError();
                     }
 
                     @Override
                     public void onConnectionFailed() {
                         Log.d(TAG, "failed");
-                        callback.onError();
+                        mCallback.onError();
                     }
                 };
         Bundle rootHints = new Bundle();
         rootHints.putBoolean(MediaBrowserService.BrowserRoot.EXTRA_RECENT, true);
-        MediaBrowser browser = new MediaBrowser(mContext,
+        mMediaBrowser = new MediaBrowser(mContext,
                 mComponentName,
-                mConnectionCallback,
+                connectionCallback,
                 rootHints);
-        browser.connect();
-        return browser;
+        mMediaBrowser.connect();
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
index a3004bd..e8f6c96 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
@@ -270,27 +270,8 @@
             return;
         }
 
-        QSMediaPlayer player = null;
         String packageName = notif.getPackageName();
-        for (QSMediaPlayer p : mMediaPlayers) {
-            if (p.getKey() == null) {
-                // No notification key = loaded via mediabrowser, so just match on package
-                if (packageName.equals(p.getMediaPlayerPackage())) {
-                    Log.d(TAG, "Found matching resume player by package: " + packageName);
-                    player = p;
-                    break;
-                }
-            } else if (p.getMediaSessionToken().equals(token)) {
-                Log.d(TAG, "Found matching player by token " + packageName);
-                player = p;
-                break;
-            } else if (packageName.equals(p.getMediaPlayerPackage()) && key.equals(p.getKey())) {
-                // Also match if it's the same package and notification key
-                Log.d(TAG, "Found matching player by package " + packageName + ", " + key);
-                player = p;
-                break;
-            }
-        }
+        QSMediaPlayer player = findMediaPlayer(packageName, token, key);
 
         int playerWidth = (int) getResources().getDimension(R.dimen.qs_media_width);
         int padding = (int) getResources().getDimension(R.dimen.qs_media_padding);
@@ -299,7 +280,7 @@
         lp.setMarginEnd(padding);
 
         if (player == null) {
-            Log.d(TAG, "creating new player");
+            Log.d(TAG, "creating new player for " + packageName);
             // Set up listener for device changes
             // TODO: integrate with MediaTransferManager?
             InfoMediaManager imm = new InfoMediaManager(mContext, notif.getPackageName(),
@@ -331,6 +312,35 @@
         }
     }
 
+    /**
+     * Check for an existing media player using the given information
+     * @param packageName
+     * @param token
+     * @param key
+     * @return a player, or null if no match found
+     */
+    private QSMediaPlayer findMediaPlayer(String packageName, MediaSession.Token token,
+            String key) {
+        for (QSMediaPlayer player : mMediaPlayers) {
+            if (player.getKey() == null || key == null) {
+                // No notification key = loaded via mediabrowser, so just match on package
+                if (packageName.equals(player.getMediaPlayerPackage())) {
+                    Log.d(TAG, "Found matching resume player by package: " + packageName);
+                    return player;
+                }
+            } else if (player.getMediaSessionToken().equals(token)) {
+                Log.d(TAG, "Found matching player by token " + packageName);
+                return player;
+            } else if (packageName.equals(player.getMediaPlayerPackage())
+                    && key.equals(player.getKey())) {
+                // Also match if it's the same package and notification key
+                Log.d(TAG, "Found matching player by package " + packageName + ", " + key);
+                return player;
+            }
+        }
+        return null;
+    }
+
     protected View getMediaPanel() {
         return mMediaCarousel;
     }
@@ -369,26 +379,30 @@
             }
 
             Log.d(TAG, "adding track from browser: " + desc + ", " + component);
-            QSMediaPlayer player = new QSMediaPlayer(mContext, QSPanel.this,
-                    null, mForegroundExecutor, mBackgroundExecutor, mActivityStarter);
 
+            // Check if there's an old player for this app
             String pkgName = component.getPackageName();
+            MediaSession.Token token = browser.getToken();
+            QSMediaPlayer player = findMediaPlayer(pkgName, token, null);
 
-            // Add controls to carousel
-            int playerWidth = (int) getResources().getDimension(R.dimen.qs_media_width);
-            int padding = (int) getResources().getDimension(R.dimen.qs_media_padding);
-            LinearLayout.LayoutParams lp = new LinearLayout.LayoutParams(playerWidth,
-                    LayoutParams.MATCH_PARENT);
-            lp.setMarginStart(padding);
-            lp.setMarginEnd(padding);
-            mMediaCarousel.addView(player.getView(), lp);
-            ((View) mMediaCarousel.getParent()).setVisibility(View.VISIBLE);
-            mMediaPlayers.add(player);
+            if (player == null) {
+                player = new QSMediaPlayer(mContext, QSPanel.this,
+                        null, mForegroundExecutor, mBackgroundExecutor, mActivityStarter);
+
+                // Add to carousel
+                int playerWidth = (int) getResources().getDimension(R.dimen.qs_media_width);
+                int padding = (int) getResources().getDimension(R.dimen.qs_media_padding);
+                LinearLayout.LayoutParams lp = new LinearLayout.LayoutParams(playerWidth,
+                        LayoutParams.MATCH_PARENT);
+                lp.setMarginStart(padding);
+                lp.setMarginEnd(padding);
+                mMediaCarousel.addView(player.getView(), lp);
+                ((View) mMediaCarousel.getParent()).setVisibility(View.VISIBLE);
+                mMediaPlayers.add(player);
+            }
 
             int iconColor = Color.DKGRAY;
             int bgColor = Color.LTGRAY;
-
-            MediaSession.Token token = browser.getToken();
             player.setMediaSession(token, desc, iconColor, bgColor, browser.getAppIntent(),
                     pkgName);
         }