Use SurfaceControl as token to register a SurfaceControlInputReceiver

Instead of getting back an IBinder that's used to unregister the
receiver, the SC itself can be the id into the map. The client token
still needs to be created and retrieved for testing to ensure we match
the token when looking for the layer in the window infos list.

Test: SurfaceControlInputReceiverTests
Bug: 320506582
Change-Id: I516bbff349c7ad3639ab14350d96112c22e1d6dc

7 files changed, 133 insertions, 88 deletions

diff --git a/core/api/current.txt b/core/api/current.txt
index 0e413c41a64a..c5d9989875b5 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -53587,13 +53587,13 @@ package android.view {
     method @Deprecated public android.view.Display getDefaultDisplay();
     method @NonNull public default android.view.WindowMetrics getMaximumWindowMetrics();
     method public default boolean isCrossWindowBlurEnabled();
-    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") @NonNull public default android.os.IBinder registerBatchedSurfaceControlInputReceiver(int, @NonNull android.os.IBinder, @NonNull android.view.SurfaceControl, @NonNull android.view.Choreographer, @NonNull android.view.SurfaceControlInputReceiver);
+    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") public default void registerBatchedSurfaceControlInputReceiver(int, @NonNull android.os.IBinder, @NonNull android.view.SurfaceControl, @NonNull android.view.Choreographer, @NonNull android.view.SurfaceControlInputReceiver);
     method @FlaggedApi("com.android.window.flags.trusted_presentation_listener_for_window") public default void registerTrustedPresentationListener(@NonNull android.os.IBinder, @NonNull android.window.TrustedPresentationThresholds, @NonNull java.util.concurrent.Executor, @NonNull java.util.function.Consumer<java.lang.Boolean>);
-    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") @NonNull public default android.os.IBinder registerUnbatchedSurfaceControlInputReceiver(int, @NonNull android.os.IBinder, @NonNull android.view.SurfaceControl, @NonNull android.os.Looper, @NonNull android.view.SurfaceControlInputReceiver);
+    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") public default void registerUnbatchedSurfaceControlInputReceiver(int, @NonNull android.os.IBinder, @NonNull android.view.SurfaceControl, @NonNull android.os.Looper, @NonNull android.view.SurfaceControlInputReceiver);
     method public default void removeCrossWindowBlurEnabledListener(@NonNull java.util.function.Consumer<java.lang.Boolean>);
     method public default void removeProposedRotationListener(@NonNull java.util.function.IntConsumer);
     method public void removeViewImmediate(android.view.View);
-    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") public default void unregisterSurfaceControlInputReceiver(@NonNull android.os.IBinder);
+    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") public default void unregisterSurfaceControlInputReceiver(@NonNull android.view.SurfaceControl);
     method @FlaggedApi("com.android.window.flags.trusted_presentation_listener_for_window") public default void unregisterTrustedPresentationListener(@NonNull java.util.function.Consumer<java.lang.Boolean>);
     field public static final String PROPERTY_ACTIVITY_EMBEDDING_ALLOW_SYSTEM_OVERRIDE = "android.window.PROPERTY_ACTIVITY_EMBEDDING_ALLOW_SYSTEM_OVERRIDE";
     field public static final String PROPERTY_ACTIVITY_EMBEDDING_SPLITS_ENABLED = "android.window.PROPERTY_ACTIVITY_EMBEDDING_SPLITS_ENABLED";
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 5d271cc10a69..5cf320ea4a54 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -3637,6 +3637,7 @@ package android.view {
 
   public interface WindowManager extends android.view.ViewManager {
     method public default int getDisplayImePolicy(int);
+    method @FlaggedApi("com.android.window.flags.surface_control_input_receiver") @Nullable public default android.os.IBinder getSurfaceControlInputClientToken(@NonNull android.view.SurfaceControl);
     method public static boolean hasWindowExtensionsEnabled();
     method public default void holdLock(android.os.IBinder, int);
     method public default boolean isGlobalKey(int);
diff --git a/core/java/android/view/WindowManager.java b/core/java/android/view/WindowManager.java
index 42355bb17c2d..427d053f754e 100644
--- a/core/java/android/view/WindowManager.java
+++ b/core/java/android/view/WindowManager.java
@@ -6022,8 +6022,8 @@ public interface WindowManager extends ViewManager {
      * This is different from
      * {@link #registerUnbatchedSurfaceControlInputReceiver(int, IBinder, SurfaceControl, Looper,
      * SurfaceControlInputReceiver)} in that the input events are received batched. The caller must
-     * invoke {@link #unregisterSurfaceControlInputReceiver(IBinder)} to clean up the resources when
-     * no longer needing to use the {@link SurfaceControlInputReceiver}
+     * invoke {@link #unregisterSurfaceControlInputReceiver(SurfaceControl)} to clean up the
+     * resources when no longer needing to use the {@link SurfaceControlInputReceiver}
      *
      * @param displayId      The display that the SurfaceControl will be placed on. Input will
      *                       only work
@@ -6035,14 +6035,9 @@ public interface WindowManager extends ViewManager {
      * @param choreographer  The Choreographer used for batching. This should match the rendering
      *                       Choreographer.
      * @param receiver       The SurfaceControlInputReceiver that will receive the input events
-     * @return an {@link IBinder} token that is used to unregister the input receiver via
-     * {@link #unregisterSurfaceControlInputReceiver(IBinder)}.
-     * @see #registerUnbatchedSurfaceControlInputReceiver(int, IBinder, SurfaceControl, Looper,
-     * SurfaceControlInputReceiver)
      */
     @FlaggedApi(Flags.FLAG_SURFACE_CONTROL_INPUT_RECEIVER)
-    @NonNull
-    default IBinder registerBatchedSurfaceControlInputReceiver(int displayId,
+    default void registerBatchedSurfaceControlInputReceiver(int displayId,
             @NonNull IBinder hostToken, @NonNull SurfaceControl surfaceControl,
             @NonNull Choreographer choreographer, @NonNull SurfaceControlInputReceiver receiver) {
         throw new UnsupportedOperationException(
@@ -6054,8 +6049,8 @@ public interface WindowManager extends ViewManager {
      * receive every input event. This is different than calling @link
      * #registerBatchedSurfaceControlInputReceiver(int, IBinder, SurfaceControl, Choreographer,
      * SurfaceControlInputReceiver)} in that the input events are received unbatched. The caller
-     * must invoke {@link #unregisterSurfaceControlInputReceiver(IBinder)} to clean up the resources
-     * when no longer needing to use the {@link SurfaceControlInputReceiver}
+     * must invoke {@link #unregisterSurfaceControlInputReceiver(SurfaceControl)} to clean up the
+     * resources when no longer needing to use the {@link SurfaceControlInputReceiver}
      *
      * @param displayId      The display that the SurfaceControl will be placed on. Input will only
      *                       work if SurfaceControl is on that display and that display was
@@ -6066,14 +6061,9 @@ public interface WindowManager extends ViewManager {
      * @param surfaceControl The SurfaceControl to register the InputChannel for
      * @param looper         The looper to use when invoking callbacks.
      * @param receiver       The SurfaceControlInputReceiver that will receive the input events
-     * @return an {@link IBinder} token that is used to unregister the input receiver via
-     * {@link #unregisterSurfaceControlInputReceiver(IBinder)}.
-     * @see #registerBatchedSurfaceControlInputReceiver(int, IBinder, SurfaceControl, Choreographer,
-     * SurfaceControlInputReceiver)
      **/
     @FlaggedApi(Flags.FLAG_SURFACE_CONTROL_INPUT_RECEIVER)
-    @NonNull
-    default IBinder registerUnbatchedSurfaceControlInputReceiver(int displayId,
+    default void registerUnbatchedSurfaceControlInputReceiver(int displayId,
             @NonNull IBinder hostToken, @NonNull SurfaceControl surfaceControl,
             @NonNull Looper looper, @NonNull SurfaceControlInputReceiver receiver) {
         throw new UnsupportedOperationException(
@@ -6091,17 +6081,32 @@ public interface WindowManager extends ViewManager {
      * {@link #registerUnbatchedSurfaceControlInputReceiver(int, IBinder, SurfaceControl, Looper,
      * SurfaceControlInputReceiver)}
      *
-     * @param token The token that was returned via
-     *              {@link #registerBatchedSurfaceControlInputReceiver(int, IBinder,
-     *              SurfaceControl,
-     *              Choreographer, SurfaceControlInputReceiver)} or
-     *              {@link #registerUnbatchedSurfaceControlInputReceiver(int, IBinder,
-     *              SurfaceControl,
-     *              Looper, SurfaceControlInputReceiver)}
+     * @param surfaceControl The SurfaceControl to remove and unregister the input channel for.
      */
     @FlaggedApi(Flags.FLAG_SURFACE_CONTROL_INPUT_RECEIVER)
-    default void unregisterSurfaceControlInputReceiver(@NonNull IBinder token) {
+    default void unregisterSurfaceControlInputReceiver(@NonNull SurfaceControl surfaceControl) {
         throw new UnsupportedOperationException(
                 "unregisterSurfaceControlInputReceiver is not implemented");
     }
+
+    /**
+     * Returns the input client token for the {@link SurfaceControl}. This will only return non null
+     * if the SurfaceControl was registered for input via
+     * { #registerBatchedSurfaceControlInputReceiver(int, IBinder, SurfaceControl, Choreographer,
+     * SurfaceControlInputReceiver)} or
+     * {@link #registerUnbatchedSurfaceControlInputReceiver(int, IBinder, SurfaceControl, Looper,
+     * SurfaceControlInputReceiver)}.
+     * <p>
+     * This is helpful for testing to ensure the test waits for the layer to be registered with
+     * SurfaceFlinger and Input before proceeding with the test.
+     *
+     * @hide
+     */
+    @FlaggedApi(Flags.FLAG_SURFACE_CONTROL_INPUT_RECEIVER)
+    @TestApi
+    @Nullable
+    default IBinder getSurfaceControlInputClientToken(@NonNull SurfaceControl surfaceControl) {
+        throw new UnsupportedOperationException(
+                "getSurfaceControlInputClientToken is not implemented");
+    }
 }
diff --git a/core/java/android/view/WindowManagerGlobal.java b/core/java/android/view/WindowManagerGlobal.java
index 8d40f9a4f7b1..c49fce5b558f 100644
--- a/core/java/android/view/WindowManagerGlobal.java
+++ b/core/java/android/view/WindowManagerGlobal.java
@@ -38,10 +38,12 @@ import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
 import android.util.Pair;
+import android.util.SparseArray;
 import android.view.inputmethod.InputMethodManager;
 import android.window.ITrustedPresentationListener;
 import android.window.TrustedPresentationThresholds;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.FastPrintWriter;
 
 import java.io.FileDescriptor;
@@ -50,7 +52,6 @@ import java.io.PrintWriter;
 import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.WeakHashMap;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 import java.util.function.IntConsumer;
@@ -156,8 +157,9 @@ public final class WindowManagerGlobal {
     private final TrustedPresentationListener mTrustedPresentationListener =
             new TrustedPresentationListener();
 
-    private final ConcurrentHashMap<IBinder, InputEventReceiver> mSurfaceControlInputReceivers =
-            new ConcurrentHashMap<>();
+    @GuardedBy("mSurfaceControlInputReceivers")
+    private final SparseArray<SurfaceControlInputReceiverInfo>
+            mSurfaceControlInputReceivers = new SparseArray<>();
 
     private WindowManagerGlobal() {
     }
@@ -816,7 +818,7 @@ public final class WindowManagerGlobal {
         mTrustedPresentationListener.removeListener(listener);
     }
 
-    IBinder registerBatchedSurfaceControlInputReceiver(int displayId,
+    void registerBatchedSurfaceControlInputReceiver(int displayId,
             @NonNull IBinder hostToken, @NonNull SurfaceControl surfaceControl,
             @NonNull Choreographer choreographer, @NonNull SurfaceControlInputReceiver receiver) {
         IBinder clientToken = new Binder();
@@ -830,19 +832,21 @@ public final class WindowManagerGlobal {
             e.rethrowAsRuntimeException();
         }
 
-        mSurfaceControlInputReceivers.put(clientToken,
-                new BatchedInputEventReceiver(inputChannel, choreographer.getLooper(),
-                        choreographer) {
-                    @Override
-                    public void onInputEvent(InputEvent event) {
-                        boolean handled = receiver.onInputEvent(event);
-                        finishInputEvent(event, handled);
-                    }
-                });
-        return clientToken;
+        synchronized (mSurfaceControlInputReceivers) {
+            mSurfaceControlInputReceivers.put(surfaceControl.getLayerId(),
+                    new SurfaceControlInputReceiverInfo(clientToken,
+                            new BatchedInputEventReceiver(inputChannel, choreographer.getLooper(),
+                                    choreographer) {
+                                @Override
+                                public void onInputEvent(InputEvent event) {
+                                    boolean handled = receiver.onInputEvent(event);
+                                    finishInputEvent(event, handled);
+                                }
+                            }));
+        }
     }
 
-    IBinder registerUnbatchedSurfaceControlInputReceiver(
+    void registerUnbatchedSurfaceControlInputReceiver(
             int displayId, @NonNull IBinder hostToken, @NonNull SurfaceControl surfaceControl,
             @NonNull Looper looper, @NonNull SurfaceControlInputReceiver receiver) {
         IBinder clientToken = new Binder();
@@ -856,32 +860,53 @@ public final class WindowManagerGlobal {
             e.rethrowAsRuntimeException();
         }
 
-        mSurfaceControlInputReceivers.put(clientToken,
-                new InputEventReceiver(inputChannel, looper) {
-                    @Override
-                    public void onInputEvent(InputEvent event) {
-                        boolean handled = receiver.onInputEvent(event);
-                        finishInputEvent(event, handled);
-                    }
-                });
-
-        return clientToken;
+        synchronized (mSurfaceControlInputReceivers) {
+            mSurfaceControlInputReceivers.put(surfaceControl.getLayerId(),
+                    new SurfaceControlInputReceiverInfo(clientToken,
+                            new InputEventReceiver(inputChannel, looper) {
+                                @Override
+                                public void onInputEvent(InputEvent event) {
+                                    boolean handled = receiver.onInputEvent(event);
+                                    finishInputEvent(event, handled);
+                                }
+                            }));
+        }
     }
 
-    void unregisterSurfaceControlInputReceiver(IBinder token) {
-        InputEventReceiver inputEventReceiver = mSurfaceControlInputReceivers.get(token);
-        if (inputEventReceiver == null) {
-            Log.w(TAG, "No registered input event receiver with token: " + token);
+    void unregisterSurfaceControlInputReceiver(SurfaceControl surfaceControl) {
+        SurfaceControlInputReceiverInfo surfaceControlInputReceiverInfo;
+        synchronized (mSurfaceControlInputReceivers) {
+            surfaceControlInputReceiverInfo = mSurfaceControlInputReceivers.removeReturnOld(
+                    surfaceControl.getLayerId());
+        }
+
+        if (surfaceControlInputReceiverInfo == null) {
+            Log.w(TAG, "No registered input event receiver with sc: " + surfaceControl);
             return;
         }
         try {
-            WindowManagerGlobal.getWindowSession().remove(token);
+            WindowManagerGlobal.getWindowSession().remove(
+                    surfaceControlInputReceiverInfo.mClientToken);
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to remove input channel", e);
             e.rethrowAsRuntimeException();
         }
 
-        inputEventReceiver.dispose();
+        surfaceControlInputReceiverInfo.mInputEventReceiver.dispose();
+    }
+
+    IBinder getSurfaceControlInputClientToken(SurfaceControl surfaceControl) {
+        SurfaceControlInputReceiverInfo surfaceControlInputReceiverInfo;
+        synchronized (mSurfaceControlInputReceivers) {
+            surfaceControlInputReceiverInfo = mSurfaceControlInputReceivers.get(
+                    surfaceControl.getLayerId());
+        }
+
+        if (surfaceControlInputReceiverInfo == null) {
+            Log.w(TAG, "No registered input event receiver with sc: " + surfaceControl);
+            return null;
+        }
+        return surfaceControlInputReceiverInfo.mClientToken;
     }
 
     private final class TrustedPresentationListener extends
@@ -976,6 +1001,17 @@ public final class WindowManagerGlobal {
             throw e.rethrowFromSystemServer();
         }
     }
+
+    private static class SurfaceControlInputReceiverInfo {
+        final IBinder mClientToken;
+        final InputEventReceiver mInputEventReceiver;
+
+        private SurfaceControlInputReceiverInfo(IBinder clientToken,
+                InputEventReceiver inputEventReceiver) {
+            mClientToken = clientToken;
+            mInputEventReceiver = inputEventReceiver;
+        }
+    }
 }
 
 final class WindowLeaked extends AndroidRuntimeException {
diff --git a/core/java/android/view/WindowManagerImpl.java b/core/java/android/view/WindowManagerImpl.java
index aaf5fcc6f095..41d181c1b10c 100644
--- a/core/java/android/view/WindowManagerImpl.java
+++ b/core/java/android/view/WindowManagerImpl.java
@@ -523,26 +523,30 @@ public final class WindowManagerImpl implements WindowManager {
         mGlobal.unregisterTrustedPresentationListener(listener);
     }
 
-    @NonNull
     @Override
-    public IBinder registerBatchedSurfaceControlInputReceiver(int displayId,
+    public void registerBatchedSurfaceControlInputReceiver(int displayId,
             @NonNull IBinder hostToken, @NonNull SurfaceControl surfaceControl,
             @NonNull Choreographer choreographer, @NonNull SurfaceControlInputReceiver receiver) {
-        return mGlobal.registerBatchedSurfaceControlInputReceiver(displayId, hostToken,
+        mGlobal.registerBatchedSurfaceControlInputReceiver(displayId, hostToken,
                 surfaceControl, choreographer, receiver);
     }
 
-    @NonNull
     @Override
-    public IBinder registerUnbatchedSurfaceControlInputReceiver(
+    public void registerUnbatchedSurfaceControlInputReceiver(
             int displayId, @NonNull IBinder hostToken, @NonNull SurfaceControl surfaceControl,
             @NonNull Looper looper, @NonNull SurfaceControlInputReceiver receiver) {
-        return mGlobal.registerUnbatchedSurfaceControlInputReceiver(displayId, hostToken,
+        mGlobal.registerUnbatchedSurfaceControlInputReceiver(displayId, hostToken,
                 surfaceControl, looper, receiver);
     }
 
     @Override
-    public void unregisterSurfaceControlInputReceiver(@NonNull IBinder token) {
-        mGlobal.unregisterSurfaceControlInputReceiver(token);
+    public void unregisterSurfaceControlInputReceiver(@NonNull SurfaceControl surfaceControl) {
+        mGlobal.unregisterSurfaceControlInputReceiver(surfaceControl);
+    }
+
+    @Override
+    @Nullable
+    public IBinder getSurfaceControlInputClientToken(@NonNull SurfaceControl surfaceControl) {
+        return mGlobal.getSurfaceControlInputClientToken(surfaceControl);
     }
 }
diff --git a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java
index 5aaf30a5b3a7..14230fe4c323 100644
--- a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java
+++ b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/EmbeddedWindowService.java
@@ -134,7 +134,7 @@ public class EmbeddedWindowService extends Service {
                 c.drawText("Remote", 250, 250, paint);
                 surface.unlockCanvasAndPost(c);
                 WindowManager wm = getSystemService(WindowManager.class);
-                mInputToken = wm.registerBatchedSurfaceControlInputReceiver(displayId, hostToken,
+                wm.registerBatchedSurfaceControlInputReceiver(displayId, hostToken,
                         mSurfaceControl,
                         Choreographer.getInstance(), event -> {
                             Log.d(TAG, "onInputEvent-remote " + event);
@@ -147,11 +147,9 @@ public class EmbeddedWindowService extends Service {
         @Override
         public void tearDownEmbeddedSurfaceControl() {
             if (mSurfaceControl != null) {
-                new SurfaceControl.Transaction().remove(mSurfaceControl);
-            }
-            if (mInputToken != null) {
                 WindowManager wm = getSystemService(WindowManager.class);
-                wm.unregisterSurfaceControlInputReceiver(mInputToken);
+                wm.unregisterSurfaceControlInputReceiver(mSurfaceControl);
+                new SurfaceControl.Transaction().remove(mSurfaceControl).apply();
             }
         }
     }
diff --git a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java
index e5f8f47aeecd..7330ec14011b 100644
--- a/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java
+++ b/tests/SurfaceControlViewHostTest/src/com/android/test/viewembed/SurfaceInputTestActivity.java
@@ -50,10 +50,11 @@ public class SurfaceInputTestActivity extends Activity {
     private static final String TAG = "SurfaceInputTestActivity";
     private SurfaceView mLocalSurfaceView;
     private SurfaceView mRemoteSurfaceView;
-    private IBinder mInputToken;
     private IAttachEmbeddedWindow mIAttachEmbeddedWindow;
     private SurfaceControl mParentSurfaceControl;
 
+    private SurfaceControl mLocalSurfaceControl;
+
     private final ServiceConnection mConnection = new ServiceConnection() {
         // Called when the connection with the service is established
         public void onServiceConnected(ComponentName className, IBinder service) {
@@ -112,30 +113,33 @@ public class SurfaceInputTestActivity extends Activity {
     @Override
     protected void onDestroy() {
         super.onDestroy();
-        getWindowManager().unregisterSurfaceControlInputReceiver(mInputToken);
+        if (mLocalSurfaceControl != null) {
+            getWindowManager().unregisterSurfaceControlInputReceiver(mLocalSurfaceControl);
+            new SurfaceControl.Transaction().remove(mLocalSurfaceControl).apply();
+        }
     }
 
     private void addLocalChildSurfaceControl(AttachedSurfaceControl attachedSurfaceControl) {
-        SurfaceControl surfaceControl = new SurfaceControl.Builder().setName("LocalSC")
+        mLocalSurfaceControl = new SurfaceControl.Builder().setName("LocalSC")
                 .setBufferSize(100, 100).build();
-        attachedSurfaceControl.buildReparentTransaction(surfaceControl)
-                .setVisibility(surfaceControl, true)
-                .setCrop(surfaceControl, new Rect(0, 0, 100, 100))
-                .setPosition(surfaceControl, 250, 1000)
-                .setLayer(surfaceControl, 1).apply();
+        attachedSurfaceControl.buildReparentTransaction(mLocalSurfaceControl)
+                .setVisibility(mLocalSurfaceControl, true)
+                .setCrop(mLocalSurfaceControl, new Rect(0, 0, 100, 100))
+                .setPosition(mLocalSurfaceControl, 250, 1000)
+                .setLayer(mLocalSurfaceControl, 1).apply();
 
         Paint paint = new Paint();
         paint.setColor(Color.WHITE);
         paint.setTextSize(20);
 
-        Surface surface = new Surface(surfaceControl);
+        Surface surface = new Surface(mLocalSurfaceControl);
         Canvas c = surface.lockCanvas(null);
         c.drawColor(Color.GREEN);
         c.drawText("Local SC", 0, 0, paint);
         surface.unlockCanvasAndPost(c);
         WindowManager wm = getSystemService(WindowManager.class);
-        mInputToken = wm.registerBatchedSurfaceControlInputReceiver(getDisplayId(),
-                attachedSurfaceControl.getHostToken(), surfaceControl,
+        wm.registerBatchedSurfaceControlInputReceiver(getDisplayId(),
+                attachedSurfaceControl.getHostToken(), mLocalSurfaceControl,
                 Choreographer.getInstance(), event -> {
                     Log.d(TAG, "onInputEvent-sc " + event);
                     return false;
@@ -143,8 +147,6 @@ public class SurfaceInputTestActivity extends Activity {
     }
 
     private final SurfaceHolder.Callback mLocalSurfaceViewCallback = new SurfaceHolder.Callback() {
-        private IBinder mInputToken;
-
         @Override
         public void surfaceCreated(@NonNull SurfaceHolder holder) {
             Paint paint = new Paint();
@@ -157,7 +159,7 @@ public class SurfaceInputTestActivity extends Activity {
             holder.unlockCanvasAndPost(c);
 
             WindowManager wm = getSystemService(WindowManager.class);
-            mInputToken = wm.registerBatchedSurfaceControlInputReceiver(getDisplayId(),
+            wm.registerBatchedSurfaceControlInputReceiver(getDisplayId(),
                     mLocalSurfaceView.getHostToken(), mLocalSurfaceView.getSurfaceControl(),
                     Choreographer.getInstance(), event -> {
                         Log.d(TAG, "onInputEvent-local " + event);
@@ -173,9 +175,8 @@ public class SurfaceInputTestActivity extends Activity {
 
         @Override
         public void surfaceDestroyed(@NonNull SurfaceHolder holder) {
-            if (mInputToken != null) {
-                getWindowManager().unregisterSurfaceControlInputReceiver(mInputToken);
-            }
+            getWindowManager().unregisterSurfaceControlInputReceiver(
+                    mLocalSurfaceView.getSurfaceControl());
         }
     };