Merge "Let blocked InputConnection APIs fail upon IInputMethod.unbindInput()" into rvc-dev

10 files changed, 791 insertions, 497 deletions

diff --git a/core/java/android/inputmethodservice/IInputMethodWrapper.java b/core/java/android/inputmethodservice/IInputMethodWrapper.java
index b52b437b4557..a298c856a0fb 100644
--- a/core/java/android/inputmethodservice/IInputMethodWrapper.java
+++ b/core/java/android/inputmethodservice/IInputMethodWrapper.java
@@ -18,6 +18,7 @@ package android.inputmethodservice;
 
 import android.annotation.BinderThread;
 import android.annotation.MainThread;
+import android.annotation.Nullable;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.content.pm.PackageManager;
@@ -37,6 +38,7 @@ import android.view.inputmethod.InputMethodSession;
 import android.view.inputmethod.InputMethodSubtype;
 
 import com.android.internal.inputmethod.IInputMethodPrivilegedOperations;
+import com.android.internal.inputmethod.CancellationGroup;
 import com.android.internal.os.HandlerCaller;
 import com.android.internal.os.SomeArgs;
 import com.android.internal.view.IInlineSuggestionsRequestCallback;
@@ -52,7 +54,6 @@ import java.io.PrintWriter;
 import java.lang.ref.WeakReference;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Implements the internal IInputMethod interface to convert incoming calls
