Merge "MediaRouter2: Clean up APIs"
diff --git a/api/current.txt b/api/current.txt
index dd06fc3..7a5753a 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -27020,12 +27020,19 @@
     method @NonNull public static android.media.MediaRouter2 getInstance(@NonNull android.content.Context);
     method @NonNull public java.util.List<android.media.MediaRoute2Info> getRoutes();
     method @NonNull public android.media.MediaRouter2.RoutingController getSystemController();
-    method public void registerControllerCallback(@NonNull java.util.concurrent.Executor, @NonNull android.media.MediaRouter2.RoutingControllerCallback);
+    method public void registerControllerCallback(@NonNull java.util.concurrent.Executor, @NonNull android.media.MediaRouter2.ControllerCallback);
     method public void registerRouteCallback(@NonNull java.util.concurrent.Executor, @NonNull android.media.MediaRouter2.RouteCallback, @NonNull android.media.RouteDiscoveryPreference);
-    method public void requestCreateController(@NonNull android.media.MediaRoute2Info);
+    method public void registerTransferCallback(@NonNull java.util.concurrent.Executor, @NonNull android.media.MediaRouter2.TransferCallback);
     method public void setOnGetControllerHintsListener(@Nullable android.media.MediaRouter2.OnGetControllerHintsListener);
-    method public void unregisterControllerCallback(@NonNull android.media.MediaRouter2.RoutingControllerCallback);
+    method public void transferTo(@Nullable android.media.MediaRoute2Info);
+    method public void unregisterControllerCallback(@NonNull android.media.MediaRouter2.ControllerCallback);
     method public void unregisterRouteCallback(@NonNull android.media.MediaRouter2.RouteCallback);
+    method public void unregisterTransferCallback(@NonNull android.media.MediaRouter2.TransferCallback);
+  }
+
+  public static class MediaRouter2.ControllerCallback {
+    ctor public MediaRouter2.ControllerCallback();
+    method public void onControllerUpdated(@NonNull android.media.MediaRouter2.RoutingController);
   }
 
   public static interface MediaRouter2.OnGetControllerHintsListener {
@@ -27046,7 +27053,6 @@
     method @NonNull public String getId();
     method @NonNull public java.util.List<android.media.MediaRoute2Info> getSelectableRoutes();
     method @NonNull public java.util.List<android.media.MediaRoute2Info> getSelectedRoutes();
-    method @NonNull public java.util.List<android.media.MediaRoute2Info> getTransferableRoutes();
     method public int getVolume();
     method public int getVolumeHandling();
     method public int getVolumeMax();
@@ -27054,15 +27060,12 @@
     method public void release();
     method public void selectRoute(@NonNull android.media.MediaRoute2Info);
     method public void setVolume(int);
-    method public void transferToRoute(@NonNull android.media.MediaRoute2Info);
   }
 
