Merge "MediaRouter2: Updates callback for MediaRouter2Manager"

5 files changed, 277 insertions, 258 deletions

diff --git a/media/java/android/media/IMediaRouter2Manager.aidl b/media/java/android/media/IMediaRouter2Manager.aidl
index b059bd3cec91..b7cb7059ce3d 100644
--- a/media/java/android/media/IMediaRouter2Manager.aidl
+++ b/media/java/android/media/IMediaRouter2Manager.aidl
@@ -25,5 +25,7 @@ import android.media.MediaRoute2Info;
 oneway interface IMediaRouter2Manager {
     void notifyRouteSelected(String packageName, in MediaRoute2Info route);
     void notifyControlCategoriesChanged(String packageName, in List<String> categories);
-    void notifyProviderInfosUpdated(in List<MediaRoute2ProviderInfo> providers);
+    void notifyRoutesAdded(in List<MediaRoute2Info> routes);
+    void notifyRoutesRemoved(in List<MediaRoute2Info> routes);
+    void notifyRoutesChanged(in List<MediaRoute2Info> routes);
 }
diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index cb6c126bc6b1..86fa9db32256 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -95,6 +95,10 @@ public class MediaRouter2 {
     private final String mPackageName;
     private final Map<String, MediaRoute2Info> mRoutes = new HashMap<>();
 
+    //TODO: Use a lock for this to cover the below use case
+    // mRouter.setControlCategories(...);
+    // routes = mRouter.getRoutes();
+    // The current implementation returns empty list
     private volatile List<String> mControlCategories = Collections.emptyList();
 
     private MediaRoute2Info mSelectedRoute;
@@ -202,6 +206,7 @@ public class MediaRouter2 {
                     } catch (RemoteException ex) {
                         Log.e(TAG, "Unable to unregister media router.", ex);
                     }
+                    //TODO: Clean up mRoutes. (onHandler?)
                     mClient = null;
                 }
             }
@@ -222,7 +227,6 @@ public class MediaRouter2 {
                 MediaRouter2.this, new ArrayList<>(controlCategories)));
     }
 
-
     /**
      * Gets the unmodifiable list of {@link MediaRoute2Info routes} currently
      * known to the media router.
diff --git a/media/java/android/media/MediaRouter2Manager.java b/media/java/android/media/MediaRouter2Manager.java
index 7e848a0bdf90..d56dd11fb8f2 100644
--- a/media/java/android/media/MediaRouter2Manager.java
+++ b/media/java/android/media/MediaRouter2Manager.java
@@ -25,18 +25,16 @@ import android.content.Context;
 import android.os.Handler;
 import android.os.RemoteException;
 import android.os.ServiceManager;
-import android.text.TextUtils;
 import android.util.Log;
 
 import com.android.internal.annotations.GuardedBy;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.Collections;
-import java.util.HashSet;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -59,13 +57,11 @@ public class MediaRouter2Manager {
     private Client mClient;
     private final IMediaRouterService mMediaRouterService;
     final Handler mHandler;
-    final List<CallbackRecord> mCallbacks = new CopyOnWriteArrayList<>();
+    final List<CallbackRecord> mCallbackRecords = new CopyOnWriteArrayList<>();
 
-    @SuppressWarnings("WeakerAccess") /* synthetic access */
-    @NonNull
-    List<MediaRoute2ProviderInfo> mProviders = Collections.emptyList();
-    @NonNull
-    List<MediaRoute2Info> mRoutes = Collections.emptyList();
+    private final Object mRoutesLock = new Object();
+    @GuardedBy("mRoutesLock")
+    private final Map<String, MediaRoute2Info> mRoutes = new HashMap<>();
     @NonNull
     final ConcurrentMap<String, List<String>> mControlCategoryMap = new ConcurrentHashMap<>();
 