@@ -90,12 +91,13 @@ class IInputMethodWrapper extends IInputMethod.Stub
      *
      * <p>This field must be set and cleared only from the binder thread(s), where the system
      * guarantees that {@link #bindInput(InputBinding)},
-     * {@link #startInput(IBinder, IInputContext, int, EditorInfo, boolean)}, and
+     * {@link #startInput(IBinder, IInputContext, int, EditorInfo, boolean, boolean)}, and
      * {@link #unbindInput()} are called with the same order as the original calls
      * in {@link com.android.server.inputmethod.InputMethodManagerService}.
      * See {@link IBinder#FLAG_ONEWAY} for detailed semantics.</p>
      */
-    AtomicBoolean mIsUnbindIssued = null;
+    @Nullable
+    CancellationGroup mCancellationGroup = null;
 
     // NOTE: we should have a cache of these.
     static final class InputMethodSessionCallbackWrapper implements InputMethod.SessionCallback {
@@ -187,11 +189,11 @@ class IInputMethodWrapper extends IInputMethod.Stub
                 final IBinder startInputToken = (IBinder) args.arg1;
                 final IInputContext inputContext = (IInputContext) args.arg2;
                 final EditorInfo info = (EditorInfo) args.arg3;
-                final AtomicBoolean isUnbindIssued = (AtomicBoolean) args.arg4;
+                final CancellationGroup cancellationGroup = (CancellationGroup) args.arg4;
                 SomeArgs moreArgs = (SomeArgs) args.arg5;
                 final InputConnection ic = inputContext != null
                         ? new InputConnectionWrapper(
-                                mTarget, inputContext, moreArgs.argi3, isUnbindIssued)
+                                mTarget, inputContext, moreArgs.argi3, cancellationGroup)
                         : null;
                 info.makeCompatible(mTargetSdkVersion);
                 inputMethod.dispatchStartInputWithToken(
@@ -295,15 +297,15 @@ class IInputMethodWrapper extends IInputMethod.Stub
     @BinderThread
     @Override
     public void bindInput(InputBinding binding) {
-        if (mIsUnbindIssued != null) {
+        if (mCancellationGroup != null) {
             Log.e(TAG, "bindInput must be paired with unbindInput.");
         }
-        mIsUnbindIssued = new AtomicBoolean();
+        mCancellationGroup = new CancellationGroup();
         // This IInputContext is guaranteed to implement all the methods.
         final int missingMethodFlags = 0;
         InputConnection ic = new InputConnectionWrapper(mTarget,
                 IInputContext.Stub.asInterface(binding.getConnectionToken()), missingMethodFlags,
-                mIsUnbindIssued);
+                mCancellationGroup);
         InputBinding nu = new InputBinding(ic, binding);
         mCaller.executeOrSendMessage(mCaller.obtainMessageO(DO_SET_INPUT_CONTEXT, nu));
     }
@@ -311,10 +313,10 @@ class IInputMethodWrapper extends IInputMethod.Stub
     @BinderThread
     @Override
     public void unbindInput() {
-        if (mIsUnbindIssued != null) {
+        if (mCancellationGroup != null) {
             // Signal the flag then forget it.
-            mIsUnbindIssued.set(true);
-            mIsUnbindIssued = null;
+            mCancellationGroup.cancelAll();
+            mCancellationGroup = null;
         } else {
             Log.e(TAG, "unbindInput must be paired with bindInput.");
         }
@@ -326,16 +328,16 @@ class IInputMethodWrapper extends IInputMethod.Stub
     public void startInput(IBinder startInputToken, IInputContext inputContext,
             @InputConnectionInspector.MissingMethodFlags final int missingMethods,
             EditorInfo attribute, boolean restarting, boolean shouldPreRenderIme) {
-        if (mIsUnbindIssued == null) {
+        if (mCancellationGroup == null) {
             Log.e(TAG, "startInput must be called after bindInput.");
-            mIsUnbindIssued = new AtomicBoolean();
+            mCancellationGroup = new CancellationGroup();
         }
         SomeArgs args = SomeArgs.obtain();
         args.argi1 = restarting ? 1 : 0;
         args.argi2 = shouldPreRenderIme ? 1 : 0;
         args.argi3 = missingMethods;
-        mCaller.executeOrSendMessage(mCaller.obtainMessageOOOOO(
-                DO_START_INPUT, startInputToken, inputContext, attribute, mIsUnbindIssued, args));
+        mCaller.executeOrSendMessage(mCaller.obtainMessageOOOOO(DO_START_INPUT, startInputToken,
+                inputContext, attribute, mCancellationGroup, args));
     }
 
     @BinderThread
diff --git a/core/java/android/inputmethodservice/MultiClientInputMethodClientCallbackAdaptor.java b/core/java/android/inputmethodservice/MultiClientInputMethodClientCallbackAdaptor.java
index ef138a0c2217..dbb669be1402 100644
--- a/core/java/android/inputmethodservice/MultiClientInputMethodClientCallbackAdaptor.java
+++ b/core/java/android/inputmethodservice/MultiClientInputMethodClientCallbackAdaptor.java
@@ -39,6 +39,7 @@ import android.view.inputmethod.ExtractedText;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.inputmethod.IMultiClientInputMethodSession;
+import com.android.internal.inputmethod.CancellationGroup;
 import com.android.internal.os.SomeArgs;
 import com.android.internal.util.function.pooled.PooledLambda;
 import com.android.internal.view.IInputContext;
@@ -46,7 +47,6 @@ import com.android.internal.view.IInputMethodSession;
 import com.android.internal.view.InputConnectionWrapper;
 
 import java.lang.ref.WeakReference;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Re-dispatches all the incoming per-client events to the specified {@link Looper} thread.
@@ -80,19 +80,19 @@ final class MultiClientInputMethodClientCallbackAdaptor {
     @Nullable
     InputEventReceiver mInputEventReceiver;
 
-    private final AtomicBoolean mFinished = new AtomicBoolean(false);
+    private final CancellationGroup mCancellationGroup = new CancellationGroup();
 
     IInputMethodSession.Stub createIInputMethodSession() {
         synchronized (mSessionLock) {
             return new InputMethodSessionImpl(
-                    mSessionLock, mCallbackImpl, mHandler, mFinished);
+                    mSessionLock, mCallbackImpl, mHandler, mCancellationGroup);
         }
     }
 
     IMultiClientInputMethodSession.Stub createIMultiClientInputMethodSession() {
         synchronized (mSessionLock) {
             return new MultiClientInputMethodSessionImpl(
-                    mSessionLock, mCallbackImpl, mHandler, mFinished);
+                    mSessionLock, mCallbackImpl, mHandler, mCancellationGroup);
         }
     }
 
@@ -105,7 +105,7 @@ final class MultiClientInputMethodClientCallbackAdaptor {
             mHandler = new Handler(looper, null, true);
             mReadChannel = readChannel;
             mInputEventReceiver = new ImeInputEventReceiver(mReadChannel, mHandler.getLooper(),
-                    mFinished, mDispatcherState, mCallbackImpl.mOriginalCallback);
+                    mCancellationGroup, mDispatcherState, mCallbackImpl.mOriginalCallback);
         }
     }
 
@@ -139,16 +139,17 @@ final class MultiClientInputMethodClientCallbackAdaptor {
     }
 
     private static final class ImeInputEventReceiver extends InputEventReceiver {
-        private final AtomicBoolean mFinished;
+        private final CancellationGroup mCancellationGroupOnFinishSession;
         private final KeyEvent.DispatcherState mDispatcherState;
         private final MultiClientInputMethodServiceDelegate.ClientCallback mClientCallback;
         private final KeyEventCallbackAdaptor mKeyEventCallbackAdaptor;
 
-        ImeInputEventReceiver(InputChannel readChannel, Looper looper, AtomicBoolean finished,
+        ImeInputEventReceiver(InputChannel readChannel, Looper looper,
+                CancellationGroup cancellationGroupOnFinishSession,
                 KeyEvent.DispatcherState dispatcherState,
                 MultiClientInputMethodServiceDelegate.ClientCallback callback) {
             super(readChannel, looper);
-            mFinished = finished;
+            mCancellationGroupOnFinishSession = cancellationGroupOnFinishSession;
             mDispatcherState = dispatcherState;
             mClientCallback = callback;
             mKeyEventCallbackAdaptor = new KeyEventCallbackAdaptor(callback);
@@ -156,7 +157,7 @@ final class MultiClientInputMethodClientCallbackAdaptor {
 
         @Override
         public void onInputEvent(InputEvent event) {
-            if (mFinished.get()) {
+            if (mCancellationGroupOnFinishSession.isCanceled()) {
                 // The session has been finished.
                 finishInputEvent(event, false);
                 return;
@@ -187,14 +188,14 @@ final class MultiClientInputMethodClientCallbackAdaptor {
         private CallbackImpl mCallbackImpl;
         @GuardedBy("mSessionLock")
         private Handler mHandler;
-        private final AtomicBoolean mSessionFinished;
+        private final CancellationGroup mCancellationGroupOnFinishSession;
 
         InputMethodSessionImpl(Object lock, CallbackImpl callback, Handler handler,
-                AtomicBoolean sessionFinished) {
+                CancellationGroup cancellationGroupOnFinishSession) {
             mSessionLock = lock;
             mCallbackImpl = callback;
             mHandler = handler;
-            mSessionFinished = sessionFinished;
+            mCancellationGroupOnFinishSession = cancellationGroupOnFinishSession;
         }
 
         @Override
@@ -272,7 +273,7 @@ final class MultiClientInputMethodClientCallbackAdaptor {
                 if (mCallbackImpl == null || mHandler == null) {
                     return;
                 }
-                mSessionFinished.set(true);
+                mCancellationGroupOnFinishSession.cancelAll();
                 mHandler.sendMessage(PooledLambda.obtainMessage(
                         CallbackImpl::finishSession, mCallbackImpl));
                 mCallbackImpl = null;
@@ -311,14 +312,14 @@ final class MultiClientInputMethodClientCallbackAdaptor {
         private CallbackImpl mCallbackImpl;
         @GuardedBy("mSessionLock")
         private Handler mHandler;
-        private final AtomicBoolean mSessionFinished;
+        private final CancellationGroup mCancellationGroupOnFinishSession;
 
         MultiClientInputMethodSessionImpl(Object lock, CallbackImpl callback,
-                Handler handler, AtomicBoolean sessionFinished) {
+                Handler handler, CancellationGroup cancellationGroupOnFinishSession) {
             mSessionLock = lock;
             mCallbackImpl = callback;
             mHandler = handler;
-            mSessionFinished = sessionFinished;
+            mCancellationGroupOnFinishSession = cancellationGroupOnFinishSession;
         }
 
         @Override
@@ -335,7 +336,7 @@ final class MultiClientInputMethodClientCallbackAdaptor {
                         new WeakReference<>(null);
                 args.arg1 = (inputContext == null) ? null
                         : new InputConnectionWrapper(fakeIMS, inputContext, missingMethods,
-                                mSessionFinished);
+                                mCancellationGroupOnFinishSession);
                 args.arg2 = editorInfo;
                 args.argi1 = controlFlags;
                 args.argi2 = softInputMode;
diff --git a/core/java/com/android/internal/inputmethod/CancellationGroup.java b/core/java/com/android/internal/inputmethod/CancellationGroup.java
new file mode 100644
index 000000000000..09c9d128553b
--- /dev/null
+++ b/core/java/com/android/internal/inputmethod/CancellationGroup.java
@@ -0,0 +1,348 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.internal.inputmethod;
+
+import android.annotation.AnyThread;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+
+import com.android.internal.annotations.GuardedBy;
+
+import java.util.ArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * A utility class, which works as both a factory class of completable objects and a cancellation
+ * signal to cancel all the completable objects created by this object.
+ */
+public final class CancellationGroup {
+    private final Object mLock = new Object();
+
+    /**
+     * List of {@link CountDownLatch}, which can be used to propagate {@link #cancelAll()} to
+     * completable objects.
+     *
+     * <p>This will be lazily instantiated to avoid unnecessary object allocations.</p>
+     */
+    @Nullable
+    @GuardedBy("mLock")
+    private ArrayList<CountDownLatch> mLatchList = null;
+
+    @GuardedBy("mLock")
+    private boolean mCanceled = false;
+
+    /**
+     * An inner class to consolidate completable object types supported by
+     * {@link CancellationGroup}.
+     */
+    public static final class Completable {
+
+        /**
+         * Not intended to be instantiated.
+         */
+        private Completable() {
+        }
+
+        /**
+         * Base class of all the completable types supported by {@link CancellationGroup}.
+         */
+        protected static class ValueBase {
+            /**
+             * {@link CountDownLatch} to be signaled to unblock {@link #await(int, TimeUnit)}.
+             */
+            private final CountDownLatch mLatch = new CountDownLatch(1);
+
+            /**
+             * {@link CancellationGroup} to which this completable object belongs.
+             */
+            @NonNull
+            private final CancellationGroup mParentGroup;
+
+            /**
+             * Lock {@link Object} to guard complete operations within this class.
+             */
+            protected final Object mValueLock = new Object();
+
+            /**
+             * {@code true} after {@link #onComplete()} gets called.
+             */
+            @GuardedBy("mValueLock")
+            protected boolean mHasValue = false;
+
+            /**
+             * Base constructor.
+             *
+             * @param parentGroup {@link CancellationGroup} to which this completable object
+             *                    belongs.
+             */
+            protected ValueBase(@NonNull CancellationGroup parentGroup) {
+                mParentGroup = parentGroup;
+            }
+
+            /**
+             * @return {@link true} if {@link #onComplete()} gets called already.
+             */
+            @AnyThread
+            public boolean hasValue() {
+                synchronized (mValueLock) {
+                    return mHasValue;
+                }
+            }
+
+            /**
+             * Called by subclasses to signale {@link #mLatch}.
+             */
+            @AnyThread
+            protected void onComplete() {
+                mLatch.countDown();
+            }
+
+            /**
+             * Blocks the calling thread until at least one of the following conditions is met.
+             *
+             * <p>
+             *     <ol>
+             *         <li>This object becomes ready to return the value.</li>
+             *         <li>{@link CancellationGroup#cancelAll()} gets called.</li>
+             *         <li>The given timeout period has passed.</li>
+             *     </ol>
+             * </p>
+             *
+             * <p>The caller can distinguish the case 1 and case 2 by calling {@link #hasValue()}.
+             * Note that the return value of {@link #hasValue()} can change from {@code false} to
+             * {@code true} at any time, even after this methods finishes with returning
+             * {@code true}.</p>
+             *
+             * @param timeout length of the timeout.
+             * @param timeUnit unit of {@code timeout}.
+             * @return {@code false} if and only if the given timeout period has passed. Otherwise
+             *         {@code true}.
+             */
+            @AnyThread
+            public boolean await(int timeout, @NonNull TimeUnit timeUnit) {
+                if (!mParentGroup.registerLatch(mLatch)) {
+                    // Already canceled when this method gets called.
+                    return false;
+                }
+                try {
+                    return mLatch.await(timeout, timeUnit);
+                } catch (InterruptedException e) {
+                    return true;
+                } finally {
+                    mParentGroup.unregisterLatch(mLatch);
+                }
+            }
+        }
+
+        /**
+         * Completable object of integer primitive.
+         */
+        public static final class Int extends ValueBase {
+            @GuardedBy("mValueLock")
+            private int mValue = 0;
+
+            private Int(@NonNull CancellationGroup factory) {
+                super(factory);
+            }
+
+            /**
+             * Notify when a value is set to this completable object.
+             *
+             * @param value value to be set.
+             */
+            @AnyThread
+            void onComplete(int value) {
+                synchronized (mValueLock) {
+                    if (mHasValue) {
+                        throw new UnsupportedOperationException(
+                                "onComplete() cannot be called multiple times");
+                    }
+                    mValue = value;
+                    mHasValue = true;
+                }
+                onComplete();
+            }
+
+            /**
+             * @return value associated with this object.
+             * @throws UnsupportedOperationException when called while {@link #hasValue()} returns
+             *                                       {@code false}.
+             */
+            @AnyThread
+            public int getValue() {
+                synchronized (mValueLock) {
+                    if (!mHasValue) {
+                        throw new UnsupportedOperationException(
+                                "getValue() is allowed only if hasValue() returns true");
+                    }
+                    return mValue;
+                }
+            }
+        }
+
+        /**
+         * Base class of completable object types.
+         *
+         * @param <T> type associated with this completable object.
+         */
+        public static class Values<T> extends ValueBase {
+            @GuardedBy("mValueLock")
+            @Nullable
+            private T mValue = null;
+
+            protected Values(@NonNull CancellationGroup factory) {
+                super(factory);
+            }
+
+            /**
+             * Notify when a value is set to this completable value object.
+             *
+             * @param value value to be set.
+             */
+            @AnyThread
+            void onComplete(@Nullable T value) {
+                synchronized (mValueLock) {
+                    if (mHasValue) {
+                        throw new UnsupportedOperationException(
+                                "onComplete() cannot be called multiple times");
+                    }
+                    mValue = value;
+                    mHasValue = true;
+                }
+                onComplete();
+            }
+
+            /**
+             * @return value associated with this object.
+             * @throws UnsupportedOperationException when called while {@link #hasValue()} returns
+             *                                       {@code false}.
+             */
+            @AnyThread
+            @Nullable
+            public T getValue() {
+                synchronized (mValueLock) {
+                    if (!mHasValue) {
+                        throw new UnsupportedOperationException(
+                                "getValue() is allowed only if hasValue() returns true");
+                    }
+                    return mValue;
+                }
+            }
+        }
+
+        /**
+         * Completable object of {@link java.lang.CharSequence}.
+         */
+        public static final class CharSequence extends Values<java.lang.CharSequence> {
+            private CharSequence(@NonNull CancellationGroup factory) {
+                super(factory);
+            }
+        }
+
+        /**
+         * Completable object of {@link android.view.inputmethod.ExtractedText}.
+         */
+        public static final class ExtractedText
+                extends Values<android.view.inputmethod.ExtractedText> {
+            private ExtractedText(@NonNull CancellationGroup factory) {
+                super(factory);
+            }
+        }
+    }
+
+    /**
+     * @return an instance of {@link Completable.Int} that is associated with this
+     *         {@link CancellationGroup}.
+     */
+    @AnyThread
+    public Completable.Int createCompletableInt() {
+        return new Completable.Int(this);
+    }
+
+    /**
+     * @return an instance of {@link Completable.CharSequence} that is associated with this
+     *         {@link CancellationGroup}.
+     */
+    @AnyThread
+    public Completable.CharSequence createCompletableCharSequence() {
+        return new Completable.CharSequence(this);
+    }
+
+    /**
+     * @return an instance of {@link Completable.ExtractedText} that is associated with this
+     *         {@link CancellationGroup}.
+     */
+    @AnyThread
+    public Completable.ExtractedText createCompletableExtractedText() {
+        return new Completable.ExtractedText(this);
+    }
+
+    @AnyThread
+    private boolean registerLatch(@NonNull CountDownLatch latch) {
+        synchronized (mLock) {
+            if (mCanceled) {
+                return false;
+            }
+            if (mLatchList == null) {
+                // Set the initial capacity to 1 with an assumption that usually there is up to 1
+                // on-going operation.
+                mLatchList = new ArrayList<>(1);
+            }
+            mLatchList.add(latch);
+            return true;
+        }
+    }
+
+    @AnyThread
+    private void unregisterLatch(@NonNull CountDownLatch latch) {
+        synchronized (mLock) {
+            if (mLatchList != null) {
+                mLatchList.remove(latch);
+            }
+        }
+    }
+
+    /**
+     * Cancel all the completable objects created from this {@link CancellationGroup}.
+     *
+     * <p>Secondary calls will be silently ignored.</p>
+     */
+    @AnyThread
+    public void cancelAll() {
+        synchronized (mLock) {
+            if (!mCanceled) {
+                mCanceled = true;
+                if (mLatchList != null) {
+                    mLatchList.forEach(CountDownLatch::countDown);
+                    mLatchList.clear();
+                    mLatchList = null;
+                }
+            }
+        }
+    }
+
+    /**
+     * @return {@code true} if {@link #cancelAll()} is already called. {@code false} otherwise.
+     */
+    @AnyThread
+    public boolean isCanceled() {
+        synchronized (mLock) {
+            return mCanceled;
+        }
+    }
+}
diff --git a/core/java/com/android/internal/inputmethod/ICharSequenceResultCallback.aidl b/core/java/com/android/internal/inputmethod/ICharSequenceResultCallback.aidl
new file mode 100644
index 000000000000..da56fd045e57
--- /dev/null
+++ b/core/java/com/android/internal/inputmethod/ICharSequenceResultCallback.aidl
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.internal.inputmethod;
+
+oneway interface ICharSequenceResultCallback {
+    void onResult(in CharSequence result);
+}
diff --git a/core/java/com/android/internal/view/IInputContextCallback.aidl b/core/java/com/android/internal/inputmethod/IExtractedTextResultCallback.aidl
index 0f40a83d7ee4..b603f6adc2d2 100644
--- a/core/java/com/android/internal/view/IInputContextCallback.aidl
+++ b/core/java/com/android/internal/inputmethod/IExtractedTextResultCallback.aidl
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2008 The Android Open Source Project
+ * Copyright (C) 2020 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,19 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.internal.view;
+package com.android.internal.inputmethod;
 
 import android.view.inputmethod.ExtractedText;
 
-/**
- * {@hide}
- */
-oneway interface IInputContextCallback {
-    void setTextBeforeCursor(CharSequence textBeforeCursor, int seq);
-    void setTextAfterCursor(CharSequence textAfterCursor, int seq);
-    void setCursorCapsMode(int capsMode, int seq);
-    void setExtractedText(in ExtractedText extractedText, int seq);
-    void setSelectedText(CharSequence selectedText, int seq);
-    void setRequestUpdateCursorAnchorInfoResult(boolean result, int seq);
-    void setCommitContentResult(boolean result, int seq);
+oneway interface IExtractedTextResultCallback {
+    void onResult(in ExtractedText result);
 }
diff --git a/core/java/com/android/internal/inputmethod/IIntResultCallback.aidl b/core/java/com/android/internal/inputmethod/IIntResultCallback.aidl
new file mode 100644
index 000000000000..bc5ed0d38633
--- /dev/null
+++ b/core/java/com/android/internal/inputmethod/IIntResultCallback.aidl
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.internal.inputmethod;
+
+oneway interface IIntResultCallback {
+    void onResult(int result);
+}
diff --git a/core/java/com/android/internal/inputmethod/ResultCallbacks.java b/core/java/com/android/internal/inputmethod/ResultCallbacks.java
new file mode 100644
index 000000000000..44a8a83b519f
--- /dev/null
+++ b/core/java/com/android/internal/inputmethod/ResultCallbacks.java
@@ -0,0 +1,132 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.internal.inputmethod;
+
+import android.annotation.AnyThread;
+import android.annotation.BinderThread;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+
+import java.lang.ref.WeakReference;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Defines a set of factory methods to create {@link android.os.IBinder}-based callbacks that are
+ * associated with completable objects defined in {@link CancellationGroup.Completable}.
+ */
+public final class ResultCallbacks {
+
+    /**
+     * Not intended to be instantiated.
+     */
+    private ResultCallbacks() {
+    }
+
+    @AnyThread
+    @Nullable
+    private static <T> T unwrap(@NonNull AtomicReference<WeakReference<T>> atomicRef) {
+        final WeakReference<T> ref = atomicRef.getAndSet(null);
+        if (ref == null) {
+            // Double-call is guaranteed to be ignored here.
+            return null;
+        }
+        final T value = ref.get();
+        ref.clear();
+        return value;
+    }
+
+    /**
+     * Creates {@link IIntResultCallback.Stub} that is to set
+     * {@link CancellationGroup.Completable.Int} when receiving the result.
+     *
+     * @param value {@link CancellationGroup.Completable.Int} to be set when receiving the result.
+     * @return {@link IIntResultCallback.Stub} that can be passed as a binder IPC parameter.
+     */
+    @AnyThread
+    public static IIntResultCallback.Stub of(@NonNull CancellationGroup.Completable.Int value) {
+        final AtomicReference<WeakReference<CancellationGroup.Completable.Int>>
+                atomicRef = new AtomicReference<>(new WeakReference<>(value));
+
+        return new IIntResultCallback.Stub() {
+            @BinderThread
+            @Override
+            public void onResult(int result) {
+                final CancellationGroup.Completable.Int value = unwrap(atomicRef);
+                if (value == null) {
+                    return;
+                }
+                value.onComplete(result);
+            }
+        };
+    }
+
+    /**
+     * Creates {@link ICharSequenceResultCallback.Stub} that is to set
+     * {@link CancellationGroup.Completable.CharSequence} when receiving the result.
+     *
+     * @param value {@link CancellationGroup.Completable.CharSequence} to be set when receiving the
+     *              result.
+     * @return {@link ICharSequenceResultCallback.Stub} that can be passed as a binder IPC
+     *         parameter.
+     */
+    @AnyThread
+    public static ICharSequenceResultCallback.Stub of(
+            @NonNull CancellationGroup.Completable.CharSequence value) {
+        final AtomicReference<WeakReference<CancellationGroup.Completable.CharSequence>> atomicRef =
+                new AtomicReference<>(new WeakReference<>(value));
+
+        return new ICharSequenceResultCallback.Stub() {
+            @BinderThread
+            @Override
+            public void onResult(CharSequence result) {
+                final CancellationGroup.Completable.CharSequence value = unwrap(atomicRef);
+                if (value == null) {
+                    return;
+                }
+                value.onComplete(result);
+            }
+        };
+    }
+
+    /**
+     * Creates {@link IExtractedTextResultCallback.Stub} that is to set
+     * {@link CancellationGroup.Completable.ExtractedText} when receiving the result.
+     *
+     * @param value {@link CancellationGroup.Completable.ExtractedText} to be set when receiving the
+     *              result.
+     * @return {@link IExtractedTextResultCallback.Stub} that can be passed as a binder IPC
+     *         parameter.
+     */
+    @AnyThread
+    public static IExtractedTextResultCallback.Stub of(
+            @NonNull CancellationGroup.Completable.ExtractedText value) {
+        final AtomicReference<WeakReference<CancellationGroup.Completable.ExtractedText>>
+                atomicRef = new AtomicReference<>(new WeakReference<>(value));
+
+        return new IExtractedTextResultCallback.Stub() {
+            @BinderThread
+            @Override
+            public void onResult(android.view.inputmethod.ExtractedText result) {
+                final CancellationGroup.Completable.ExtractedText value = unwrap(atomicRef);
+                if (value == null) {
+                    return;
+                }
+                value.onComplete(result);
+            }
+        };
+    }
+}
diff --git a/core/java/com/android/internal/view/IInputConnectionWrapper.java b/core/java/com/android/internal/view/IInputConnectionWrapper.java
index 6278d4a35329..9257c6d19148 100644
--- a/core/java/com/android/internal/view/IInputConnectionWrapper.java
+++ b/core/java/com/android/internal/view/IInputConnectionWrapper.java
@@ -29,6 +29,7 @@ import android.util.Log;
 import android.view.KeyEvent;
 import android.view.inputmethod.CompletionInfo;
 import android.view.inputmethod.CorrectionInfo;
+import android.view.inputmethod.ExtractedText;
 import android.view.inputmethod.ExtractedTextRequest;
 import android.view.inputmethod.InputConnection;
 import android.view.inputmethod.InputConnectionInspector;
@@ -36,6 +37,9 @@ import android.view.inputmethod.InputConnectionInspector.MissingMethodFlags;
 import android.view.inputmethod.InputContentInfo;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.inputmethod.ICharSequenceResultCallback;
+import com.android.internal.inputmethod.IExtractedTextResultCallback;
+import com.android.internal.inputmethod.IIntResultCallback;
 import com.android.internal.os.SomeArgs;
 
 public abstract class IInputConnectionWrapper extends IInputContext.Stub {
@@ -111,28 +115,31 @@ public abstract class IInputConnectionWrapper extends IInputContext.Stub {
 
     abstract protected boolean isActive();
 
-    public void getTextAfterCursor(int length, int flags, int seq, IInputContextCallback callback) {
-        dispatchMessage(obtainMessageIISC(DO_GET_TEXT_AFTER_CURSOR, length, flags, seq, callback));
+    public void getTextAfterCursor(int length, int flags, ICharSequenceResultCallback callback) {
+        dispatchMessage(mH.obtainMessage(DO_GET_TEXT_AFTER_CURSOR, length, flags, callback));
     }
-    
-    public void getTextBeforeCursor(int length, int flags, int seq, IInputContextCallback callback) {
-        dispatchMessage(obtainMessageIISC(DO_GET_TEXT_BEFORE_CURSOR, length, flags, seq, callback));
+
+    public void getTextBeforeCursor(int length, int flags, ICharSequenceResultCallback callback) {
+        dispatchMessage(mH.obtainMessage(DO_GET_TEXT_BEFORE_CURSOR, length, flags, callback));
     }
 
-    public void getSelectedText(int flags, int seq, IInputContextCallback callback) {
-        dispatchMessage(obtainMessageISC(DO_GET_SELECTED_TEXT, flags, seq, callback));
+    public void getSelectedText(int flags, ICharSequenceResultCallback callback) {
+        dispatchMessage(mH.obtainMessage(DO_GET_SELECTED_TEXT, flags, 0 /* unused */, callback));
     }
 
-    public void getCursorCapsMode(int reqModes, int seq, IInputContextCallback callback) {
-        dispatchMessage(obtainMessageISC(DO_GET_CURSOR_CAPS_MODE, reqModes, seq, callback));
+    public void getCursorCapsMode(int reqModes, IIntResultCallback callback) {
+        dispatchMessage(
+                mH.obtainMessage(DO_GET_CURSOR_CAPS_MODE, reqModes, 0 /* unused */, callback));
     }
 
-    public void getExtractedText(ExtractedTextRequest request,
-            int flags, int seq, IInputContextCallback callback) {
-        dispatchMessage(obtainMessageIOSC(DO_GET_EXTRACTED_TEXT, flags,
-                request, seq, callback));
+    public void getExtractedText(ExtractedTextRequest request, int flags,
+            IExtractedTextResultCallback callback) {
+        final SomeArgs args = SomeArgs.obtain();
+        args.arg1 = request;
+        args.arg2 = callback;
+        dispatchMessage(mH.obtainMessage(DO_GET_EXTRACTED_TEXT, flags, 0 /* unused */, args));
     }
-    
+
     public void commitText(CharSequence text, int newCursorPosition) {
         dispatchMessage(obtainMessageIO(DO_COMMIT_TEXT, newCursorPosition, text));
     }
@@ -199,10 +206,9 @@ public abstract class IInputConnectionWrapper extends IInputContext.Stub {
         dispatchMessage(obtainMessageOO(DO_PERFORM_PRIVATE_COMMAND, action, data));
     }
 
-    public void requestUpdateCursorAnchorInfo(int cursorUpdateMode, int seq,
-            IInputContextCallback callback) {
-        dispatchMessage(obtainMessageISC(DO_REQUEST_UPDATE_CURSOR_ANCHOR_INFO, cursorUpdateMode,
-                seq, callback));
+    public void requestUpdateCursorAnchorInfo(int cursorUpdateMode, IIntResultCallback callback) {
+        dispatchMessage(mH.obtainMessage(DO_REQUEST_UPDATE_CURSOR_ANCHOR_INFO, cursorUpdateMode,
+                0 /* unused */, callback));
     }
 
     public void closeConnection() {
@@ -210,9 +216,12 @@ public abstract class IInputConnectionWrapper extends IInputContext.Stub {
     }
 
     public void commitContent(InputContentInfo inputContentInfo, int flags, Bundle opts,
-            int seq, IInputContextCallback callback) {
-        dispatchMessage(obtainMessageIOOSC(DO_COMMIT_CONTENT, flags, inputContentInfo, opts, seq,
-                callback));
+            IIntResultCallback callback) {
+        final SomeArgs args = SomeArgs.obtain();
+        args.arg1 = inputContentInfo;
+        args.arg2 = opts;
+        args.arg3 = callback;
+        dispatchMessage(mH.obtainMessage(DO_COMMIT_CONTENT, flags, 0 /* unused */, args));
     }
 
     void dispatchMessage(Message msg) {
@@ -231,100 +240,97 @@ public abstract class IInputConnectionWrapper extends IInputContext.Stub {
     void executeMessage(Message msg) {
         switch (msg.what) {
             case DO_GET_TEXT_AFTER_CURSOR: {
-                SomeArgs args = (SomeArgs)msg.obj;
+                final ICharSequenceResultCallback callback = (ICharSequenceResultCallback) msg.obj;
+                final InputConnection ic = getInputConnection();
+                final CharSequence result;
+                if (ic == null || !isActive()) {
+                    Log.w(TAG, "getTextAfterCursor on inactive InputConnection");
+                    result = null;
+                } else {
+                    result = ic.getTextAfterCursor(msg.arg1, msg.arg2);
+                }
                 try {
-                    final IInputContextCallback callback = (IInputContextCallback) args.arg6;
-                    final int callbackSeq = args.argi6;
-                    InputConnection ic = getInputConnection();
-                    if (ic == null || !isActive()) {
-                        Log.w(TAG, "getTextAfterCursor on inactive InputConnection");
-                        callback.setTextAfterCursor(null, callbackSeq);
-                        return;
-                    }
-                    callback.setTextAfterCursor(ic.getTextAfterCursor(
-                            msg.arg1, msg.arg2), callbackSeq);
+                    callback.onResult(result);
                 } catch (RemoteException e) {
-                    Log.w(TAG, "Got RemoteException calling setTextAfterCursor", e);
-                } finally {
-                    args.recycle();
+                    Log.w(TAG, "Failed to return the result to getTextAfterCursor()."
+                            + " result=" + result, e);
                 }
                 return;
             }
             case DO_GET_TEXT_BEFORE_CURSOR: {
-                SomeArgs args = (SomeArgs)msg.obj;
+                final ICharSequenceResultCallback callback = (ICharSequenceResultCallback) msg.obj;
+                final InputConnection ic = getInputConnection();
+                final CharSequence result;
+                if (ic == null || !isActive()) {
+                    Log.w(TAG, "getTextBeforeCursor on inactive InputConnection");
+                    result = null;
+                } else {
+                    result = ic.getTextBeforeCursor(msg.arg1, msg.arg2);
+                }
                 try {
-                    final IInputContextCallback callback = (IInputContextCallback) args.arg6;
-                    final int callbackSeq = args.argi6;
-                    InputConnection ic = getInputConnection();
-                    if (ic == null || !isActive()) {
-                        Log.w(TAG, "getTextBeforeCursor on inactive InputConnection");
-                        callback.setTextBeforeCursor(null, callbackSeq);
-                        return;
-                    }
-                    callback.setTextBeforeCursor(ic.getTextBeforeCursor(
-                            msg.arg1, msg.arg2), callbackSeq);
+                    callback.onResult(result);
                 } catch (RemoteException e) {
-                    Log.w(TAG, "Got RemoteException calling setTextBeforeCursor", e);
-                } finally {
-                    args.recycle();
+                    Log.w(TAG, "Failed to return the result to getTextBeforeCursor()."
+                            + " result=" + result, e);
                 }
                 return;
             }
             case DO_GET_SELECTED_TEXT: {
-                SomeArgs args = (SomeArgs)msg.obj;
+                final ICharSequenceResultCallback callback = (ICharSequenceResultCallback) msg.obj;
+                final InputConnection ic = getInputConnection();
+                final CharSequence result;
+                if (ic == null || !isActive()) {
+                    Log.w(TAG, "getSelectedText on inactive InputConnection");
+                    result = null;
+                } else {
+                    result = ic.getSelectedText(msg.arg1);
+                }
                 try {
-                    final IInputContextCallback callback = (IInputContextCallback) args.arg6;
-                    final int callbackSeq = args.argi6;
-                    InputConnection ic = getInputConnection();
-                    if (ic == null || !isActive()) {
-                        Log.w(TAG, "getSelectedText on inactive InputConnection");
-                        callback.setSelectedText(null, callbackSeq);
-                        return;
-                    }
-                    callback.setSelectedText(ic.getSelectedText(
-                            msg.arg1), callbackSeq);
+                    callback.onResult(result);
                 } catch (RemoteException e) {
-                    Log.w(TAG, "Got RemoteException calling setSelectedText", e);
-                } finally {
-                    args.recycle();
+                    Log.w(TAG, "Failed to return the result to getSelectedText()."
+                            + " result=" + result, e);
                 }
                 return;
             }
             case DO_GET_CURSOR_CAPS_MODE: {
-                SomeArgs args = (SomeArgs)msg.obj;
+                final IIntResultCallback callback = (IIntResultCallback) msg.obj;
+                final InputConnection ic = getInputConnection();
+                final int result;
+                if (ic == null || !isActive()) {
+                    Log.w(TAG, "getCursorCapsMode on inactive InputConnection");
+                    result = 0;
+                } else {
+                    result = ic.getCursorCapsMode(msg.arg1);
+                }
                 try {
-                    final IInputContextCallback callback = (IInputContextCallback) args.arg6;
-                    final int callbackSeq = args.argi6;
-                    InputConnection ic = getInputConnection();
-                    if (ic == null || !isActive()) {
-                        Log.w(TAG, "getCursorCapsMode on inactive InputConnection");
-                        callback.setCursorCapsMode(0, callbackSeq);
-                        return;
-                    }
-                    callback.setCursorCapsMode(ic.getCursorCapsMode(msg.arg1),
-                            callbackSeq);
+                    callback.onResult(result);
                 } catch (RemoteException e) {
-                    Log.w(TAG, "Got RemoteException calling setCursorCapsMode", e);
-                } finally {
-                    args.recycle();
+                    Log.w(TAG, "Failed to return the result to getCursorCapsMode()."
+                            + " result=" + result, e);
                 }
                 return;
             }
             case DO_GET_EXTRACTED_TEXT: {
-                SomeArgs args = (SomeArgs)msg.obj;
+                final SomeArgs args = (SomeArgs) msg.obj;
                 try {
-                    final IInputContextCallback callback = (IInputContextCallback) args.arg6;
-                    final int callbackSeq = args.argi6;
-                    InputConnection ic = getInputConnection();
+                    final ExtractedTextRequest request = (ExtractedTextRequest) args.arg1;
+                    final IExtractedTextResultCallback callback =
+                            (IExtractedTextResultCallback) args.arg2;
+                    final InputConnection ic = getInputConnection();
+                    final ExtractedText result;
                     if (ic == null || !isActive()) {
                         Log.w(TAG, "getExtractedText on inactive InputConnection");
-                        callback.setExtractedText(null, callbackSeq);
-                        return;
+                        result = null;
+                    } else {
+                        result = ic.getExtractedText(request, msg.arg1);
+                    }
+                    try {
+                        callback.onResult(result);
+                    } catch (RemoteException e) {
+                        Log.w(TAG, "Failed to return the result to getExtractedText()."
+                                + " result=" + result, e);
                     }
-                    callback.setExtractedText(ic.getExtractedText(
-                            (ExtractedTextRequest)args.arg1, msg.arg1), callbackSeq);
-                } catch (RemoteException e) {
-                    Log.w(TAG, "Got RemoteException calling setExtractedText", e);
                 } finally {
                     args.recycle();
                 }
@@ -494,22 +500,20 @@ public abstract class IInputConnectionWrapper extends IInputContext.Stub {
                 return;
             }
             case DO_REQUEST_UPDATE_CURSOR_ANCHOR_INFO: {
-                SomeArgs args = (SomeArgs)msg.obj;
+                final IIntResultCallback callback = (IIntResultCallback) msg.obj;
+                final InputConnection ic = getInputConnection();
+                final boolean result;
+                if (ic == null || !isActive()) {
+                    Log.w(TAG, "requestCursorAnchorInfo on inactive InputConnection");
+                    result = false;
+                } else {
+                    result = ic.requestCursorUpdates(msg.arg1);
+                }
                 try {
-                    final IInputContextCallback callback = (IInputContextCallback) args.arg6;
-                    final int callbackSeq = args.argi6;
-                    InputConnection ic = getInputConnection();
-                    if (ic == null || !isActive()) {
-                        Log.w(TAG, "requestCursorAnchorInfo on inactive InputConnection");
-                        callback.setRequestUpdateCursorAnchorInfoResult(false, callbackSeq);
-                        return;
-                    }
-                    callback.setRequestUpdateCursorAnchorInfoResult(
-                            ic.requestCursorUpdates(msg.arg1), callbackSeq);
+                    callback.onResult(result ? 1 : 0);
                 } catch (RemoteException e) {
-                    Log.w(TAG, "Got RemoteException calling requestCursorAnchorInfo", e);
-                } finally {
-                    args.recycle();
+                    Log.w(TAG, "Failed to return the result to requestCursorUpdates()."
+                            + " result=" + result, e);
                 }
                 return;
             }
@@ -547,26 +551,28 @@ public abstract class IInputConnectionWrapper extends IInputContext.Stub {
                 final int flags = msg.arg1;
                 SomeArgs args = (SomeArgs) msg.obj;
                 try {
-                    final IInputContextCallback callback = (IInputContextCallback) args.arg6;
-                    final int callbackSeq = args.argi6;
-                    InputConnection ic = getInputConnection();
+                    final IIntResultCallback callback = (IIntResultCallback) args.arg3;
+                    final InputConnection ic = getInputConnection();
+                    final boolean result;
                     if (ic == null || !isActive()) {
                         Log.w(TAG, "commitContent on inactive InputConnection");
-                        callback.setCommitContentResult(false, callbackSeq);
-                        return;
+                        result = false;
+                    } else {
+                        final InputContentInfo inputContentInfo = (InputContentInfo) args.arg1;
+                        if (inputContentInfo == null || !inputContentInfo.validate()) {
+                            Log.w(TAG, "commitContent with invalid inputContentInfo="
+                                    + inputContentInfo);
+                            result = false;
+                        } else {
+                            result = ic.commitContent(inputContentInfo, flags, (Bundle) args.arg2);
+                        }
                     }
-                    final InputContentInfo inputContentInfo = (InputContentInfo) args.arg1;
-                    if (inputContentInfo == null || !inputContentInfo.validate()) {
-                        Log.w(TAG, "commitContent with invalid inputContentInfo="
-                                + inputContentInfo);
-                        callback.setCommitContentResult(false, callbackSeq);
-                        return;
+                    try {
+                        callback.onResult(result ? 1 : 0);
+                    } catch (RemoteException e) {
+                        Log.w(TAG, "Failed to return the result to commitContent()."
+                                + " result=" + result, e);
                     }
-                    final boolean result =
-                            ic.commitContent(inputContentInfo, flags, (Bundle) args.arg2);
-                    callback.setCommitContentResult(result, callbackSeq);
-                } catch (RemoteException e) {
-                    Log.w(TAG, "Got RemoteException calling commitContent", e);
                 } finally {
                     args.recycle();
                 }
@@ -588,40 +594,6 @@ public abstract class IInputConnectionWrapper extends IInputContext.Stub {
         return mH.obtainMessage(what, 0, 0, arg1);
     }
 
-    Message obtainMessageISC(int what, int arg1, int callbackSeq, IInputContextCallback callback) {
-        final SomeArgs args = SomeArgs.obtain();
-        args.arg6 = callback;
-        args.argi6 = callbackSeq;
-        return mH.obtainMessage(what, arg1, 0, args);
-    }
-
-    Message obtainMessageIISC(int what, int arg1, int arg2, int callbackSeq,
-            IInputContextCallback callback) {
-        final SomeArgs args = SomeArgs.obtain();
-        args.arg6 = callback;
-        args.argi6 = callbackSeq;
-        return mH.obtainMessage(what, arg1, arg2, args);
-    }
-
-    Message obtainMessageIOOSC(int what, int arg1, Object objArg1, Object objArg2, int callbackSeq,
-            IInputContextCallback callback) {
-        final SomeArgs args = SomeArgs.obtain();
-        args.arg1 = objArg1;
-        args.arg2 = objArg2;
-        args.arg6 = callback;
-        args.argi6 = callbackSeq;
-        return mH.obtainMessage(what, arg1, 0, args);
-    }
-
-    Message obtainMessageIOSC(int what, int arg1, Object arg2, int callbackSeq,
-            IInputContextCallback callback) {
-        final SomeArgs args = SomeArgs.obtain();
-        args.arg1 = arg2;
-        args.arg6 = callback;
-        args.argi6 = callbackSeq;
-        return mH.obtainMessage(what, arg1, 0, args);
-    }
-
     Message obtainMessageIO(int what, int arg1, Object arg2) {
         return mH.obtainMessage(what, arg1, 0, arg2);
     }
diff --git a/core/java/com/android/internal/view/IInputContext.aidl b/core/java/com/android/internal/view/IInputContext.aidl
index c22799179b72..86f1293c014f 100644
--- a/core/java/com/android/internal/view/IInputContext.aidl
+++ b/core/java/com/android/internal/view/IInputContext.aidl
@@ -23,7 +23,9 @@ import android.view.inputmethod.CorrectionInfo;
 import android.view.inputmethod.ExtractedTextRequest;
 import android.view.inputmethod.InputContentInfo;
 
-import com.android.internal.view.IInputContextCallback;
+import com.android.internal.inputmethod.ICharSequenceResultCallback;
+import com.android.internal.inputmethod.IExtractedTextResultCallback;
+import com.android.internal.inputmethod.IIntResultCallback;
 
 /**
  * Interface from an input method to the application, allowing it to perform
@@ -31,14 +33,14 @@ import com.android.internal.view.IInputContextCallback;
  * {@hide}
  */
  oneway interface IInputContext {
-    void getTextBeforeCursor(int length, int flags, int seq, IInputContextCallback callback); 
+    void getTextBeforeCursor(int length, int flags, ICharSequenceResultCallback callback);
 
-    void getTextAfterCursor(int length, int flags, int seq, IInputContextCallback callback);
-    
-    void getCursorCapsMode(int reqModes, int seq, IInputContextCallback callback);
-    
-    void getExtractedText(in ExtractedTextRequest request, int flags, int seq,
-            IInputContextCallback callback);
+    void getTextAfterCursor(int length, int flags, ICharSequenceResultCallback callback);
+
+    void getCursorCapsMode(int reqModes, IIntResultCallback callback);
+
+    void getExtractedText(in ExtractedTextRequest request, int flags,
+            IExtractedTextResultCallback callback);
 
     void deleteSurroundingText(int beforeLength, int afterLength);
     void deleteSurroundingTextInCodePoints(int beforeLength, int afterLength);
@@ -71,11 +73,10 @@ import com.android.internal.view.IInputContextCallback;
 
     void setComposingRegion(int start, int end);
 
-    void getSelectedText(int flags, int seq, IInputContextCallback callback);
+    void getSelectedText(int flags, ICharSequenceResultCallback callback);
 
-    void requestUpdateCursorAnchorInfo(int cursorUpdateMode, int seq,
-            IInputContextCallback callback);
+    void requestUpdateCursorAnchorInfo(int cursorUpdateMode, IIntResultCallback callback);
 
-    void commitContent(in InputContentInfo inputContentInfo, int flags, in Bundle opts, int sec,
-            IInputContextCallback callback);
+    void commitContent(in InputContentInfo inputContentInfo, int flags, in Bundle opts,
+            IIntResultCallback callback);
 }
diff --git a/core/java/com/android/internal/view/InputConnectionWrapper.java b/core/java/com/android/internal/view/InputConnectionWrapper.java
index a41048c0f426..0bf52345bc7e 100644
--- a/core/java/com/android/internal/view/InputConnectionWrapper.java
+++ b/core/java/com/android/internal/view/InputConnectionWrapper.java
@@ -17,14 +17,12 @@
 package com.android.internal.view;
 
 import android.annotation.AnyThread;
-import android.annotation.BinderThread;
 import android.annotation.NonNull;
-import android.compat.annotation.UnsupportedAppUsage;
+import android.annotation.Nullable;
 import android.inputmethodservice.AbstractInputMethodService;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.RemoteException;
-import android.os.SystemClock;
 import android.util.Log;
 import android.view.KeyEvent;
 import android.view.inputmethod.CompletionInfo;
@@ -36,10 +34,15 @@ import android.view.inputmethod.InputConnectionInspector;
 import android.view.inputmethod.InputConnectionInspector.MissingMethodFlags;
 import android.view.inputmethod.InputContentInfo;
 
+import com.android.internal.inputmethod.CancellationGroup;
+import com.android.internal.inputmethod.ResultCallbacks;
+
 import java.lang.ref.WeakReference;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.TimeUnit;
 
 public class InputConnectionWrapper implements InputConnection {
+    private static final String TAG = "InputConnectionWrapper";
+
     private static final int MAX_WAIT_TIME_MILLIS = 2000;
     private final IInputContext mIInputContext;
     @NonNull
@@ -49,257 +52,94 @@ public class InputConnectionWrapper implements InputConnection {
     private final int mMissingMethods;
 
     /**
-     * {@code true} if the system already decided to take away IME focus from the target app. This
-     * can be signaled even when the corresponding signal is in the task queue and
-     * {@link InputMethodService#onUnbindInput()} is not yet called back on the UI thread.
+     * Signaled when the system decided to take away IME focus from the target app.
+     *
+     * <p>This is expected to be signaled immediately when the IME process receives
+     * {@link IInputMethod#unbindInput()}.</p>
      */
     @NonNull
-    private final AtomicBoolean mIsUnbindIssued;
-
-    static class InputContextCallback extends IInputContextCallback.Stub {
-        private static final String TAG = "InputConnectionWrapper.ICC";
-        public int mSeq;
-        public boolean mHaveValue;
-        public CharSequence mTextBeforeCursor;
-        public CharSequence mTextAfterCursor;
-        public CharSequence mSelectedText;
-        public ExtractedText mExtractedText;
-        public int mCursorCapsMode;
-        public boolean mRequestUpdateCursorAnchorInfoResult;
-        public boolean mCommitContentResult;
-
-        // A 'pool' of one InputContextCallback.  Each ICW request will attempt to gain
-        // exclusive access to this object.
-        private static InputContextCallback sInstance = new InputContextCallback();
-        private static int sSequenceNumber = 1;
-        
-        /**
-         * Returns an InputContextCallback object that is guaranteed not to be in use by
-         * any other thread.  The returned object's 'have value' flag is cleared and its expected
-         * sequence number is set to a new integer.  We use a sequence number so that replies that
-         * occur after a timeout has expired are not interpreted as replies to a later request.
-         */
-        @UnsupportedAppUsage
-        @AnyThread
-        private static InputContextCallback getInstance() {
-            synchronized (InputContextCallback.class) {
-                // Return sInstance if it's non-null, otherwise construct a new callback
-                InputContextCallback callback;
-                if (sInstance != null) {
-                    callback = sInstance;
-                    sInstance = null;
-                    
-                    // Reset the callback
-                    callback.mHaveValue = false;
-                } else {
-                    callback = new InputContextCallback();
-                }
-                
-                // Set the sequence number
-                callback.mSeq = sSequenceNumber++;
-                return callback;
-            }
-        }
-        
-        /**
-         * Makes the given InputContextCallback available for use in the future.
-         */
-        @UnsupportedAppUsage
-        @AnyThread
-        private void dispose() {
-            synchronized (InputContextCallback.class) {
-                // If sInstance is non-null, just let this object be garbage-collected
-                if (sInstance == null) {
-                    // Allow any objects being held to be gc'ed
-                    mTextAfterCursor = null;
-                    mTextBeforeCursor = null;
-                    mExtractedText = null;
-                    sInstance = this;
-                }
-            }
-        }
-
-        @BinderThread
-        public void setTextBeforeCursor(CharSequence textBeforeCursor, int seq) {
-            synchronized (this) {
-                if (seq == mSeq) {
-                    mTextBeforeCursor = textBeforeCursor;
-                    mHaveValue = true;
-                    notifyAll();
-                } else {
-                    Log.i(TAG, "Got out-of-sequence callback " + seq + " (expected " + mSeq
-                            + ") in setTextBeforeCursor, ignoring.");
-                }
-            }
-        }
-
-        @BinderThread
-        public void setTextAfterCursor(CharSequence textAfterCursor, int seq) {
-            synchronized (this) {
-                if (seq == mSeq) {
-                    mTextAfterCursor = textAfterCursor;
-                    mHaveValue = true;
-                    notifyAll();
-                } else {
-                    Log.i(TAG, "Got out-of-sequence callback " + seq + " (expected " + mSeq
-                            + ") in setTextAfterCursor, ignoring.");
-                }
-            }
-        }
-
-        @BinderThread
-        public void setSelectedText(CharSequence selectedText, int seq) {
-            synchronized (this) {
-                if (seq == mSeq) {
-                    mSelectedText = selectedText;
-                    mHaveValue = true;
-                    notifyAll();
-                } else {
-                    Log.i(TAG, "Got out-of-sequence callback " + seq + " (expected " + mSeq
-                            + ") in setSelectedText, ignoring.");
-                }
-            }
-        }
-
-        @BinderThread
-        public void setCursorCapsMode(int capsMode, int seq) {
-            synchronized (this) {
-                if (seq == mSeq) {
-                    mCursorCapsMode = capsMode; 
-                    mHaveValue = true;  
-                    notifyAll();
-                } else {
-                    Log.i(TAG, "Got out-of-sequence callback " + seq + " (expected " + mSeq
-                            + ") in setCursorCapsMode, ignoring.");
-                }
-            }
-        }
-
-        @BinderThread
-        public void setExtractedText(ExtractedText extractedText, int seq) {
-            synchronized (this) {
-                if (seq == mSeq) {
-                    mExtractedText = extractedText;
-                    mHaveValue = true;
-                    notifyAll();
-                } else {
-                    Log.i(TAG, "Got out-of-sequence callback " + seq + " (expected " + mSeq
-                            + ") in setExtractedText, ignoring.");
-                }
-            }
-        }
-
-        @BinderThread
-        public void setRequestUpdateCursorAnchorInfoResult(boolean result, int seq) {
-            synchronized (this) {
-                if (seq == mSeq) {
-                    mRequestUpdateCursorAnchorInfoResult = result;
-                    mHaveValue = true;
-                    notifyAll();
-                } else {
-                    Log.i(TAG, "Got out-of-sequence callback " + seq + " (expected " + mSeq
-                            + ") in setCursorAnchorInfoRequestResult, ignoring.");
-                }
-            }
-        }
-
-        @BinderThread
-        public void setCommitContentResult(boolean result, int seq) {
-            synchronized (this) {
-                if (seq == mSeq) {
-                    mCommitContentResult = result;
-                    mHaveValue = true;
-                    notifyAll();
-                } else {
-                    Log.i(TAG, "Got out-of-sequence callback " + seq + " (expected " + mSeq
-                            + ") in setCommitContentResult, ignoring.");
-                }
-            }
-        }
-
-        /**
-         * Waits for a result for up to {@link #MAX_WAIT_TIME_MILLIS} milliseconds.
-         * 
-         * <p>The caller must be synchronized on this callback object.
-         */
-        @AnyThread
-        void waitForResultLocked() {
-            long startTime = SystemClock.uptimeMillis();
-            long endTime = startTime + MAX_WAIT_TIME_MILLIS;
-
-            while (!mHaveValue) {
-                long remainingTime = endTime - SystemClock.uptimeMillis();
-                if (remainingTime <= 0) {
-                    Log.w(TAG, "Timed out waiting on IInputContextCallback");
-                    return;
-                }
-                try {
-                    wait(remainingTime);
-                } catch (InterruptedException e) {
-                }
-            }
-        }
-    }
+    private final CancellationGroup mCancellationGroup;
 
     public InputConnectionWrapper(
             @NonNull WeakReference<AbstractInputMethodService> inputMethodService,
-            IInputContext inputContext, @MissingMethodFlags final int missingMethods,
-            @NonNull AtomicBoolean isUnbindIssued) {
+            IInputContext inputContext, @MissingMethodFlags int missingMethods,
+            @NonNull CancellationGroup cancellationGroup) {
         mInputMethodService = inputMethodService;
         mIInputContext = inputContext;
         mMissingMethods = missingMethods;
-        mIsUnbindIssued = isUnbindIssued;
+        mCancellationGroup = cancellationGroup;
+    }
+
+    @AnyThread
+    private static void logInternal(@Nullable String methodName, boolean timedOut,
+            @Nullable Object defaultValue) {
+        if (timedOut) {
+            Log.w(TAG, methodName + " didn't respond in " + MAX_WAIT_TIME_MILLIS + " msec."
+                    + " Returning default: " + defaultValue);
+        } else {
+            Log.w(TAG, methodName + " was canceled before complete. Returning default: "
+                    + defaultValue);
+        }
+    }
+
+    @AnyThread
+    private static int getResultOrZero(@NonNull CancellationGroup.Completable.Int value,
+             @NonNull String methodName) {
+        final boolean timedOut = value.await(MAX_WAIT_TIME_MILLIS,  TimeUnit.MILLISECONDS);
+        if (value.hasValue()) {
+            return value.getValue();
+        }
+        logInternal(methodName, timedOut, 0);
+        return 0;
+    }
+
+    @AnyThread
+    @Nullable
+    private static <T> T getResultOrNull(@NonNull CancellationGroup.Completable.Values<T> value,
+            @NonNull String methodName) {
+        final boolean timedOut = value.await(MAX_WAIT_TIME_MILLIS,  TimeUnit.MILLISECONDS);
+        if (value.hasValue()) {
+            return value.getValue();
+        }
+        logInternal(methodName, timedOut, null);
+        return null;
     }
 
     @AnyThread
     public CharSequence getTextAfterCursor(int length, int flags) {
-        if (mIsUnbindIssued.get()) {
+        if (mCancellationGroup.isCanceled()) {
             return null;
         }
 
-        CharSequence value = null;
+        final CancellationGroup.Completable.CharSequence value =
+                mCancellationGroup.createCompletableCharSequence();
         try {
-            InputContextCallback callback = InputContextCallback.getInstance();
-            mIInputContext.getTextAfterCursor(length, flags, callback.mSeq, callback);
-            synchronized (callback) {
-                callback.waitForResultLocked();
-                if (callback.mHaveValue) {
-                    value = callback.mTextAfterCursor;
-                }
-            }
-            callback.dispose();
+            mIInputContext.getTextAfterCursor(length, flags, ResultCallbacks.of(value));
         } catch (RemoteException e) {
             return null;
         }
-        return value;
+        return getResultOrNull(value, "getTextAfterCursor()");
     }
 
     @AnyThread
     public CharSequence getTextBeforeCursor(int length, int flags) {
-        if (mIsUnbindIssued.get()) {
+        if (mCancellationGroup.isCanceled()) {
             return null;
         }
 
-        CharSequence value = null;
+        final CancellationGroup.Completable.CharSequence value =
+                mCancellationGroup.createCompletableCharSequence();
         try {
-            InputContextCallback callback = InputContextCallback.getInstance();
-            mIInputContext.getTextBeforeCursor(length, flags, callback.mSeq, callback);
-            synchronized (callback) {
-                callback.waitForResultLocked();
-                if (callback.mHaveValue) {
-                    value = callback.mTextBeforeCursor;
-                }
-            }
-            callback.dispose();
+            mIInputContext.getTextBeforeCursor(length, flags, ResultCallbacks.of(value));
         } catch (RemoteException e) {
             return null;
         }
-        return value;
+        return getResultOrNull(value, "getTextBeforeCursor()");
     }
 
     @AnyThread
     public CharSequence getSelectedText(int flags) {
-        if (mIsUnbindIssued.get()) {
+        if (mCancellationGroup.isCanceled()) {
             return null;
         }
 
@@ -307,67 +147,46 @@ public class InputConnectionWrapper implements InputConnection {
             // This method is not implemented.
             return null;
         }
-        CharSequence value = null;
+        final CancellationGroup.Completable.CharSequence value =
+                mCancellationGroup.createCompletableCharSequence();
         try {
-            InputContextCallback callback = InputContextCallback.getInstance();
-            mIInputContext.getSelectedText(flags, callback.mSeq, callback);
-            synchronized (callback) {
-                callback.waitForResultLocked();
-                if (callback.mHaveValue) {
-                    value = callback.mSelectedText;
-                }
-            }
-            callback.dispose();
+            mIInputContext.getSelectedText(flags, ResultCallbacks.of(value));
         } catch (RemoteException e) {
             return null;
         }
-        return value;
+        return getResultOrNull(value, "getSelectedText()");
     }
 
     @AnyThread
     public int getCursorCapsMode(int reqModes) {
-        if (mIsUnbindIssued.get()) {
+        if (mCancellationGroup.isCanceled()) {
             return 0;
         }
 
-        int value = 0;
+        final CancellationGroup.Completable.Int value =
+                mCancellationGroup.createCompletableInt();
         try {
-            InputContextCallback callback = InputContextCallback.getInstance();
-            mIInputContext.getCursorCapsMode(reqModes, callback.mSeq, callback);
-            synchronized (callback) {
-                callback.waitForResultLocked();
-                if (callback.mHaveValue) {
-                    value = callback.mCursorCapsMode;
-                }
-            }
-            callback.dispose();
+            mIInputContext.getCursorCapsMode(reqModes, ResultCallbacks.of(value));
         } catch (RemoteException e) {
             return 0;
         }
-        return value;
+        return getResultOrZero(value, "getCursorCapsMode()");
     }
 
     @AnyThread
     public ExtractedText getExtractedText(ExtractedTextRequest request, int flags) {
-        if (mIsUnbindIssued.get()) {
+        if (mCancellationGroup.isCanceled()) {
             return null;
         }
 
-        ExtractedText value = null;
+        final CancellationGroup.Completable.ExtractedText value =
+                mCancellationGroup.createCompletableExtractedText();
         try {
-            InputContextCallback callback = InputContextCallback.getInstance();
-            mIInputContext.getExtractedText(request, flags, callback.mSeq, callback);
-            synchronized (callback) {
-                callback.waitForResultLocked();
-                if (callback.mHaveValue) {
-                    value = callback.mExtractedText;
-                }
-            }
-            callback.dispose();
+            mIInputContext.getExtractedText(request, flags, ResultCallbacks.of(value));
         } catch (RemoteException e) {
             return null;
         }
-        return value;
+        return getResultOrNull(value, "getExtractedText()");
     }
 
     @AnyThread
@@ -563,29 +382,22 @@ public class InputConnectionWrapper implements InputConnection {
 
     @AnyThread
     public boolean requestCursorUpdates(int cursorUpdateMode) {
-        if (mIsUnbindIssued.get()) {
+        if (mCancellationGroup.isCanceled()) {
             return false;
         }
 
-        boolean result = false;
         if (isMethodMissing(MissingMethodFlags.REQUEST_CURSOR_UPDATES)) {
             // This method is not implemented.
             return false;
         }
+        final CancellationGroup.Completable.Int value = mCancellationGroup.createCompletableInt();
         try {
-            InputContextCallback callback = InputContextCallback.getInstance();
-            mIInputContext.requestUpdateCursorAnchorInfo(cursorUpdateMode, callback.mSeq, callback);
-            synchronized (callback) {
-                callback.waitForResultLocked();
-                if (callback.mHaveValue) {
-                    result = callback.mRequestUpdateCursorAnchorInfoResult;
-                }
-            }
-            callback.dispose();
+            mIInputContext.requestUpdateCursorAnchorInfo(cursorUpdateMode,
+                    ResultCallbacks.of(value));
         } catch (RemoteException e) {
             return false;
         }
-        return result;
+        return getResultOrZero(value, "requestUpdateCursorAnchorInfo()") != 0;
     }
 
     @AnyThread
@@ -601,38 +413,31 @@ public class InputConnectionWrapper implements InputConnection {
 
     @AnyThread
     public boolean commitContent(InputContentInfo inputContentInfo, int flags, Bundle opts) {
-        if (mIsUnbindIssued.get()) {
+        if (mCancellationGroup.isCanceled()) {
             return false;
         }
 
-        boolean result = false;
         if (isMethodMissing(MissingMethodFlags.COMMIT_CONTENT)) {
             // This method is not implemented.
             return false;
         }
-        try {
-            if ((flags & InputConnection.INPUT_CONTENT_GRANT_READ_URI_PERMISSION) != 0) {
-                final AbstractInputMethodService inputMethodService = mInputMethodService.get();
-                if (inputMethodService == null) {
-                    // This basically should not happen, because it's the the caller of this method.
-                    return false;
-                }
-                inputMethodService.exposeContent(inputContentInfo, this);
-            }
 
-            InputContextCallback callback = InputContextCallback.getInstance();
-            mIInputContext.commitContent(inputContentInfo, flags, opts, callback.mSeq, callback);
-            synchronized (callback) {
-                callback.waitForResultLocked();
-                if (callback.mHaveValue) {
-                    result = callback.mCommitContentResult;
-                }
+        if ((flags & InputConnection.INPUT_CONTENT_GRANT_READ_URI_PERMISSION) != 0) {
+            final AbstractInputMethodService inputMethodService = mInputMethodService.get();
+            if (inputMethodService == null) {
+                // This basically should not happen, because it's the caller of this method.
+                return false;
             }
-            callback.dispose();
+            inputMethodService.exposeContent(inputContentInfo, this);
+        }
+
+        final CancellationGroup.Completable.Int value = mCancellationGroup.createCompletableInt();
+        try {
+            mIInputContext.commitContent(inputContentInfo, flags, opts, ResultCallbacks.of(value));
         } catch (RemoteException e) {
             return false;
         }
-        return result;
+        return getResultOrZero(value, "commitContent()") != 0;
     }
 
     @AnyThread