-  public static class MediaRouter2.RoutingControllerCallback {
-    ctor public MediaRouter2.RoutingControllerCallback();
-    method public void onControllerCreated(@NonNull android.media.MediaRouter2.RoutingController);
-    method public void onControllerCreationFailed(@NonNull android.media.MediaRoute2Info);
-    method public void onControllerReleased(@NonNull android.media.MediaRouter2.RoutingController);
-    method public void onControllerUpdated(@NonNull android.media.MediaRouter2.RoutingController);
+  public static class MediaRouter2.TransferCallback {
+    ctor public MediaRouter2.TransferCallback();
+    method public void onTransferFailed(@NonNull android.media.MediaRoute2Info);
+    method public void onTransferred(@NonNull android.media.MediaRouter2.RoutingController, @Nullable android.media.MediaRouter2.RoutingController);
   }
 
   public class MediaScannerConnection implements android.content.ServiceConnection {
diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index 5111411..6281ccd 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -65,7 +65,8 @@
 
     private final CopyOnWriteArrayList<RouteCallbackRecord> mRouteCallbackRecords =
             new CopyOnWriteArrayList<>();
-
+    private final CopyOnWriteArrayList<TransferCallbackRecord> mTransferCallbackRecords =
+            new CopyOnWriteArrayList<>();
     private final CopyOnWriteArrayList<ControllerCallbackRecord> mControllerCallbackRecords =
             new CopyOnWriteArrayList<>();
 
@@ -280,22 +281,21 @@
     }
 
     /**
-     * Registers a callback to get updates on creations and changes of
-     * {@link RoutingController routing controllers}.
+     * Registers a callback to get the result of {@link #transferTo(MediaRoute2Info)}.
      * If you register the same callback twice or more, it will be ignored.
      *
      * @param executor the executor to execute the callback on
      * @param callback the callback to register
-     * @see #unregisterControllerCallback
+     * @see #unregisterTransferCallback
      */
-    public void registerControllerCallback(@NonNull @CallbackExecutor Executor executor,
-            @NonNull RoutingControllerCallback callback) {
+    public void registerTransferCallback(@NonNull @CallbackExecutor Executor executor,
+            @NonNull TransferCallback callback) {
         Objects.requireNonNull(executor, "executor must not be null");
         Objects.requireNonNull(callback, "callback must not be null");
 
-        ControllerCallbackRecord record = new ControllerCallbackRecord(executor, callback);
-        if (!mControllerCallbackRecords.addIfAbsent(record)) {
-            Log.w(TAG, "Ignoring the same controller callback");
+        TransferCallbackRecord record = new TransferCallbackRecord(executor, callback);
+        if (!mTransferCallbackRecords.addIfAbsent(record)) {
+            Log.w(TAG, "registerTransferCallback: Ignoring the same callback");
             return;
         }
     }
@@ -305,13 +305,45 @@
      * If the callback has not been added or been removed already, it is ignored.
      *
      * @param callback the callback to unregister
-     * @see #registerControllerCallback
+     * @see #registerTransferCallback
      */
-    public void unregisterControllerCallback(@NonNull RoutingControllerCallback callback) {
+    public void unregisterTransferCallback(@NonNull TransferCallback callback) {
+        Objects.requireNonNull(callback, "callback must not be null");
+
+        if (!mTransferCallbackRecords.remove(new TransferCallbackRecord(null, callback))) {
+            Log.w(TAG, "unregisterTransferCallback: Ignoring an unknown callback");
+            return;
+        }
+    }
+
+    /**
+     * Registers a {@link ControllerCallback}.
+     * If you register the same callback twice or more, it will be ignored.
+     * @see #unregisterControllerCallback(ControllerCallback)
+     */
+    public void registerControllerCallback(@NonNull @CallbackExecutor Executor executor,
+            @NonNull ControllerCallback callback) {
+        Objects.requireNonNull(executor, "executor must not be null");
+        Objects.requireNonNull(callback, "callback must not be null");
+
+        ControllerCallbackRecord record = new ControllerCallbackRecord(executor, callback);
+        if (!mControllerCallbackRecords.addIfAbsent(record)) {
+            Log.w(TAG, "registerControllerCallback: Ignoring the same callback");
+            return;
+        }
+    }
+
+    /**
+     * Unregisters a {@link ControllerCallback}. The callback will no longer receive
+     * events. If the callback has not been added or been removed already, it is ignored.
+     * @see #registerControllerCallback(Executor, ControllerCallback)
+     */
+    public void unregisterControllerCallback(
+            @NonNull ControllerCallback callback) {
         Objects.requireNonNull(callback, "callback must not be null");
 
         if (!mControllerCallbackRecords.remove(new ControllerCallbackRecord(null, callback))) {
-            Log.w(TAG, "Ignoring unknown controller callback");
+            Log.w(TAG, "unregisterControllerCallback: Ignoring an unknown callback");
             return;
         }
     }
@@ -319,7 +351,7 @@
     /**
      * Sets an {@link OnGetControllerHintsListener} to send hints when creating a
      * {@link RoutingController}. To send the hints, listener should be set <em>BEFORE</em> calling
-     * {@link #requestCreateController(MediaRoute2Info)}.
+     * {@link #transferTo(MediaRoute2Info)}.
      *
      * @param listener A listener to send optional app-specific hints when creating a controller.
      *                 {@code null} for unset.
@@ -329,28 +361,52 @@
     }
 
     /**
-     * Requests the media route provider service to create a {@link RoutingController}
-     * with the given route.
+     * Transfers the current media to the given route.
+     * If it's necessary a new {@link RoutingController} is created or it is handled within
+     * the current controller.
      *
-     * @param route the route you want to create a controller with.
-     * @throws IllegalArgumentException if the given route is
-     * {@link MediaRoute2Info#isSystemRoute() system route}
+     * @param route the route you want to transfer the current media to. Pass {@code null} to
+     *              stop routing of the current media.
      *
-     * @see RoutingControllerCallback#onControllerCreated
-     * @see RoutingControllerCallback#onControllerCreationFailed
+     * @see TransferCallback#onTransferred
+     * @see TransferCallback#onTransferFailed
      */
-    public void requestCreateController(@NonNull MediaRoute2Info route) {
-        Objects.requireNonNull(route, "route must not be null");
-        if (route.isSystemRoute()) {
-            throw new IllegalArgumentException("Can't create a route controller with "
-                    + "a system route. Use getSystemController().");
+    public void transferTo(@Nullable MediaRoute2Info route) {
+        List<RoutingController> controllers = getControllers();
+        RoutingController controller = controllers.get(controllers.size() - 1);
+
+        transfer(controller, route);
+    }
+
+    /**
+     * Transfers the media of a routing controller to the given route.
+     * @param controller a routing controller controlling media routing.
+     * @param route the route you want to transfer the media to. Pass {@code null} to stop
+     *              routing controlled by the given controller.
+     * @hide
+     */
+    void transfer(@NonNull RoutingController controller, @Nullable MediaRoute2Info route) {
+        Objects.requireNonNull(controller, "controller must not be null");
+
+        if (route == null) {
+            controller.release();
+            return;
         }
+
         // TODO: Check the given route exists
+        // TODO: Check thread-safety
+        if (controller.getRoutingSessionInfo().getTransferableRoutes().contains(route.getId())) {
+            controller.transferToRoute(route);
+            return;
+        }
+
+        controller.release();
 
         final int requestId;
         requestId = mControllerCreationRequestCnt.getAndIncrement();
 
-        ControllerCreationRequest request = new ControllerCreationRequest(requestId, route);
+        ControllerCreationRequest request =
+                new ControllerCreationRequest(requestId, controller, route);
         mControllerCreationRequests.add(request);
 
         OnGetControllerHintsListener listener = mOnGetControllerHintsListener;
@@ -371,7 +427,7 @@
                 mMediaRouterService.requestCreateSessionWithRouter2(
                         stub, route, requestId, controllerHints);
             } catch (RemoteException ex) {
-                Log.e(TAG, "Unable to request to create controller.", ex);
+                Log.e(TAG, "transfer: Unable to request to create controller.", ex);
                 mHandler.sendMessage(obtainMessage(MediaRouter2::createControllerOnHandler,
                         MediaRouter2.this, null, requestId));
             }
@@ -385,7 +441,7 @@
      * Note: The system controller can't be released. Calling {@link RoutingController#release()}
      * will be ignored.
      * <p>
-     * This method will always return the same instance.
+     * This method always returns the same instance.
      */
     @NonNull
     public RoutingController getSystemController() {
@@ -493,9 +549,9 @@
     }
 
     /**
-     * Creates a controller and calls the {@link RoutingControllerCallback#onControllerCreated}.
+     * Creates a controller and calls the {@link TransferCallback#onTransferred}.
      * If the controller creation has failed, then it calls
-     * {@link RoutingControllerCallback#onControllerCreationFailed}.
+     * {@link TransferCallback#onTransferFailed}.
      * <p>
      * Pass {@code null} to sessionInfo for the failure case.
      */
@@ -516,14 +572,14 @@
             if (sessionInfo == null) {
                 // TODO: We may need to distinguish between failure and rejection.
                 //       One way can be introducing 'reason'.
-                notifyControllerCreationFailed(requestedRoute);
+                notifyTransferFailed(requestedRoute);
                 return;
             } else if (!sessionInfo.getSelectedRoutes().contains(requestedRoute.getId())) {
                 Log.w(TAG, "The session does not contain the requested route. "
                         + "(requestedRouteId=" + requestedRoute.getId()
                         + ", actualRoutes=" + sessionInfo.getSelectedRoutes()
                         + ")");
-                notifyControllerCreationFailed(requestedRoute);
+                notifyTransferFailed(requestedRoute);
                 return;
             } else if (!TextUtils.equals(requestedRoute.getProviderId(),
                     sessionInfo.getProviderId())) {
@@ -531,17 +587,18 @@
                         + "(requested route's providerId=" + requestedRoute.getProviderId()
                         + ", actual providerId=" + sessionInfo.getProviderId()
                         + ")");
-                notifyControllerCreationFailed(requestedRoute);
+                notifyTransferFailed(requestedRoute);
                 return;
             }
         }
 
         if (sessionInfo != null) {
-            RoutingController controller = new RoutingController(sessionInfo);
+            RoutingController newController = new RoutingController(sessionInfo);
             synchronized (sRouterLock) {
-                mRoutingControllers.put(controller.getId(), controller);
+                mRoutingControllers.put(newController.getId(), newController);
             }
-            notifyControllerCreated(controller);
+            notifyTransferred(matchingRequest != null ? matchingRequest.mController :
+                    getSystemController(), newController);
         }
     }
 
@@ -657,31 +714,32 @@
         }
     }
 
-    private void notifyControllerCreated(RoutingController controller) {
-        for (ControllerCallbackRecord record: mControllerCallbackRecords) {
+    private void notifyTransferred(RoutingController oldController,
+            RoutingController newController) {
+        for (TransferCallbackRecord record: mTransferCallbackRecords) {
             record.mExecutor.execute(
-                    () -> record.mControllerCallback.onControllerCreated(controller));
+                    () -> record.mTransferCallback.onTransferred(oldController,
+                            newController));
         }
     }
 
-    private void notifyControllerCreationFailed(MediaRoute2Info route) {
-        for (ControllerCallbackRecord record: mControllerCallbackRecords) {
+    private void notifyTransferFailed(MediaRoute2Info route) {
+        for (TransferCallbackRecord record: mTransferCallbackRecords) {
             record.mExecutor.execute(
-                    () -> record.mControllerCallback.onControllerCreationFailed(route));
+                    () -> record.mTransferCallback.onTransferFailed(route));
         }
     }
 
     private void notifyControllerUpdated(RoutingController controller) {
         for (ControllerCallbackRecord record: mControllerCallbackRecords) {
-            record.mExecutor.execute(
-                    () -> record.mControllerCallback.onControllerUpdated(controller));
+            record.mExecutor.execute(() -> record.mCallback.onControllerUpdated(controller));
         }
     }
 
     private void notifyControllerReleased(RoutingController controller) {
-        for (ControllerCallbackRecord record: mControllerCallbackRecords) {
+        for (TransferCallbackRecord record: mTransferCallbackRecords) {
             record.mExecutor.execute(
-                    () -> record.mControllerCallback.onControllerReleased(controller));
+                    () -> record.mTransferCallback.onTransferred(controller, null));
         }
     }
 
@@ -714,48 +772,28 @@
     }
 
     /**
-     * Callback for receiving a result of {@link RoutingController} creation and updates.
+     * Callback for receiving events on media transfer.
      */
-    public static class RoutingControllerCallback {
+    public static class TransferCallback {
         /**
-         * Called when the {@link RoutingController} is created.
-         * A {@link RoutingController} can be created by calling
-         * {@link #requestCreateController(MediaRoute2Info)}, or by the system.
+         * Called when a media is transferred between two different routing controllers.
+         * This can happen by calling {@link #transferTo(MediaRoute2Info)} or
+         * {@link RoutingController#release()}.
          *
-         * @param controller the controller to control routes
+         * @param oldController the previous controller that controlled routing.
+         * @param newController the new controller to control routing or {@code null} if the
+         *                      previous controller is released.
+         * @see #transferTo(MediaRoute2Info)
          */
-        public void onControllerCreated(@NonNull RoutingController controller) {}
+        public void onTransferred(@NonNull RoutingController oldController,
+                @Nullable RoutingController newController) {}
 
         /**
-         * Called when the controller creation request failed.
+         * Called when {@link #transferTo(MediaRoute2Info)} failed.
          *
-         * @param requestedRoute the route info which was used for the creation request
+         * @param requestedRoute the route info which was used for the transfer.
          */
-        public void onControllerCreationFailed(@NonNull MediaRoute2Info requestedRoute) {}
-
-        /**
-         * Called when the controller is updated.
-         *
-         * @param controller the updated controller. Can be the system controller.
-         * @see #getSystemController()
-         */
-        public void onControllerUpdated(@NonNull RoutingController controller) {}
-
-        /**
-         * Called when a routing controller is released. It can be released in two cases:
-         * <ul>
-         *     <li>When {@link RoutingController#release()} is called.</li>
-         *     <li>When the remote session in the provider is destroyed.</li>
-         * </ul>
-         * {@link RoutingController#isReleased()} will always return {@code true}
-         * for the {@code controller} here.
-         *
-         * @see RoutingController#release()
-         * @see RoutingController#isReleased()
-         */
-        // TODO: Add tests for checking whether this method is called.
-        // TODO: When service process dies, this should be called.
-        public void onControllerReleased(@NonNull RoutingController controller) {}
+        public void onTransferFailed(@NonNull MediaRoute2Info requestedRoute) {}
     }
 
     /**
@@ -769,8 +807,8 @@
          * The {@link Bundle} returned here will be sent to media route provider service as a hint.
          * <p>
          * To send hints when creating the controller, set the listener before calling
-         * {@link #requestCreateController(MediaRoute2Info)}. The method will be called
-         * on the same thread which calls {@link #requestCreateController(MediaRoute2Info)}.
+         * {@link #transferTo(MediaRoute2Info)}. The method will be called
+         * on the same thread which calls {@link #transferTo(MediaRoute2Info)}.
          *
          * @param route The route to create controller with
          * @return An optional bundle of app-specific arguments to send to the provider,
@@ -783,9 +821,23 @@
     }
 
     /**
+     * Callback for receiving {@link RoutingController} updates.
+     */
+    public static class ControllerCallback {
+        /**
+         * Called when a controller is updated. (e.g., the selected routes of the
+         * controller is changed or the volume of the controller is changed.)
+         *
+         * @param controller the updated controller. Can be the system controller.
+         * @see #getSystemController()
+         */
+        public void onControllerUpdated(@NonNull RoutingController controller) { }
+    }
+
+    /**
      * A class to control media routing session in media route provider.
-     * For example, selecting/deselcting/transferring routes to session can be done through this
-     * class. Instances are created by {@link #requestCreateController(MediaRoute2Info)}.
+     * For example, selecting/deselecting/transferring routes to a session can be done through this
+     * class. Instances are created by {@link #transferTo(MediaRoute2Info)}.
      */
     public class RoutingController {
         private final Object mControllerLock = new Object();
@@ -801,7 +853,7 @@
         }
 
         /**
-         * @return the ID of the controller
+         * @return the ID of the controller. It is globally unique.
          */
         @NonNull
         public String getId() {
@@ -851,16 +903,6 @@
         }
 
         /**
-         * @return the unmodifiable list of transferable routes for the session.
-         */
-        @NonNull
-        public List<MediaRoute2Info> getTransferableRoutes() {
-            synchronized (mControllerLock) {
-                return getRoutesWithIdsLocked(mSessionInfo.getTransferableRoutes());
-            }
-        }
-
-        /**
          * Gets information about how volume is handled on the session.
          *
          * @return {@link MediaRoute2Info#PLAYBACK_VOLUME_FIXED} or
@@ -919,9 +961,10 @@
          * </ul>
          * If the route doesn't meet any of above conditions, it will be ignored.
          *
+         * @see #deselectRoute(MediaRoute2Info)
          * @see #getSelectedRoutes()
          * @see #getSelectableRoutes()
-         * @see RoutingControllerCallback#onControllerUpdated
+         * @see ControllerCallback#onControllerUpdated
          */
         public void selectRoute(@NonNull MediaRoute2Info route) {
             Objects.requireNonNull(route, "route must not be null");
@@ -968,7 +1011,7 @@
          *
          * @see #getSelectedRoutes()
          * @see #getDeselectableRoutes()
-         * @see RoutingControllerCallback#onControllerUpdated
+         * @see ControllerCallback#onControllerUpdated
          */
         public void deselectRoute(@NonNull MediaRoute2Info route) {
             Objects.requireNonNull(route, "route must not be null");
@@ -1008,35 +1051,33 @@
          * Transfers to a given route for the remote session. The given route must satisfy
          * all of the following conditions:
          * <ul>
-         * <li>ID should not be included in {@link #getSelectedRoutes()}</li>
-         * <li>ID should be included in {@link #getTransferableRoutes()}</li>
+         * <li>ID should not be included in {@link RoutingSessionInfo#getSelectedRoutes()}</li>
+         * <li>ID should be included in {@link RoutingSessionInfo#getTransferableRoutes()}</li>
          * </ul>
          * If the route doesn't meet any of above conditions, it will be ignored.
          *
-         * @see #getSelectedRoutes()
-         * @see #getTransferableRoutes()
-         * @see RoutingControllerCallback#onControllerUpdated
+         * @see RoutingSessionInfo#getSelectedRoutes()
+         * @see RoutingSessionInfo#getTransferableRoutes()
+         * @see ControllerCallback#onControllerUpdated
          */
-        public void transferToRoute(@NonNull MediaRoute2Info route) {
+        void transferToRoute(@NonNull MediaRoute2Info route) {
             Objects.requireNonNull(route, "route must not be null");
             synchronized (mControllerLock) {
                 if (mIsReleased) {
                     Log.w(TAG, "transferToRoute() called on released controller. Ignoring.");
                     return;
                 }
-            }
 
-            List<MediaRoute2Info> selectedRoutes = getSelectedRoutes();
-            if (checkRouteListContainsRouteId(selectedRoutes, route.getId())) {
-                Log.w(TAG, "Ignoring transferring to a route that is already added. route="
-                        + route);
-                return;
-            }
+                if (mSessionInfo.getSelectedRoutes().contains(route.getId())) {
+                    Log.w(TAG, "Ignoring transferring to a route that is already added. "
+                            + "route=" + route);
+                    return;
+                }
 
-            List<MediaRoute2Info> transferableRoutes = getTransferableRoutes();
-            if (!checkRouteListContainsRouteId(transferableRoutes, route.getId())) {
-                Log.w(TAG, "Ignoring transferring to a non-transferable route=" + route);
-                return;
+                if (!mSessionInfo.getTransferableRoutes().contains(route.getId())) {
+                    Log.w(TAG, "Ignoring transferring to a non-transferrable route=" + route);
+                    return;
+                }
             }
 
             MediaRouter2Stub stub;
@@ -1132,8 +1173,6 @@
                     .map(MediaRoute2Info::getId).collect(Collectors.toList());
             List<String> deselectableRoutes = getDeselectableRoutes().stream()
                     .map(MediaRoute2Info::getId).collect(Collectors.toList());
-            List<String> transferableRoutes = getTransferableRoutes().stream()
-                    .map(MediaRoute2Info::getId).collect(Collectors.toList());
 
             StringBuilder result = new StringBuilder()
                     .append("RoutingController{ ")
@@ -1147,19 +1186,12 @@
                     .append(", deselectableRoutes={")
                     .append(deselectableRoutes)
                     .append("}")
-                    .append(", transferableRoutes={")
-                    .append(transferableRoutes)
-                    .append("}")
                     .append(" }");
             return result.toString();
         }
 
-        /**
-         * TODO: Change this to package private. (Hidden for debugging purposes)
-         * @hide
-         */
         @NonNull
-        public RoutingSessionInfo getRoutingSessionInfo() {
+        RoutingSessionInfo getRoutingSessionInfo() {
             synchronized (mControllerLock) {
                 return mSessionInfo;
             }
@@ -1205,7 +1237,7 @@
         }
     }
 
-    final class RouteCallbackRecord {
+    static final class RouteCallbackRecord {
         public final Executor mExecutor;
         public final RouteCallback mRouteCallback;
         public final RouteDiscoveryPreference mPreference;
@@ -1234,13 +1266,41 @@
         }
     }
 
-    final class ControllerCallbackRecord {
+    static final class TransferCallbackRecord {
         public final Executor mExecutor;
-        public final RoutingControllerCallback mControllerCallback;
+        public final TransferCallback mTransferCallback;
 
-        ControllerCallbackRecord(@NonNull Executor executor,
-                @NonNull RoutingControllerCallback controllerCallback) {
-            mControllerCallback = controllerCallback;
+        TransferCallbackRecord(@NonNull Executor executor,
+                @NonNull TransferCallback transferCallback) {
+            mTransferCallback = transferCallback;
+            mExecutor = executor;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (this == obj) {
+                return true;
+            }
+            if (!(obj instanceof TransferCallbackRecord)) {
+                return false;
+            }
+            return mTransferCallback
+                    == ((TransferCallbackRecord) obj).mTransferCallback;
+        }
+
+        @Override
+        public int hashCode() {
+            return mTransferCallback.hashCode();
+        }
+    }
+
+    static final class ControllerCallbackRecord {
+        public final Executor mExecutor;
+        public final ControllerCallback mCallback;
+
+        ControllerCallbackRecord(@Nullable Executor executor,
+                @NonNull ControllerCallback callback) {
+            mCallback = callback;
             mExecutor = executor;
         }
 
@@ -1252,23 +1312,25 @@
             if (!(obj instanceof ControllerCallbackRecord)) {
                 return false;
             }
-            return mControllerCallback
-                    == ((ControllerCallbackRecord) obj).mControllerCallback;
+            return mCallback == ((ControllerCallbackRecord) obj).mCallback;
         }
 
         @Override
         public int hashCode() {
-            return mControllerCallback.hashCode();
+            return mCallback.hashCode();
         }
     }
 
-    final class ControllerCreationRequest {
-        public final MediaRoute2Info mRoute;
+    static final class ControllerCreationRequest {
         public final int mRequestId;
+        public final RoutingController mController;
+        public final MediaRoute2Info mRoute;
 
-        ControllerCreationRequest(int requestId, @NonNull MediaRoute2Info route) {
-            mRoute = route;
+        ControllerCreationRequest(int requestId, @NonNull RoutingController controller,
+                @NonNull MediaRoute2Info route) {
             mRequestId = requestId;
+            mController = controller;
+            mRoute = route;
         }
     }
 
diff --git a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
index ca7c36c..e80562b 100644
--- a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
+++ b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
@@ -39,7 +39,7 @@
 import android.media.MediaRoute2Info;
 import android.media.MediaRouter2;
 import android.media.MediaRouter2.RouteCallback;
-import android.media.MediaRouter2.RoutingControllerCallback;
+import android.media.MediaRouter2.TransferCallback;
 import android.media.MediaRouter2Manager;
 import android.media.MediaRouter2Utils;
 import android.media.RouteDiscoveryPreference;
@@ -79,7 +79,7 @@
 
     private final List<MediaRouter2Manager.Callback> mManagerCallbacks = new ArrayList<>();
     private final List<RouteCallback> mRouteCallbacks = new ArrayList<>();
-    private final List<RoutingControllerCallback> mControllerCallbacks = new ArrayList<>();
+    private final List<MediaRouter2.TransferCallback> mTransferCallbacks = new ArrayList<>();
 
     public static final List<String> FEATURES_ALL = new ArrayList();
     public static final List<String> FEATURES_SPECIAL = new ArrayList();
@@ -153,18 +153,14 @@
 
         MediaRoute2Info routeToRemove = routes.get(ROUTE_ID2);
 
-        try {
-            SampleMediaRoute2ProviderService sInstance =
-                    SampleMediaRoute2ProviderService.getInstance();
-            assertNotNull(sInstance);
-            sInstance.removeRoute(ROUTE_ID2);
-            assertTrue(removedLatch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        SampleMediaRoute2ProviderService sInstance =
+                SampleMediaRoute2ProviderService.getInstance();
+        assertNotNull(sInstance);
+        sInstance.removeRoute(ROUTE_ID2);
+        assertTrue(removedLatch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
 
-            sInstance.addRoute(routeToRemove);
-            assertTrue(addedLatch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
-        } finally {
-            mRouter2.unregisterRouteCallback(routeCallback);
-        }
+        sInstance.addRoute(routeToRemove);
+        assertTrue(addedLatch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
     }
 
     /**
@@ -198,10 +194,14 @@
         addManagerCallback(new MediaRouter2Manager.Callback());
         //TODO: remove this when it's not necessary.
         addRouterCallback(new MediaRouter2.RouteCallback());
-        addSessionCallback(new RoutingControllerCallback() {
+        addTransferCallback(new MediaRouter2.TransferCallback() {
             @Override
-            public void onControllerCreated(MediaRouter2.RoutingController controller) {
-                if (createRouteMap(controller.getSelectedRoutes()).containsKey(ROUTE_ID1)) {
+            public void onTransferred(MediaRouter2.RoutingController oldController,
+                    MediaRouter2.RoutingController newController) {
+                if (newController == null) {
+                    return;
+                }
+                if (createRouteMap(newController.getSelectedRoutes()).containsKey(ROUTE_ID1)) {
                     latch.countDown();
                 }
             }
@@ -363,7 +363,7 @@
         int currentVolume = sessionInfo.getVolume();
         int targetVolume = (currentVolume == 0) ? 1 : (currentVolume - 1);
 
-        RoutingControllerCallback routingControllerCallback = new RoutingControllerCallback() {
+        MediaRouter2.ControllerCallback controllerCallback = new MediaRouter2.ControllerCallback() {
             @Override
             public void onControllerUpdated(MediaRouter2.RoutingController controller) {
                 if (!TextUtils.equals(sessionInfo.getId(), controller.getId())) {
@@ -374,7 +374,6 @@
                 }
             }
         };
-        mRouter2.registerControllerCallback(mExecutor, routingControllerCallback);
 
         addManagerCallback(new MediaRouter2Manager.Callback() {
             @Override
@@ -390,12 +389,12 @@
             }
         });
 
-        mManager.setSessionVolume(sessionInfo, targetVolume);
-
         try {
+            mRouter2.registerControllerCallback(mExecutor, controllerCallback);
+            mManager.setSessionVolume(sessionInfo, targetVolume);
             assertTrue(volumeChangedLatch.await(TIMEOUT_MS, TimeUnit.MILLISECONDS));
         } finally {
-            mRouter2.unregisterControllerCallback(routingControllerCallback);
+            mRouter2.unregisterControllerCallback(controllerCallback);
         }
     }
 
@@ -491,9 +490,9 @@
         mRouter2.registerRouteCallback(mExecutor, routeCallback, RouteDiscoveryPreference.EMPTY);
     }
 
-    private void addSessionCallback(RoutingControllerCallback controllerCallback) {
-        mControllerCallbacks.add(controllerCallback);
-        mRouter2.registerControllerCallback(mExecutor, controllerCallback);
+    private void addTransferCallback(TransferCallback transferCallback) {
+        mTransferCallbacks.add(transferCallback);
+        mRouter2.registerTransferCallback(mExecutor, transferCallback);
     }
 
     private void clearCallbacks() {
@@ -507,10 +506,10 @@
         }
         mRouteCallbacks.clear();
 
-        for (RoutingControllerCallback controllerCallback : mControllerCallbacks) {
-            mRouter2.unregisterControllerCallback(controllerCallback);
+        for (MediaRouter2.TransferCallback transferCallback : mTransferCallbacks) {
+            mRouter2.unregisterTransferCallback(transferCallback);
         }
-        mControllerCallbacks.clear();
+        mTransferCallbacks.clear();
     }
 
     private void releaseAllSessions() {