diff --git a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
index 1dd796504ea37..fb5847929a459 100644
--- a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
+++ b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
@@ -143,10 +143,7 @@ class MediaRouter2ServiceImpl {
         mContext.registerReceiver(mScreenOnOffReceiver, screenOnOffIntentFilter);
     }
 
-    ////////////////////////////////////////////////////////////////
-    ////  Calls from MediaRouter2
-    ////   - Should not have @NonNull/@Nullable on any arguments
-    ////////////////////////////////////////////////////////////////
+    // Methods that implement MediaRouter2 operations.
 
     @NonNull
     public void enforceMediaContentControlPermission() {
@@ -234,7 +231,7 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void registerRouter2(IMediaRouter2 router, String packageName) {
+    public void registerRouter2(@NonNull IMediaRouter2 router, @NonNull String packageName) {
         Objects.requireNonNull(router, "router must not be null");
         if (TextUtils.isEmpty(packageName)) {
             throw new IllegalArgumentException("packageName must not be empty");
@@ -261,7 +258,7 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void unregisterRouter2(IMediaRouter2 router) {
+    public void unregisterRouter2(@NonNull IMediaRouter2 router) {
         Objects.requireNonNull(router, "router must not be null");
 
         final long token = Binder.clearCallingIdentity();
@@ -274,8 +271,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void setDiscoveryRequestWithRouter2(IMediaRouter2 router,
-            RouteDiscoveryPreference preference) {
+    public void setDiscoveryRequestWithRouter2(@NonNull IMediaRouter2 router,
+            @NonNull RouteDiscoveryPreference preference) {
         Objects.requireNonNull(router, "router must not be null");
         Objects.requireNonNull(preference, "preference must not be null");
 
@@ -294,8 +291,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void setRouteVolumeWithRouter2(IMediaRouter2 router,
-            MediaRoute2Info route, int volume) {
+    public void setRouteVolumeWithRouter2(@NonNull IMediaRouter2 router,
+            @NonNull MediaRoute2Info route, int volume) {
         Objects.requireNonNull(router, "router must not be null");
         Objects.requireNonNull(route, "route must not be null");
 
@@ -309,9 +306,9 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void requestCreateSessionWithRouter2(IMediaRouter2 router, int requestId,
-            long managerRequestId, RoutingSessionInfo oldSession,
-            MediaRoute2Info route, Bundle sessionHints) {
+    public void requestCreateSessionWithRouter2(@NonNull IMediaRouter2 router, int requestId,
+            long managerRequestId, @NonNull RoutingSessionInfo oldSession,
+            @NonNull MediaRoute2Info route, Bundle sessionHints) {
         Objects.requireNonNull(router, "router must not be null");
         Objects.requireNonNull(oldSession, "oldSession must not be null");
         Objects.requireNonNull(route, "route must not be null");
@@ -327,8 +324,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void selectRouteWithRouter2(IMediaRouter2 router, String uniqueSessionId,
-            MediaRoute2Info route) {
+    public void selectRouteWithRouter2(@NonNull IMediaRouter2 router,
+            @NonNull String uniqueSessionId, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(router, "router must not be null");
         Objects.requireNonNull(route, "route must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
@@ -345,8 +342,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void deselectRouteWithRouter2(IMediaRouter2 router, String uniqueSessionId,
-            MediaRoute2Info route) {
+    public void deselectRouteWithRouter2(@NonNull IMediaRouter2 router,
+            @NonNull String uniqueSessionId, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(router, "router must not be null");
         Objects.requireNonNull(route, "route must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
@@ -363,8 +360,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void transferToRouteWithRouter2(IMediaRouter2 router, String uniqueSessionId,
-            MediaRoute2Info route) {
+    public void transferToRouteWithRouter2(@NonNull IMediaRouter2 router,
+            @NonNull String uniqueSessionId, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(router, "router must not be null");
         Objects.requireNonNull(route, "route must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
@@ -381,8 +378,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void setSessionVolumeWithRouter2(IMediaRouter2 router, String uniqueSessionId,
-            int volume) {
+    public void setSessionVolumeWithRouter2(@NonNull IMediaRouter2 router,
+            @NonNull String uniqueSessionId, int volume) {
         Objects.requireNonNull(router, "router must not be null");
         Objects.requireNonNull(uniqueSessionId, "uniqueSessionId must not be null");
 
@@ -396,7 +393,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void releaseSessionWithRouter2(IMediaRouter2 router, String uniqueSessionId) {
+    public void releaseSessionWithRouter2(@NonNull IMediaRouter2 router,
+            @NonNull String uniqueSessionId) {
         Objects.requireNonNull(router, "router must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
             throw new IllegalArgumentException("uniqueSessionId must not be empty");
@@ -412,13 +410,10 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    ////////////////////////////////////////////////////////////////
-    ////  Calls from MediaRouter2Manager
-    ////   - Should not have @NonNull/@Nullable on any arguments
-    ////////////////////////////////////////////////////////////////
+    // Methods that implement MediaRouter2Manager operations.
 
     @NonNull
-    public List<RoutingSessionInfo> getRemoteSessions(IMediaRouter2Manager manager) {
+    public List<RoutingSessionInfo> getRemoteSessions(@NonNull IMediaRouter2Manager manager) {
         Objects.requireNonNull(manager, "manager must not be null");
         final long token = Binder.clearCallingIdentity();
         try {
@@ -430,7 +425,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void registerManager(IMediaRouter2Manager manager, String packageName) {
+    public void registerManager(@NonNull IMediaRouter2Manager manager,
+            @NonNull String packageName) {
         Objects.requireNonNull(manager, "manager must not be null");
         if (TextUtils.isEmpty(packageName)) {
             throw new IllegalArgumentException("packageName must not be empty");
@@ -450,7 +446,7 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void unregisterManager(IMediaRouter2Manager manager) {
+    public void unregisterManager(@NonNull IMediaRouter2Manager manager) {
         Objects.requireNonNull(manager, "manager must not be null");
 
         final long token = Binder.clearCallingIdentity();
@@ -463,7 +459,7 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void startScan(IMediaRouter2Manager manager) {
+    public void startScan(@NonNull IMediaRouter2Manager manager) {
         Objects.requireNonNull(manager, "manager must not be null");
         final long token = Binder.clearCallingIdentity();
         try {
@@ -475,7 +471,7 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void stopScan(IMediaRouter2Manager manager) {
+    public void stopScan(@NonNull IMediaRouter2Manager manager) {
         Objects.requireNonNull(manager, "manager must not be null");
         final long token = Binder.clearCallingIdentity();
         try {
@@ -487,8 +483,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void setRouteVolumeWithManager(IMediaRouter2Manager manager, int requestId,
-            MediaRoute2Info route, int volume) {
+    public void setRouteVolumeWithManager(@NonNull IMediaRouter2Manager manager, int requestId,
+            @NonNull MediaRoute2Info route, int volume) {
         Objects.requireNonNull(manager, "manager must not be null");
         Objects.requireNonNull(route, "route must not be null");
 
@@ -502,10 +498,11 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void requestCreateSessionWithManager(IMediaRouter2Manager manager, int requestId,
-            RoutingSessionInfo oldSession, MediaRoute2Info route) {
+    public void requestCreateSessionWithManager(@NonNull IMediaRouter2Manager manager,
+            int requestId, @NonNull RoutingSessionInfo oldSession, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(manager, "manager must not be null");
         Objects.requireNonNull(oldSession, "oldSession must not be null");
+        Objects.requireNonNull(route, "route must not be null");
 
         final long token = Binder.clearCallingIdentity();
         try {
@@ -517,8 +514,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void selectRouteWithManager(IMediaRouter2Manager manager, int requestId,
-            String uniqueSessionId, MediaRoute2Info route) {
+    public void selectRouteWithManager(@NonNull IMediaRouter2Manager manager, int requestId,
+            @NonNull String uniqueSessionId, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(manager, "manager must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
             throw new IllegalArgumentException("uniqueSessionId must not be empty");
@@ -535,8 +532,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void deselectRouteWithManager(IMediaRouter2Manager manager, int requestId,
-            String uniqueSessionId, MediaRoute2Info route) {
+    public void deselectRouteWithManager(@NonNull IMediaRouter2Manager manager, int requestId,
+            @NonNull String uniqueSessionId, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(manager, "manager must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
             throw new IllegalArgumentException("uniqueSessionId must not be empty");
@@ -553,8 +550,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void transferToRouteWithManager(IMediaRouter2Manager manager, int requestId,
-            String uniqueSessionId, MediaRoute2Info route) {
+    public void transferToRouteWithManager(@NonNull IMediaRouter2Manager manager, int requestId,
+            @NonNull String uniqueSessionId, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(manager, "manager must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
             throw new IllegalArgumentException("uniqueSessionId must not be empty");
@@ -571,8 +568,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void setSessionVolumeWithManager(IMediaRouter2Manager manager, int requestId,
-            String uniqueSessionId, int volume) {
+    public void setSessionVolumeWithManager(@NonNull IMediaRouter2Manager manager, int requestId,
+            @NonNull String uniqueSessionId, int volume) {
         Objects.requireNonNull(manager, "manager must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
             throw new IllegalArgumentException("uniqueSessionId must not be empty");
@@ -588,8 +585,8 @@ class MediaRouter2ServiceImpl {
         }
     }
 
-    public void releaseSessionWithManager(IMediaRouter2Manager manager, int requestId,
-            String uniqueSessionId) {
+    public void releaseSessionWithManager(@NonNull IMediaRouter2Manager manager, int requestId,
+            @NonNull String uniqueSessionId) {
         Objects.requireNonNull(manager, "manager must not be null");
         if (TextUtils.isEmpty(uniqueSessionId)) {
             throw new IllegalArgumentException("uniqueSessionId must not be empty");
@@ -668,11 +665,6 @@ class MediaRouter2ServiceImpl {
         return mUserManagerInternal.getProfileParentId(userId) == mCurrentActiveUserId;
     }
 
-    ////////////////////////////////////////////////////////////////
-    ////  ***Locked methods related to MediaRouter2
-    ////   - Should have @NonNull/@Nullable on all arguments
-    ////////////////////////////////////////////////////////////////
-
     @GuardedBy("mLock")
     private void registerRouter2Locked(@NonNull IMediaRouter2 router, int uid, int pid,
             @NonNull String packageName, int userId, boolean hasConfigureWifiDisplayPermission,
@@ -900,11 +892,6 @@ class MediaRouter2ServiceImpl {
                         DUMMY_REQUEST_ID, routerRecord, uniqueSessionId));
     }
 
-    ////////////////////////////////////////////////////////////
-    ////  ***Locked methods related to MediaRouter2Manager
-    ////   - Should have @NonNull/@Nullable on all arguments
-    ////////////////////////////////////////////////////////////
-
     private List<RoutingSessionInfo> getRemoteSessionsLocked(
             @NonNull IMediaRouter2Manager manager) {
         final IBinder binder = manager.asBinder();
@@ -1013,8 +1000,8 @@ class MediaRouter2ServiceImpl {
     }
 
     private void requestCreateSessionWithManagerLocked(int requestId,
-            @NonNull IMediaRouter2Manager manager,
-            @NonNull RoutingSessionInfo oldSession, @NonNull MediaRoute2Info route) {
+            @NonNull IMediaRouter2Manager manager, @NonNull RoutingSessionInfo oldSession,
+            @NonNull MediaRoute2Info route) {
         ManagerRecord managerRecord = mAllManagerRecords.get(manager.asBinder());
         if (managerRecord == null) {
             return;
@@ -1136,8 +1123,7 @@ class MediaRouter2ServiceImpl {
     }
 
     private void releaseSessionWithManagerLocked(int requestId,
-            @NonNull IMediaRouter2Manager manager,
-            @NonNull String uniqueSessionId) {
+            @NonNull IMediaRouter2Manager manager, @NonNull String uniqueSessionId) {
         final IBinder binder = manager.asBinder();
         ManagerRecord managerRecord = mAllManagerRecords.get(binder);
 
@@ -1155,11 +1141,6 @@ class MediaRouter2ServiceImpl {
                         uniqueRequestId, routerRecord, uniqueSessionId));
     }
 
-    ////////////////////////////////////////////////////////////
-    ////  ***Locked methods used by both router2 and manager
-    ////   - Should have @NonNull/@Nullable on all arguments
-    ////////////////////////////////////////////////////////////
-
     @GuardedBy("mLock")
     private UserRecord getOrCreateUserRecordLocked(int userId) {
         UserRecord userRecord = mUserRecords.get(userId);