MediaSession2: Use Executor for callback handling

This also simplifies future work for adding more functions
Test: Run all MediaComponents test once

Change-Id: Ib9aebd9212368d616dba99792d6ed13b24617885
diff --git a/packages/MediaComponents/src/com/android/media/MediaController2Impl.java b/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
index 43f8473..edba88a 100644
--- a/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
+++ b/packages/MediaComponents/src/com/android/media/MediaController2Impl.java
@@ -222,38 +222,38 @@
 
     @Override
     public void play_impl() {
-        sendCommand(MediaSession2.COMMAND_CODE_PLAYBACK_START);
+        sendTransportControlCommand(MediaSession2.COMMAND_CODE_PLAYBACK_START);
     }
 
     @Override
     public void pause_impl() {
-        sendCommand(MediaSession2.COMMAND_CODE_PLAYBACK_PAUSE);
+        sendTransportControlCommand(MediaSession2.COMMAND_CODE_PLAYBACK_PAUSE);
     }
 
     @Override
     public void stop_impl() {
-        sendCommand(MediaSession2.COMMAND_CODE_PLAYBACK_STOP);
+        sendTransportControlCommand(MediaSession2.COMMAND_CODE_PLAYBACK_STOP);
     }
 
     @Override
     public void skipToPrevious_impl() {
-        sendCommand(MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_PREV_ITEM);
+        sendTransportControlCommand(MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_PREV_ITEM);
     }
 
     @Override
     public void skipToNext_impl() {
-        sendCommand(MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_NEXT_ITEM);
+        sendTransportControlCommand(MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_NEXT_ITEM);
     }
 