@@ -104,13 +100,13 @@ public class MediaRouter2Manager {
         Objects.requireNonNull(callback, "callback must not be null");
 
         CallbackRecord callbackRecord;
-        synchronized (mCallbacks) {
+        synchronized (mCallbackRecords) {
             if (findCallbackRecordIndexLocked(callback) >= 0) {
                 Log.w(TAG, "Ignoring to add the same callback twice.");
                 return;
             }
             callbackRecord = new CallbackRecord(executor, callback);
-            mCallbacks.add(callbackRecord);
+            mCallbackRecords.add(callbackRecord);
         }
 
         synchronized (sLock) {
@@ -136,32 +132,32 @@ public class MediaRouter2Manager {
     public void unregisterCallback(@NonNull Callback callback) {
         Objects.requireNonNull(callback, "callback must not be null");
 
-        synchronized (mCallbacks) {
+        synchronized (mCallbackRecords) {
             final int index = findCallbackRecordIndexLocked(callback);
             if (index < 0) {
                 Log.w(TAG, "Ignore removing unknown callback. " + callback);
                 return;
             }
-            mCallbacks.remove(index);
+            mCallbackRecords.remove(index);
             synchronized (sLock) {
-                if (mCallbacks.size() == 0 && mClient != null) {
+                if (mCallbackRecords.size() == 0 && mClient != null) {
                     try {
                         mMediaRouterService.unregisterManager(mClient);
                     } catch (RemoteException ex) {
                         Log.e(TAG, "Unable to unregister media router manager", ex);
                     }
-                    mClient.notifyProviderInfosUpdated(Collections.emptyList());
+                    //TODO: clear mRoutes?
                     mClient = null;
                 }
             }
         }
     }
 
-    @GuardedBy("mCallbacks")
+    @GuardedBy("mCallbackRecords")
     private int findCallbackRecordIndexLocked(Callback callback) {
-        final int count = mCallbacks.size();
+        final int count = mCallbackRecords.size();
         for (int i = 0; i < count; i++) {
-            if (mCallbacks.get(i).mCallback == callback) {
+            if (mCallbackRecords.get(i).mCallback == callback) {
                 return i;
             }
         }
@@ -184,11 +180,14 @@ public class MediaRouter2Manager {
             return Collections.emptyList();
         }
         List<MediaRoute2Info> routes = new ArrayList<>();
-        for (MediaRoute2Info route : mRoutes) {
-            if (route.supportsControlCategory(controlCategories)) {
-                routes.add(route);
+        synchronized (mRoutesLock) {
+            for (MediaRoute2Info route : mRoutes.values()) {
+                if (route.supportsControlCategory(controlCategories)) {
+                    routes.add(route);
+                }
             }
         }
+        //TODO: Should we cache this?
         return routes;
     }
 
@@ -282,136 +281,96 @@ public class MediaRouter2Manager {
         }
     }
 
-    int findProviderIndex(MediaRoute2ProviderInfo provider) {
-        final int count = mProviders.size();
-        for (int i = 0; i < count; i++) {
-            if (TextUtils.equals(mProviders.get(i).getUniqueId(), provider.getUniqueId())) {
-                return i;
+    void addRoutesOnHandler(List<MediaRoute2Info> routes) {
+        synchronized (mRoutesLock) {
+            for (MediaRoute2Info route : routes) {
+                mRoutes.put(route.getUniqueId(), route);
             }
         }
-        return -1;
-    }
-
-    void updateProvider(@NonNull MediaRoute2ProviderInfo provider) {
-        if (provider == null || !provider.isValid()) {
-            Log.w(TAG, "Ignoring invalid provider : " + provider);
-            return;
+        if (routes.size() > 0) {
+            notifyRoutesAdded(routes);
         }
+    }
 
-        final Collection<MediaRoute2Info> routes = provider.getRoutes();
-
-        final int index = findProviderIndex(provider);
-        if (index >= 0) {
-            final MediaRoute2ProviderInfo prevProvider = mProviders.get(index);
-            final Set<String> updatedRouteIds = new HashSet<>();
-            for (MediaRoute2Info routeInfo : routes) {
-                if (!routeInfo.isValid()) {
-                    Log.w(TAG, "Ignoring invalid route : " + routeInfo);
-                    continue;
-                }
-                final MediaRoute2Info prevRoute = prevProvider.getRoute(routeInfo.getId());
-                if (prevRoute == null) {
-                    notifyRouteAdded(routeInfo);
-                } else {
-                    if (!Objects.equals(prevRoute, routeInfo)) {
-                        notifyRouteChanged(routeInfo);
-                    }
-                    updatedRouteIds.add(routeInfo.getId());
-                }
-            }
-            final Collection<MediaRoute2Info> prevRoutes = prevProvider.getRoutes();
-
-            for (MediaRoute2Info prevRoute : prevRoutes) {
-                if (!updatedRouteIds.contains(prevRoute.getId())) {
-                    notifyRouteRemoved(prevRoute);
-                }
-            }
-        } else {
-            for (MediaRoute2Info routeInfo: routes) {
-                notifyRouteAdded(routeInfo);
+    void removeRoutesOnHandler(List<MediaRoute2Info> routes) {
+        synchronized (mRoutesLock) {
+            for (MediaRoute2Info route : routes) {
+                mRoutes.remove(route.getUniqueId());
             }
         }
-    }
-
-    void notifyRouteAdded(MediaRoute2Info routeInfo) {
-        for (CallbackRecord record : mCallbacks) {
-            record.mExecutor.execute(
-                    () -> record.mCallback.onRouteAdded(routeInfo));
+        if (routes.size() > 0) {
+            notifyRoutesRemoved(routes);
         }
     }
 
-    void notifyRouteChanged(MediaRoute2Info routeInfo) {
-        for (CallbackRecord record : mCallbacks) {
-            record.mExecutor.execute(
-                    () -> record.mCallback.onRouteChanged(routeInfo));
+    void changeRoutesOnHandler(List<MediaRoute2Info> routes) {
+        synchronized (mRoutesLock) {
+            for (MediaRoute2Info route : routes) {
+                mRoutes.put(route.getUniqueId(), route);
+            }
+        }
+        if (routes.size() > 0) {
+            notifyRoutesChanged(routes);
         }
     }
 
-    void notifyRouteRemoved(MediaRoute2Info routeInfo) {
-        for (CallbackRecord record : mCallbacks) {
+    private void notifyRoutesAdded(List<MediaRoute2Info> routes) {
+        for (CallbackRecord record: mCallbackRecords) {
             record.mExecutor.execute(
-                    () -> record.mCallback.onRouteRemoved(routeInfo));
+                    () -> record.mCallback.onRoutesAdded(routes));
         }
     }
 
-    void notifyRouteListChanged() {
-        for (CallbackRecord record: mCallbacks) {
+    private void notifyRoutesRemoved(List<MediaRoute2Info> routes) {
+        for (CallbackRecord record: mCallbackRecords) {
             record.mExecutor.execute(
-                    () -> record.mCallback.onRoutesChanged(mRoutes));
+                    () -> record.mCallback.onRoutesRemoved(routes));
         }
     }
 
-    void notifyProviderInfosUpdated(List<MediaRoute2ProviderInfo> providers) {
-        if (providers == null) {
-            Log.w(TAG, "Providers info is null.");
-            return;
-        }
-
-        ArrayList<MediaRoute2Info> routes = new ArrayList<>();
-
-        for (MediaRoute2ProviderInfo provider : providers) {
-            updateProvider(provider);
-            //TODO: Should we do this in updateProvider()?
-            routes.addAll(provider.getRoutes());
+    private void notifyRoutesChanged(List<MediaRoute2Info> routes) {
+        for (CallbackRecord record: mCallbackRecords) {
+            record.mExecutor.execute(
+                    () -> record.mCallback.onRoutesChanged(routes));
         }
-        //TODO: Call notifyRouteRemoved for the routes of the removed providers.
-
-        //TODO: Filter invalid providers and invalid routes.
-        mProviders = providers;
-        mRoutes = routes;
-
-        //TODO: Call this when only the list is modified.
-        notifyRouteListChanged();
     }
 
     void notifyRouteSelected(String packageName, MediaRoute2Info route) {
-        for (CallbackRecord record : mCallbacks) {
+        for (CallbackRecord record : mCallbackRecords) {
             record.mExecutor.execute(() -> record.mCallback.onRouteSelected(packageName, route));
         }
     }
 
     void updateControlCategories(String packageName, List<String> categories) {
         mControlCategoryMap.put(packageName, categories);
+        for (CallbackRecord record : mCallbackRecords) {
+            record.mExecutor.execute(
+                    () -> record.mCallback.onControlCategoriesChanged(packageName));
+        }
     }
 
     /**
      * Interface for receiving events about media routing changes.
      */
     public static class Callback {
+
         /**
-         * Called when a route is added.
+         * Called when routes are added.
+         * @param routes the list of routes that have been added. It's never empty.
          */
-        public void onRouteAdded(@NonNull MediaRoute2Info routeInfo) {}
+        public void onRoutesAdded(@NonNull List<MediaRoute2Info> routes) {}
 
         /**
-         * Called when a route is changed.
+         * Called when routes are removed.
+         * @param routes the list of routes that have been removed. It's never empty.
          */
-        public void onRouteChanged(@NonNull MediaRoute2Info routeInfo) {}
+        public void onRoutesRemoved(@NonNull List<MediaRoute2Info> routes) {}
 
         /**
-         * Called when a route is removed.
+         * Called when routes are changed.
+         * @param routes the list of routes that have been changed. It's never empty.
          */
-        public void onRouteRemoved(@NonNull MediaRoute2Info routeInfo) {}
+        public void onRoutesChanged(@NonNull List<MediaRoute2Info> routes) {}
 
         /**
          * Called when a route is selected for an application.
@@ -422,13 +381,13 @@ public class MediaRouter2Manager {
          */
         public void onRouteSelected(@NonNull String packageName, @Nullable MediaRoute2Info route) {}
 
+
         /**
-         * Called when the list of routes is changed.
-         * A client may refresh available routes for each application.
+         * Called when the control categories of an app is changed.
+         *
+         * @param packageName the package name of the application
          */
-        public void onRoutesChanged(@NonNull List<MediaRoute2Info> routes) {}
-
-        //TODO: add onControlCategoriesChanged to notify available routes are changed
+        public void onControlCategoriesChanged(@NonNull String packageName) {}
     }
 
     final class CallbackRecord {
@@ -441,10 +400,12 @@ public class MediaRouter2Manager {
         }
 
         void notifyRoutes() {
-            mExecutor.execute(() -> mCallback.onRoutesChanged(mRoutes));
-            for (MediaRoute2Info routeInfo : mRoutes) {
-                mExecutor.execute(
-                        () -> mCallback.onRouteAdded(routeInfo));
+            List<MediaRoute2Info> routes;
+            synchronized (mRoutesLock) {
+                routes = new ArrayList<>(mRoutes.values());
+            }
+            if (routes.size() > 0) {
+                mExecutor.execute(() -> mCallback.onRoutesAdded(routes));
             }
         }
     }
@@ -463,9 +424,21 @@ public class MediaRouter2Manager {
         }
 
         @Override
-        public void notifyProviderInfosUpdated(List<MediaRoute2ProviderInfo> info) {
-            mHandler.sendMessage(obtainMessage(MediaRouter2Manager::notifyProviderInfosUpdated,
-                    MediaRouter2Manager.this, info));
+        public void notifyRoutesAdded(List<MediaRoute2Info> routes) {
+            mHandler.sendMessage(obtainMessage(MediaRouter2Manager::addRoutesOnHandler,
+                    MediaRouter2Manager.this, routes));
+        }
+
+        @Override
+        public void notifyRoutesRemoved(List<MediaRoute2Info> routes) {
+            mHandler.sendMessage(obtainMessage(MediaRouter2Manager::removeRoutesOnHandler,
+                    MediaRouter2Manager.this, routes));
+        }
+
+        @Override
+        public void notifyRoutesChanged(List<MediaRoute2Info> routes) {
+            mHandler.sendMessage(obtainMessage(MediaRouter2Manager::changeRoutesOnHandler,
+                    MediaRouter2Manager.this, routes));
         }
     }
 }
diff --git a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java
index acf899881600..240e7d1b6357 100644
--- a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java
+++ b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouterManagerTest.java
@@ -24,6 +24,7 @@ import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
@@ -31,7 +32,6 @@ import static org.mockito.Mockito.verify;
 import android.content.Context;
 import android.content.Intent;
 import android.media.MediaRoute2Info;
-import android.media.MediaRouter;
 import android.media.MediaRouter2;
 import android.media.MediaRouter2Manager;
 import android.support.test.InstrumentationRegistry;
@@ -52,6 +52,7 @@ import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Predicate;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -85,7 +86,6 @@ public class MediaRouterManagerTest {
 
     private Context mContext;
     private MediaRouter2Manager mManager;
-    private MediaRouter mRouter;
     private MediaRouter2 mRouter2;
     private Executor mExecutor;
     private String mPackageName;
@@ -103,7 +103,6 @@ public class MediaRouterManagerTest {
     public void setUp() throws Exception {
         mContext = InstrumentationRegistry.getTargetContext();
         mManager = MediaRouter2Manager.getInstance(mContext);
-        mRouter = (MediaRouter) mContext.getSystemService(Context.MEDIA_ROUTER_SERVICE);
         mRouter2 = MediaRouter2.getInstance(mContext);
         //TODO: If we need to support thread pool executors, change this to thread pool executor.
         mExecutor = Executors.newSingleThreadExecutor();
@@ -124,64 +123,58 @@ public class MediaRouterManagerTest {
         assertNotEquals(routeInfo1, routeInfo3);
     }
 
-    //TODO: Test onRouteChanged when it's properly implemented.
     @Test
-    public void testRouteAdded() {
-        MediaRouter2Manager.Callback mockCallback = mock(MediaRouter2Manager.Callback.class);
-
-        mManager.registerCallback(mExecutor, mockCallback);
-
-        verify(mockCallback, timeout(TIMEOUT_MS)).onRouteAdded(argThat(
-                (MediaRoute2Info info) ->
-                        info.getId().equals(ROUTE_ID1) && info.getName().equals(ROUTE_NAME1)));
-        mManager.unregisterCallback(mockCallback);
-    }
-
-    //TODO: Recover this test when media router 2 is finalized.
-    public void testRouteRemoved() {
-        MediaRouter2Manager.Callback mockCallback = mock(MediaRouter2Manager.Callback.class);
-        mManager.registerCallback(mExecutor, mockCallback);
+    public void testOnRoutesAdded() throws Exception {
+        CountDownLatch latch = new CountDownLatch(1);
+        MediaRouter2Manager.Callback callback = new MediaRouter2Manager.Callback() {
+            @Override
+            public void onRoutesAdded(List<MediaRoute2Info> routes) {
+                assertTrue(routes.size() > 0);
+                for (MediaRoute2Info route : routes) {
+                    if (route.getId().equals(ROUTE_ID1) && route.getName().equals(ROUTE_NAME1)) {
+                        latch.countDown();
+                    }
+                }
+            }
+        };
+        mManager.registerCallback(mExecutor, callback);
 
-        MediaRouter2.Callback mockRouterCallback = mock(MediaRouter2.Callback.class);
+        assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
 
-        //TODO: Figure out a more proper way to test.
-        // (Control requests shouldn't be used in this way.)
-        mRouter2.setControlCategories(CONTROL_CATEGORIES_ALL);
-        mRouter2.registerCallback(mExecutor, mockRouterCallback);
-        mRouter2.sendControlRequest(
-                new MediaRoute2Info.Builder(ROUTE_ID2, ROUTE_NAME2).build(),
-                new Intent(ACTION_REMOVE_ROUTE));
-        mRouter2.unregisterCallback(mockRouterCallback);
-
-        verify(mockCallback, timeout(TIMEOUT_MS)).onRouteRemoved(argThat(
-                (MediaRoute2Info info) ->
-                        info.getId().equals(ROUTE_ID2) && info.getName().equals(ROUTE_NAME2)));
-        mManager.unregisterCallback(mockCallback);
+        mManager.unregisterCallback(callback);
     }
 
-    /**
-     * Tests if we get proper routes for application that has special control category.
-     */
     @Test
-    public void testControlCategoryWithMediaRouter() throws Exception {
+    public void testOnRoutesRemoved() throws Exception {
         MediaRouter2Manager.Callback mockCallback = mock(MediaRouter2Manager.Callback.class);
         mManager.registerCallback(mExecutor, mockCallback);
 
-        MediaRouter.Callback mockRouterCallback = mock(MediaRouter.Callback.class);
-
-        mRouter.setControlCategories(CONTROL_CATEGORIES_SPECIAL);
-        mRouter.addCallback(MediaRouter.ROUTE_TYPE_USER, mockRouterCallback);
-
-        verify(mockCallback, timeout(TIMEOUT_MS))
-                .onRoutesChanged(argThat(routes -> routes.size() > 0));
+        MediaRouter2.Callback routerCallback = new MediaRouter2.Callback();
+        mRouter2.registerCallback(mExecutor, routerCallback);
 
         Map<String, MediaRoute2Info> routes =
-                createRouteMap(mManager.getAvailableRoutes(mPackageName));
+                waitAndGetRoutesWithManager(CONTROL_CATEGORIES_ALL);
 
-        Assert.assertEquals(1, routes.size());
-        Assert.assertNotNull(routes.get(ROUTE_ID_SPECIAL_CATEGORY));
+        CountDownLatch latch = new CountDownLatch(1);
+        MediaRouter2Manager.Callback callback = new MediaRouter2Manager.Callback() {
+            @Override
+            public void onRoutesRemoved(List<MediaRoute2Info> routes) {
+                assertTrue(routes.size() > 0);
+                for (MediaRoute2Info route : routes) {
+                    if (route.getId().equals(ROUTE_ID2) && route.getName().equals(ROUTE_NAME2)) {
+                        latch.countDown();
+                    }
+                }
+            }
+        };
+        mManager.registerCallback(mExecutor, callback);
 
-        mRouter.removeCallback(mockRouterCallback);
+        //TODO: Figure out a more proper way to test.
+        // (Control requests shouldn't be used in this way.)
+        mRouter2.sendControlRequest(routes.get(ROUTE_ID2), new Intent(ACTION_REMOVE_ROUTE));
+        assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+
+        mRouter2.unregisterCallback(routerCallback);
         mManager.unregisterCallback(mockCallback);
     }
 
@@ -233,7 +226,8 @@ public class MediaRouterManagerTest {
 
         mManager.selectRoute(mPackageName, routeToSelect);
         verify(managerCallback, timeout(TIMEOUT_MS))
-                .onRouteAdded(argThat(route -> route.equals(routeToSelect)));
+                .onRouteSelected(eq(mPackageName),
+                        argThat(route -> route != null && route.equals(routeToSelect)));
 
         mRouter2.unregisterCallback(routerCallback);
         mManager.unregisterCallback(managerCallback);
@@ -244,31 +238,28 @@ public class MediaRouterManagerTest {
      */
     @Test
     public void testSingleProviderSelect() throws Exception {
-        MediaRouter2Manager.Callback managerCallback = mock(MediaRouter2Manager.Callback.class);
         MediaRouter2.Callback routerCallback = mock(MediaRouter2.Callback.class);
 
-        mManager.registerCallback(mExecutor, managerCallback);
         mRouter2.registerCallback(mExecutor, routerCallback);
 
         Map<String, MediaRoute2Info> routes = waitAndGetRoutesWithManager(CONTROL_CATEGORIES_ALL);
 
-        mManager.selectRoute(mPackageName, routes.get(ROUTE_ID1));
-        verify(managerCallback, timeout(TIMEOUT_MS))
-                .onRouteChanged(argThat(routeInfo -> TextUtils.equals(ROUTE_ID1, routeInfo.getId())
-                        && TextUtils.equals(routeInfo.getClientPackageName(), mPackageName)));
+        awaitOnRouteChangedManager(
+                () -> mManager.selectRoute(mPackageName, routes.get(ROUTE_ID1)),
+                ROUTE_ID1,
+                route -> TextUtils.equals(route.getClientPackageName(), mPackageName));
 
-        mManager.selectRoute(mPackageName, routes.get(ROUTE_ID2));
-        verify(managerCallback, timeout(TIMEOUT_MS))
-                .onRouteChanged(argThat(routeInfo -> TextUtils.equals(ROUTE_ID2, routeInfo.getId())
-                        && TextUtils.equals(routeInfo.getClientPackageName(), mPackageName)));
+        awaitOnRouteChangedManager(
+                () -> mManager.selectRoute(mPackageName, routes.get(ROUTE_ID2)),
+                ROUTE_ID2,
+                route -> TextUtils.equals(route.getClientPackageName(), mPackageName));
 
-        mManager.unselectRoute(mPackageName);
-        verify(managerCallback, timeout(TIMEOUT_MS))
-                .onRouteChanged(argThat(routeInfo -> TextUtils.equals(ROUTE_ID2, routeInfo.getId())
-                        && TextUtils.equals(routeInfo.getClientPackageName(), null)));
+        awaitOnRouteChangedManager(
+                () -> mManager.unselectRoute(mPackageName),
+                ROUTE_ID2,
+                route -> TextUtils.equals(route.getClientPackageName(), null));
 
         mRouter2.unregisterCallback(routerCallback);
-        mManager.unregisterCallback(managerCallback);
     }
 
     @Test
@@ -279,52 +270,39 @@ public class MediaRouterManagerTest {
         int originalVolume = volRoute.getVolume();
         int deltaVolume = (originalVolume == volRoute.getVolumeMax() ? -1 : 1);
 
-        CountDownLatch latch1 = new CountDownLatch(1);
-        MediaRouter2.Callback callback1 =
-                createVolumeChangeCallback(ROUTE_ID_VARIABLE_VOLUME,
-                        originalVolume + deltaVolume, latch1);
-        mRouter2.registerCallback(mExecutor, callback1);
-        mRouter2.requestUpdateVolume(volRoute, deltaVolume);
-        assertTrue(latch1.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
-        mRouter2.unregisterCallback(callback1);
-
-        CountDownLatch latch2 = new CountDownLatch(1);
-        MediaRouter2.Callback callback2 =
-                createVolumeChangeCallback(ROUTE_ID_VARIABLE_VOLUME, originalVolume, latch2);
-        mRouter2.registerCallback(mExecutor, callback2);
-        mRouter2.requestSetVolume(volRoute, originalVolume);
-        assertTrue(latch1.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
-        mRouter2.unregisterCallback(callback1);
+        awaitOnRouteChanged(
+                () -> mRouter2.requestUpdateVolume(volRoute, deltaVolume),
+                ROUTE_ID_VARIABLE_VOLUME,
+                (route -> route.getVolume() == originalVolume + deltaVolume));
+
+        awaitOnRouteChanged(
+                () -> mRouter2.requestSetVolume(volRoute, originalVolume),
+                ROUTE_ID_VARIABLE_VOLUME,
+                (route -> route.getVolume() == originalVolume));
     }
 
     @Test
     public void testControlVolumeWithManager() throws Exception {
-        MediaRouter2Manager.Callback managerCallback = mock(MediaRouter2Manager.Callback.class);
         MediaRouter2.Callback mockCallback = mock(MediaRouter2.Callback.class);
 
-        mManager.registerCallback(mExecutor, managerCallback);
         mRouter2.registerCallback(mExecutor, mockCallback);
         Map<String, MediaRoute2Info> routes = waitAndGetRoutesWithManager(CONTROL_CATEGORIES_ALL);
 
         MediaRoute2Info volRoute = routes.get(ROUTE_ID_VARIABLE_VOLUME);
         int originalVolume = volRoute.getVolume();
         int deltaVolume = (originalVolume == volRoute.getVolumeMax() ? -1 : 1);
-        int targetVolume = originalVolume + deltaVolume;
 
-        mManager.requestSetVolume(volRoute, targetVolume);
-        verify(managerCallback, timeout(TIMEOUT_MS).atLeastOnce())
-                .onRouteChanged(argThat(route ->
-                        route.getId().equals(volRoute.getId())
-                                && route.getVolume() == targetVolume));
+        awaitOnRouteChangedManager(
+                () -> mManager.requestUpdateVolume(volRoute, deltaVolume),
+                ROUTE_ID_VARIABLE_VOLUME,
+                (route -> route.getVolume() == originalVolume + deltaVolume));
 
-        mManager.requestUpdateVolume(volRoute, -deltaVolume);
-        verify(managerCallback, timeout(TIMEOUT_MS).atLeastOnce())
-                .onRouteChanged(argThat(route ->
-                        route.getId().equals(volRoute.getId())
-                                && route.getVolume() == originalVolume));
+        awaitOnRouteChangedManager(
+                () -> mManager.requestSetVolume(volRoute, originalVolume),
+                ROUTE_ID_VARIABLE_VOLUME,
+                (route -> route.getVolume() == originalVolume));
 
         mRouter2.unregisterCallback(mockCallback);
-        mManager.unregisterCallback(managerCallback);
     }
 
     @Test
@@ -363,25 +341,29 @@ public class MediaRouterManagerTest {
 
     Map<String, MediaRoute2Info> waitAndGetRoutesWithManager(List<String> controlCategories)
             throws Exception {
-        CountDownLatch latch = new CountDownLatch(1);
+        CountDownLatch latch = new CountDownLatch(2);
 
-        // Dummy callback is required to send control category info.
+        // A dummy callback is required to send control category info.
         MediaRouter2.Callback routerCallback = new MediaRouter2.Callback();
         MediaRouter2Manager.Callback managerCallback = new MediaRouter2Manager.Callback() {
             @Override
-            public void onRoutesChanged(List<MediaRoute2Info> routes) {
+            public void onRoutesAdded(List<MediaRoute2Info> routes) {
                 if (routes.size() > 0) {
                     latch.countDown();
                 }
             }
+            @Override
+            public void onControlCategoriesChanged(String packageName) {
+                if (TextUtils.equals(mPackageName, packageName)) {
+                    latch.countDown();
+                }
+            }
         };
         mManager.registerCallback(mExecutor, managerCallback);
         mRouter2.setControlCategories(controlCategories);
         mRouter2.registerCallback(mExecutor, routerCallback);
         try {
             assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
-            //TODO: currently this returns empty list occasionally.
-            //Maybe due to control category is not set yet
             return createRouteMap(mManager.getAvailableRoutes(mPackageName));
         } finally {
             mRouter2.unregisterCallback(routerCallback);
@@ -389,18 +371,46 @@ public class MediaRouterManagerTest {
         }
     }
 
-    MediaRouter2.Callback createVolumeChangeCallback(String routeId,
-            int targetVolume, CountDownLatch latch) {
+    void awaitOnRouteChanged(Runnable task, String routeId,
+            Predicate<MediaRoute2Info> predicate) throws Exception {
+        CountDownLatch latch = new CountDownLatch(1);
         MediaRouter2.Callback callback = new MediaRouter2.Callback() {
             @Override
             public void onRoutesChanged(List<MediaRoute2Info> changed) {
-                MediaRoute2Info volRoute = createRouteMap(changed).get(routeId);
-                if (volRoute != null && volRoute.getVolume() == targetVolume) {
+                MediaRoute2Info route = createRouteMap(changed).get(routeId);
+                if (route != null && predicate.test(route)) {
                     latch.countDown();
                 }
             }
         };
-        return callback;
+        mRouter2.registerCallback(mExecutor, callback);
+        try {
+            new Thread(task).start();
+            assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        } finally {
+            mRouter2.unregisterCallback(callback);
+        }
+    }
+
+    void awaitOnRouteChangedManager(Runnable task, String routeId,
+            Predicate<MediaRoute2Info> predicate) throws Exception {
+        CountDownLatch latch = new CountDownLatch(1);
+        MediaRouter2Manager.Callback callback = new MediaRouter2Manager.Callback() {
+            @Override
+            public void onRoutesChanged(List<MediaRoute2Info> changed) {
+                MediaRoute2Info route = createRouteMap(changed).get(routeId);
+                if (route != null && predicate.test(route)) {
+                    latch.countDown();
+                }
+            }
+        };
+        mManager.registerCallback(mExecutor, callback);
+        try {
+            new Thread(task).start();
+            assertTrue(latch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        } finally {
+            mManager.unregisterCallback(callback);
+        }
     }
 
     // Helper for getting routes easily
diff --git a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
index 44642d4f85cd..2e97f6a559d6 100644
--- a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
+++ b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
@@ -469,7 +469,7 @@ class MediaRouter2ServiceImpl {
             mAllManagerRecords.put(binder, managerRecord);
 
             userRecord.mHandler.sendMessage(
-                    obtainMessage(UserHandler::notifyProviderInfosUpdatedToManager,
+                    obtainMessage(UserHandler::notifyRoutesToManager,
                             userRecord.mHandler, manager));
 
             for (ClientRecord clientRecord : userRecord.mClientRecords) {
@@ -808,17 +808,20 @@ class MediaRouter2ServiceImpl {
                 }
 
                 List<IMediaRouter2Client> clients = getClients();
+                List<IMediaRouter2Manager> managers = getManagers();
                 if (addedRoutes.size() > 0) {
                     notifyRoutesAddedToClients(clients, addedRoutes);
+                    notifyRoutesAddedToManagers(managers, addedRoutes);
                 }
                 if (removedRoutes.size() > 0) {
                     notifyRoutesRemovedToClients(clients, removedRoutes);
+                    notifyRoutesRemovedToManagers(managers, removedRoutes);
                 }
                 if (changedRoutes.size() > 0) {
                     notifyRoutesChangedToClients(clients, changedRoutes);
+                    notifyRoutesChangedToManagers(managers, changedRoutes);
                 }
             }
-            scheduleUpdateProviderInfos();
         }
 
         private int getProviderInfoIndex(String providerId) {
@@ -874,33 +877,6 @@ class MediaRouter2ServiceImpl {
             }
         }
 
-        private void scheduleUpdateProviderInfos() {
-            if (!mProviderInfosUpdateScheduled) {
-                mProviderInfosUpdateScheduled = true;
-                sendMessage(PooledLambda.obtainMessage(UserHandler::updateProviderInfos, this));
-            }
-        }
-
-        //TODO: should be replaced into notifyRoutes...ToManagers
-        private void updateProviderInfos() {
-            mProviderInfosUpdateScheduled = false;
-
-            MediaRouter2ServiceImpl service = mServiceRef.get();
-            if (service == null) {
-                return;
-            }
-
-            final List<IMediaRouter2Manager> managers = new ArrayList<>();
-            synchronized (service.mLock) {
-                for (ManagerRecord managerRecord : mUserRecord.mManagerRecords) {
-                    managers.add(managerRecord.mManager);
-                }
-            }
-            for (IMediaRouter2Manager manager : managers) {
-                notifyProviderInfosUpdatedToManager(manager);
-            }
-        }
-
         private List<IMediaRouter2Client> getClients() {
             final List<IMediaRouter2Client> clients = new ArrayList<>();
             MediaRouter2ServiceImpl service = mServiceRef.get();
@@ -917,6 +893,20 @@ class MediaRouter2ServiceImpl {
             return clients;
         }
 
+        private List<IMediaRouter2Manager> getManagers() {
+            final List<IMediaRouter2Manager> managers = new ArrayList<>();
+            MediaRouter2ServiceImpl service = mServiceRef.get();
+            if (service == null) {
+                return managers;
+            }
+            synchronized (service.mLock) {
+                for (ManagerRecord managerRecord : mUserRecord.mManagerRecords) {
+                    managers.add(managerRecord.mManager);
+                }
+            }
+            return managers;
+        }
+
         private void notifyRoutesToClient(IMediaRouter2Client client) {
             List<MediaRoute2Info> routes = new ArrayList<>();
             for (MediaRoute2ProviderInfo providerInfo : mProviderInfos) {
@@ -965,11 +955,51 @@ class MediaRouter2ServiceImpl {
             }
         }
 
-        private void notifyProviderInfosUpdatedToManager(IMediaRouter2Manager manager) {
+        private void notifyRoutesToManager(IMediaRouter2Manager manager) {
+            List<MediaRoute2Info> routes = new ArrayList<>();
+            for (MediaRoute2ProviderInfo providerInfo : mProviderInfos) {
+                routes.addAll(providerInfo.getRoutes());
+            }
+            if (routes.size() == 0) {
+                return;
+            }
             try {
-                manager.notifyProviderInfosUpdated(mProviderInfos);
+                manager.notifyRoutesAdded(routes);
             } catch (RemoteException ex) {
-                Slog.w(TAG, "Failed to notify provider infos updated. Manager probably died.");
+                Slog.w(TAG, "Failed to notify all routes. Manager probably died.", ex);
+            }
+        }
+
+        private void notifyRoutesAddedToManagers(List<IMediaRouter2Manager> managers,
+                List<MediaRoute2Info> routes) {
+            for (IMediaRouter2Manager manager : managers) {
+                try {
+                    manager.notifyRoutesAdded(routes);
+                } catch (RemoteException ex) {
+                    Slog.w(TAG, "Failed to notify routes added. Manager probably died.", ex);
+                }
+            }
+        }
+
+        private void notifyRoutesRemovedToManagers(List<IMediaRouter2Manager> managers,
+                List<MediaRoute2Info> routes) {
+            for (IMediaRouter2Manager manager : managers) {
+                try {
+                    manager.notifyRoutesRemoved(routes);
+                } catch (RemoteException ex) {
+                    Slog.w(TAG, "Failed to notify routes removed. Manager probably died.", ex);
+                }
+            }
+        }
+
+        private void notifyRoutesChangedToManagers(List<IMediaRouter2Manager> managers,
+                List<MediaRoute2Info> routes) {
+            for (IMediaRouter2Manager manager : managers) {
+                try {
+                    manager.notifyRoutesChanged(routes);
+                } catch (RemoteException ex) {
+                    Slog.w(TAG, "Failed to notify routes changed. Manager probably died.", ex);
+                }
             }
         }