-    private void sendCommand(int code) {
-        // TODO(jaewan): optimization) Cache Command objects?
-        Command command = new Command(code);
-        // TODO(jaewan): Check if the command is in the allowed group.
+    private void sendTransportControlCommand(int commandCode) {
+        sendTransportControlCommand(commandCode, 0);
+    }
 
+    private void sendTransportControlCommand(int commandCode, long arg) {
         final IMediaSession2 binder = mSessionBinder;
         if (binder != null) {
             try {
-                binder.sendCommand(mSessionCallbackStub, command.toBundle(), null);
+                binder.sendTransportControlCommand(mSessionCallbackStub, commandCode, arg);
             } catch (RemoteException e) {
                 Log.w(TAG, "Cannot connect to the service or the session is gone", e);
             }
diff --git a/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java b/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java
index 22a3187..d7d29a6 100644
--- a/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java
+++ b/packages/MediaComponents/src/com/android/media/MediaSession2Impl.java
@@ -25,6 +25,7 @@
 import android.media.IMediaSession2Callback;
 import android.media.MediaItem2;
 import android.media.MediaPlayerBase;
+import android.media.MediaPlayerBase.PlaybackListener;
 import android.media.MediaSession2;
 import android.media.MediaSession2.Builder;
 import android.media.MediaSession2.Command;
@@ -39,11 +40,11 @@
 import android.media.session.MediaSessionManager;
 import android.media.update.MediaSession2Provider;
 import android.os.Bundle;
-import android.os.Handler;
 import android.os.IBinder;
-import android.os.Looper;
 import android.os.ResultReceiver;
+import android.support.annotation.GuardedBy;
 import android.util.Log;
+
 import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.List;
@@ -53,20 +54,21 @@
     private static final String TAG = "MediaSession2";
     private static final boolean DEBUG = true;//Log.isLoggable(TAG, Log.DEBUG);
 
-    private final MediaSession2 mInstance;
+    private final Object mLock = new Object();
 
+    private final MediaSession2 mInstance;
     private final Context mContext;
     private final String mId;
-    private final Handler mHandler;
     private final Executor mCallbackExecutor;
+    private final SessionCallback mCallback;
     private final MediaSession2Stub mSessionStub;
     private final SessionToken2 mSessionToken;
-
-    private MediaPlayerBase mPlayer;
-
     private final List<PlaybackListenerHolder> mListeners = new ArrayList<>();
+
+    @GuardedBy("mLock")
+    private MediaPlayerBase mPlayer;
+    @GuardedBy("mLock")
     private MyPlaybackListener mListener;
-    private MediaSession2 instance;
 
     /**
      * Can be only called by the {@link Builder#build()}.
@@ -90,9 +92,9 @@
         // Initialize finals first.
         mContext = context;
         mId = id;
-        mHandler = new Handler(Looper.myLooper());
+        mCallback = callback;
         mCallbackExecutor = callbackExecutor;
-        mSessionStub = new MediaSession2Stub(this, callback);
+        mSessionStub = new MediaSession2Stub(this);
         // Ask server to create session token for following reasons.
         //   1. Make session ID unique per package.
         //      Server can only know if the package has another process and has another session
@@ -118,7 +120,8 @@
     //               setPlayer(null). Token can be available when player is null, and
     //               controller can also attach to session.
     @Override
-    public void setPlayer_impl(MediaPlayerBase player, VolumeProvider volumeProvider) throws IllegalArgumentException {
+    public void setPlayer_impl(MediaPlayerBase player, VolumeProvider volumeProvider)
+            throws IllegalArgumentException {
         ensureCallingThread();
         if (player == null) {
             throw new IllegalArgumentException("player shouldn't be null");
@@ -127,26 +130,24 @@
     }
 
     private void setPlayerInternal(MediaPlayerBase player) {
-        if (mPlayer == player) {
-            // Player didn't changed. No-op.
-            return;
+        synchronized (mLock) {
+            if (mPlayer == player) {
+                // Player didn't changed. No-op.
+                return;
+            }
+            if (mPlayer != null && mListener != null) {
+                // This might not work for a poorly implemented player.
+                mPlayer.removePlaybackListener(mListener);
+            }
+            mListener = new MyPlaybackListener(this, player);
+            player.addPlaybackListener(mCallbackExecutor, mListener);
+            mPlayer = player;
         }
-        // TODO(jaewan): Find equivalent for the executor
-        //mHandler.removeCallbacksAndMessages(null);
-        if (mPlayer != null && mListener != null) {
-            // This might not work for a poorly implemented player.
-            mPlayer.removePlaybackListener(mListener);
-        }
-        mListener = new MyPlaybackListener(this, player);
-        player.addPlaybackListener(mCallbackExecutor, mListener);
-        notifyPlaybackStateChanged(player.getPlaybackState());
-        mPlayer = player;
+        notifyPlaybackStateChangedNotLocked(player.getPlaybackState());
     }
 
     @Override
     public void close_impl() {
-        // Flush any pending messages.
-        mHandler.removeCallbacksAndMessages(null);
         if (mSessionStub != null) {
             if (DEBUG) {
                 Log.d(TAG, "session is now unavailable, id=" + mId);
@@ -154,6 +155,14 @@
             // Invalidate previously published session stub.
             mSessionStub.destroyNotLocked();
         }
+        synchronized (mLock) {
+            if (mPlayer != null) {
+                // close can be called multiple times
+                mPlayer.removePlaybackListener(mListener);
+                mPlayer = null;
+                return;
+            }
+        }
     }
 
     @Override
@@ -321,14 +330,14 @@
         }
     }
 
-    Handler getHandler() {
-        return mHandler;
-    }
-
-    private void notifyPlaybackStateChanged(PlaybackState2 state) {
+    private void notifyPlaybackStateChangedNotLocked(PlaybackState2 state) {
+        List<PlaybackListenerHolder> listeners = new ArrayList<>();
+        synchronized (mLock) {
+            listeners.addAll(mListeners);
+        }
         // Notify to listeners added directly to this session
-        for (int i = 0; i < mListeners.size(); i++) {
-            mListeners.get(i).postPlaybackChange(state);
+        for (int i = 0; i < listeners.size(); i++) {
+            listeners.get(i).postPlaybackChange(state);
         }
         // Notify to controllers as well.
         mSessionStub.notifyPlaybackStateChangedNotLocked(state);
@@ -346,6 +355,14 @@
         return mPlayer;
     }
 
+    Executor getCallbackExecutor() {
+        return mCallbackExecutor;
+    }
+
+    SessionCallback getCallback() {
+        return mCallback;
+    }
+
     private static class MyPlaybackListener implements MediaPlayerBase.PlaybackListener {
         private final WeakReference<MediaSession2Impl> mSession;
         private final MediaPlayerBase mPlayer;
@@ -363,7 +380,7 @@
                         new IllegalStateException());
                 return;
             }
-            session.notifyPlaybackStateChanged(state);
+            session.notifyPlaybackStateChangedNotLocked(state);
         }
     }
 
diff --git a/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java b/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java
index 2f75dfa..64c8571 100644
--- a/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java
+++ b/packages/MediaComponents/src/com/android/media/MediaSession2Stub.java
@@ -51,30 +51,19 @@
     private static final boolean DEBUG = true; // TODO(jaewan): Rename.
 
     private final Object mLock = new Object();
-    private final CommandHandler mCommandHandler;
     private final WeakReference<MediaSession2Impl> mSession;
-    private final Context mContext;
-    private final SessionCallback mSessionCallback;
-    private final MediaLibrarySessionCallback mLibraryCallback;
 
     @GuardedBy("mLock")
     private final ArrayMap<IBinder, ControllerInfo> mControllers = new ArrayMap<>();
 
-    public MediaSession2Stub(MediaSession2Impl session, SessionCallback callback) {
+    public MediaSession2Stub(MediaSession2Impl session) {
         mSession = new WeakReference<>(session);
-        mContext = session.getContext();
-        // TODO(jaewan): Should be executor from the session builder
-        mCommandHandler = new CommandHandler(session.getHandler().getLooper());
-        mSessionCallback = callback;
-        mLibraryCallback = (callback instanceof MediaLibrarySessionCallback)
-                ? (MediaLibrarySessionCallback) callback : null;
     }
 
     public void destroyNotLocked() {
         final List<ControllerInfo> list;
         synchronized (mLock) {
             mSession.clear();
-            mCommandHandler.removeCallbacksAndMessages(null);
             list = getControllers();
             mControllers.clear();
         }
@@ -99,14 +88,43 @@
     }
 
     @Override
-    public void connect(String callingPackage, IMediaSession2Callback callback) {
-        if (callback == null) {
-            // Requesting connect without callback to receive result.
-            return;
-        }
-        ControllerInfo request = new ControllerInfo(mContext,
+    public void connect(String callingPackage, IMediaSession2Callback callback)
+            throws RuntimeException {
+        final MediaSession2Impl sessionImpl = getSession();
+        final ControllerInfo request = new ControllerInfo(sessionImpl.getContext(),
                 Binder.getCallingUid(), Binder.getCallingPid(), callingPackage, callback);
-        mCommandHandler.postConnect(request);
+        sessionImpl.getCallbackExecutor().execute(() -> {
+            final MediaSession2Impl session = mSession.get();
+            if (session == null) {
+                return;
+            }
+            CommandGroup allowedCommands = session.getCallback().onConnect(request);
+            // Don't reject connection for the request from trusted app.
+            // Otherwise server will fail to retrieve session's information to dispatch
+            // media keys to.
+            boolean accept = allowedCommands != null || request.isTrusted();
+            ControllerInfoImpl impl = ControllerInfoImpl.from(request);
+            if (accept) {
+                synchronized (mLock) {
+                    mControllers.put(impl.getId(), request);
+                }
+                if (allowedCommands == null) {
+                    // For trusted apps, send non-null allowed commands to keep connection.
+                    allowedCommands = new CommandGroup();
+                }
+            }
+            if (DEBUG) {
+                Log.d(TAG, "onConnectResult, request=" + request
+                        + " accept=" + accept);
+            }
+            try {
+                impl.getControllerBinder().onConnectionChanged(
+                        accept ? MediaSession2Stub.this : null,
+                        allowedCommands == null ? null : allowedCommands.toBundle());
+            } catch (RemoteException e) {
+                // Controller may be died prematurely.
+            }
+        });
     }
 
     @Override
@@ -122,20 +140,64 @@
     @Override
     public void sendCommand(IMediaSession2Callback caller, Bundle command, Bundle args)
             throws RuntimeException {
-        ControllerInfo controller = getController(caller);
+        // TODO(jaewan): Generic command
+    }
+
+    @Override
+    public void sendTransportControlCommand(IMediaSession2Callback caller,
+            int commandCode, long arg) throws RuntimeException {
+        final MediaSession2Impl sessionImpl = getSession();
+        final ControllerInfo controller = getController(caller);
         if (controller == null) {
             if (DEBUG) {
                 Log.d(TAG, "Command from a controller that hasn't connected. Ignore");
             }
             return;
         }
-        mCommandHandler.postCommand(controller, Command.fromBundle(command), args);
+        sessionImpl.getCallbackExecutor().execute(() -> {
+            final MediaSession2Impl session = mSession.get();
+            if (session == null) {
+                return;
+            }
+            // TODO(jaewan): Sanity check.
+            Command command = new Command(commandCode);
+            boolean accepted = session.getCallback().onCommandRequest(controller, command);
+            if (!accepted) {
+                // Don't run rejected command.
+                if (DEBUG) {
+                    Log.d(TAG, "Command " + commandCode + " from "
+                            + controller + " was rejected by " + session);
+                }
+                return;
+            }
+
+            switch (commandCode) {
+                case MediaSession2.COMMAND_CODE_PLAYBACK_START:
+                    session.getInstance().play();
+                    break;
+                case MediaSession2.COMMAND_CODE_PLAYBACK_PAUSE:
+                    session.getInstance().pause();
+                    break;
+                case MediaSession2.COMMAND_CODE_PLAYBACK_STOP:
+                    session.getInstance().stop();
+                    break;
+                case MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_PREV_ITEM:
+                    session.getInstance().skipToPrevious();
+                    break;
+                case MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_NEXT_ITEM:
+                    session.getInstance().skipToNext();
+                    break;
+                default:
+                    // TODO(jaewan): Resend unknown (new) commands through the custom command.
+            }
+        });
     }
 
     @Override
     public void getBrowserRoot(IMediaSession2Callback caller, Bundle rootHints)
             throws RuntimeException {
-        if (mLibraryCallback == null) {
+        final MediaSession2Impl sessionImpl = getSession();
+        if (!(sessionImpl.getCallback() instanceof MediaLibrarySessionCallback)) {
             if (DEBUG) {
                 Log.d(TAG, "Session cannot hand getBrowserRoot()");
             }
@@ -148,7 +210,24 @@
             }
             return;
         }
-        mCommandHandler.postOnGetRoot(controller, rootHints);
+        sessionImpl.getCallbackExecutor().execute(() -> {
+            final MediaSession2Impl session = mSession.get();
+            if (session == null) {
+                return;
+            }
+            final MediaLibrarySessionCallback libraryCallback =
+                    (MediaLibrarySessionCallback) session.getCallback();
+            final ControllerInfoImpl controllerImpl = ControllerInfoImpl.from(controller);
+            BrowserRoot root = libraryCallback.onGetRoot(controller, rootHints);
+            try {
+                controllerImpl.getControllerBinder().onGetRootResult(rootHints,
+                        root == null ? null : root.getRootId(),
+                        root == null ? null : root.getExtras());
+            } catch (RemoteException e) {
+                // Controller may be died prematurely.
+                // TODO(jaewan): Handle this.
+            }
+        });
     }
 
     @Deprecated
@@ -250,131 +329,4 @@
             // TODO(jaewan): What to do when the controller is gone?
         }
     }
-
-    // TODO(jaewan): Remove this. We should use Executor given by the session builder.
-    private class CommandHandler extends Handler {
-        public static final int MSG_CONNECT = 1000;
-        public static final int MSG_COMMAND = 1001;
-        public static final int MSG_ON_GET_ROOT = 2000;
-
-        public CommandHandler(Looper looper) {
-            super(looper);
-        }
-
-        @Override
-        public void handleMessage(Message msg) {
-            final MediaSession2Impl session = MediaSession2Stub.this.mSession.get();
-            if (session == null || session.getPlayer() == null) {
-                return;
-            }
-
-            switch (msg.what) {
-                case MSG_CONNECT: {
-                    ControllerInfo request = (ControllerInfo) msg.obj;
-                    CommandGroup allowedCommands = mSessionCallback.onConnect(request);
-                    // Don't reject connection for the request from trusted app.
-                    // Otherwise server will fail to retrieve session's information to dispatch
-                    // media keys to.
-                    boolean accept = allowedCommands != null || request.isTrusted();
-                    ControllerInfoImpl impl = ControllerInfoImpl.from(request);
-                    if (accept) {
-                        synchronized (mLock) {
-                            mControllers.put(impl.getId(), request);
-                        }
-                        if (allowedCommands == null) {
-                            // For trusted apps, send non-null allowed commands to keep connection.
-                            allowedCommands = new CommandGroup();
-                        }
-                    }
-                    if (DEBUG) {
-                        Log.d(TAG, "onConnectResult, request=" + request
-                                + " accept=" + accept);
-                    }
-                    try {
-                        impl.getControllerBinder().onConnectionChanged(
-                                accept ? MediaSession2Stub.this : null,
-                                allowedCommands == null ? null : allowedCommands.toBundle());
-                    } catch (RemoteException e) {
-                        // Controller may be died prematurely.
-                    }
-                    break;
-                }
-                case MSG_COMMAND: {
-                    CommandParam param = (CommandParam) msg.obj;
-                    Command command = param.command;
-                    boolean accepted = mSessionCallback.onCommandRequest(
-                            param.controller, command);
-                    if (!accepted) {
-                        // Don't run rejected command.
-                        if (DEBUG) {
-                            Log.d(TAG, "Command " + command + " from "
-                                    + param.controller + " was rejected by " + session);
-                        }
-                        return;
-                    }
-
-                    switch (param.command.getCommandCode()) {
-                        case MediaSession2.COMMAND_CODE_PLAYBACK_START:
-                            session.getInstance().play();
-                            break;
-                        case MediaSession2.COMMAND_CODE_PLAYBACK_PAUSE:
-                            session.getInstance().pause();
-                            break;
-                        case MediaSession2.COMMAND_CODE_PLAYBACK_STOP:
-                            session.getInstance().stop();
-                            break;
-                        case MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_PREV_ITEM:
-                            session.getInstance().skipToPrevious();
-                            break;
-                        case MediaSession2.COMMAND_CODE_PLAYBACK_SKIP_NEXT_ITEM:
-                            session.getInstance().skipToNext();
-                            break;
-                        default:
-                            // TODO(jaewan): Handle custom command.
-                    }
-                    break;
-                }
-                case MSG_ON_GET_ROOT: {
-                    final CommandParam param = (CommandParam) msg.obj;
-                    final ControllerInfoImpl controller = ControllerInfoImpl.from(param.controller);
-                    BrowserRoot root = mLibraryCallback.onGetRoot(param.controller, param.args);
-                    try {
-                        controller.getControllerBinder().onGetRootResult(param.args,
-                                root == null ? null : root.getRootId(),
-                                root == null ? null : root.getExtras());
-                    } catch (RemoteException e) {
-                        // Controller may be died prematurely.
-                        // TODO(jaewan): Handle this.
-                    }
-                    break;
-                }
-            }
-        }
-
-        public void postConnect(ControllerInfo request) {
-            obtainMessage(MSG_CONNECT, request).sendToTarget();
-        }
-
-        public void postCommand(ControllerInfo controller, Command command, Bundle args) {
-            CommandParam param = new CommandParam(controller, command, args);
-            obtainMessage(MSG_COMMAND, param).sendToTarget();
-        }
-
-        public void postOnGetRoot(ControllerInfo controller, Bundle rootHints) {
-            CommandParam param = new CommandParam(controller, null, rootHints);
-            obtainMessage(MSG_ON_GET_ROOT, param).sendToTarget();
-        }
-    }
-
-    private static class CommandParam {
-        public final ControllerInfo controller;
-        public final Command command;
-        public final Bundle args;
-
-        private CommandParam(ControllerInfo controller, Command command, Bundle args) {
-            this.controller = controller;
-            this.command = command;
-            this.args = args;
-        }
-    }
 }
diff --git a/packages/MediaComponents/src/com/android/media/MediaSessionService2Impl.java b/packages/MediaComponents/src/com/android/media/MediaSessionService2Impl.java
index 9d24082..b53ad72 100644
--- a/packages/MediaComponents/src/com/android/media/MediaSessionService2Impl.java
+++ b/packages/MediaComponents/src/com/android/media/MediaSessionService2Impl.java
@@ -152,11 +152,6 @@
                 return;
             }
             MediaSession2Impl impl = (MediaSession2Impl) mSession.getProvider();
-            if (impl.getHandler().getLooper() != Looper.myLooper()) {
-                Log.w(TAG, "Ignoring " + state + ". Expected " + impl.getHandler().getLooper()
-                        + " but " + Looper.myLooper());
-                return;
-            }
             updateNotification(state);
         }
     }
diff --git a/packages/MediaComponents/test/src/android/media/MediaController2Test.java b/packages/MediaComponents/test/src/android/media/MediaController2Test.java
index ae67a95..09df42b 100644
--- a/packages/MediaComponents/test/src/android/media/MediaController2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaController2Test.java
@@ -45,6 +45,7 @@
  */
 // TODO(jaewan): Implement host-side test so controller and session can run in different processes.
 // TODO(jaewan): Fix flaky failure -- see MediaController2Impl.getController()
+// TODO(jaeawn): Revisit create/close session in the sHandler. It's no longer necessary.
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 @FlakyTest
@@ -60,11 +61,10 @@
     public void setUp() throws Exception {
         super.setUp();
         // Create this test specific MediaSession2 to use our own Handler.
-        sHandler.postAndSync(()->{
-            mPlayer = new MockPlayer(1);
-            mSession = new MediaSession2.Builder(mContext, mPlayer).setId(TAG).build();
-        });
-
+        mPlayer = new MockPlayer(1);
+        mSession = new MediaSession2.Builder(mContext, mPlayer)
+                .setSessionCallback(sHandlerExecutor, new SessionCallback())
+                .setId(TAG).build();
         mController = createController(mSession.getToken());
         TestServiceRegistry.getInstance().setHandler(sHandler);
     }
@@ -73,11 +73,9 @@
     @Override
     public void cleanUp() throws Exception {
         super.cleanUp();
-        sHandler.postAndSync(() -> {
-            if (mSession != null) {
-                mSession.close();
-            }
-        });
+        if (mSession != null) {
+            mSession.close();
+        }
         TestServiceRegistry.getInstance().cleanUp();
     }
 
@@ -275,6 +273,7 @@
             final MockPlayer player = new MockPlayer(0);
             sessionHandler.postAndSync(() -> {
                 mSession = new MediaSession2.Builder(mContext, mPlayer)
+                        .setSessionCallback(sHandlerExecutor, new SessionCallback())
                         .setId("testDeadlock").build();
             });
             final MediaController2 controller = createController(mSession.getToken());
@@ -462,7 +461,9 @@
         sHandler.postAndSync(() -> {
             // Recreated session has different session stub, so previously created controller
             // shouldn't be available.
-            mSession = new MediaSession2.Builder(mContext, mPlayer).setId(id).build();
+            mSession = new MediaSession2.Builder(mContext, mPlayer)
+                    .setSessionCallback(sHandlerExecutor, new SessionCallback())
+                    .setId(id).build();
         });
         testNoInteraction();
     }
diff --git a/packages/MediaComponents/test/src/android/media/MediaSession2Test.java b/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
index 045dcd5..35becbb 100644
--- a/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
+++ b/packages/MediaComponents/test/src/android/media/MediaSession2Test.java
@@ -54,19 +54,16 @@
     @Override
     public void setUp() throws Exception {
         super.setUp();
-        sHandler.postAndSync(() -> {
-            mPlayer = new MockPlayer(0);
-            mSession = new MediaSession2.Builder(mContext, mPlayer).build();
-        });
+        mPlayer = new MockPlayer(0);
+        mSession = new MediaSession2.Builder(mContext, mPlayer)
+                .setSessionCallback(sHandlerExecutor, new SessionCallback()).build();
     }
 
     @After
     @Override
     public void cleanUp() throws Exception {
         super.cleanUp();
-        sHandler.postAndSync(() -> {
-            mSession.close();
-        });
+        mSession.close();
     }
 
     @Test
diff --git a/packages/MediaComponents/test/src/android/media/MediaSessionManager_MediaSession2.java b/packages/MediaComponents/test/src/android/media/MediaSessionManager_MediaSession2.java
index 192cbc2..6037619 100644
--- a/packages/MediaComponents/test/src/android/media/MediaSessionManager_MediaSession2.java
+++ b/packages/MediaComponents/test/src/android/media/MediaSessionManager_MediaSession2.java
@@ -59,9 +59,8 @@
         // Specify TAG here so {@link MediaSession2.getInstance()} doesn't complaint about
         // per test thread differs across the {@link MediaSession2} with the same TAG.
         final MockPlayer player = new MockPlayer(1);
-        sHandler.postAndSync(() -> {
-            mSession = new MediaSession2.Builder(mContext, player).setId(TAG).build();
-        });
+        mSession = new MediaSession2.Builder(mContext, player)
+                .setSessionCallback(sHandlerExecutor, new SessionCallback()).setId(TAG).build();
         ensureChangeInSession();
     }
 
@@ -70,9 +69,7 @@
     public void cleanUp() throws Exception {
         super.cleanUp();
         sHandler.removeCallbacksAndMessages(null);
-        sHandler.postAndSync(() -> {
-            mSession.close();
-        });
+        mSession.close();
     }
 
     // TODO(jaewan): Make this host-side test to see per-user behavior.
diff --git a/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java b/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java
index b058117..5c5c7d2 100644
--- a/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java
+++ b/packages/MediaComponents/test/src/android/media/MockMediaSessionService2.java
@@ -28,6 +28,8 @@
 import android.media.session.PlaybackState;
 import android.os.Process;
 
+import java.util.concurrent.Executor;
+
 /**
  * Mock implementation of {@link android.media.MediaSessionService2} for testing.
  */
@@ -46,11 +48,12 @@
     public MediaSession2 onCreateSession(String sessionId) {
         final MockPlayer player = new MockPlayer(1);
         final SyncHandler handler = (SyncHandler) TestServiceRegistry.getInstance().getHandler();
+        final Executor executor = (runnable) -> handler.post(runnable);
         try {
             handler.postAndSync(() -> {
                 mSession = new MediaSession2.Builder(MockMediaSessionService2.this, player)
-                        .setId(sessionId).setSessionCallback((runnable)->handler.post(runnable),
-                                new MySessionCallback()).build();
+                        .setSessionCallback(executor, new MySessionCallback())
+                        .setId(sessionId).build();
             });
         } catch (InterruptedException e) {
             fail(e.toString());