Fixing a SpellChecker leak in Android M

AndroidLeaks.kt

/**
 * Every editable TextView has an Editor instance which has a SpellChecker instance. SpellChecker
 * is in charge of displaying the little squiggle spans that show typos. SpellChecker starts a
 * SpellCheckerSession as needed and then closes it when the TextView is detached from the window.
 * A SpellCheckerSession is in charge of communicating with the spell checker service (which lives
 * in another process) through TextServicesManager.
 *
 * The SpellChecker sends the TextView content to the spell checker service every 400ms, ie every
 * time the service calls back with a result the SpellChecker schedules another check for 400ms
 * later.
 *
 * When the TextView is detached from the window, the spell checker closes the session. In practice,
 * SpellCheckerSessionListenerImpl.mHandler is set to null and when the service calls
 * SpellCheckerSessionListenerImpl.onGetSuggestions or
 * SpellCheckerSessionListenerImpl.onGetSentenceSuggestions back from another process, there's a
 * null check for SpellCheckerSessionListenerImpl.mHandler and the callback is dropped.
 *
 * Unfortunately, on Android M there's a race condition in how that's done. When the service calls
 * back into our app process, the IPC call is received on a binder thread. That's when the null
 * check happens. If the session is not closed at this point (mHandler not null), the callback is
 * then posted to the main thread. If on the main thread the session is closed after that post but
 * prior to that post being handled, then the post will still be processed, after the session has
 * been closed.
 *
 * When the post is processed, SpellCheckerSession calls back into SpellChecker which in turns
 * schedules a new spell check to be ran in 400ms. The check is an anonymous inner class
 * (SpellChecker$1) stored as SpellChecker.mSpellRunnable and implementing Runnable. It is scheduled
 * by calling [View.postDelayed]. As we've seen, at this point the session may be closed which means
 * that the view has been detached. [View.postDelayed] behaves differently when a view is detached:
 * instead of posting to the single [Handler] used by the view hierarchy, it enqueues the Runnable
 * into ViewRootImpl.RunQueue, a static queue that holds on to "actions" to be executed. As soon as
 * a view hierarchy is attached, the ViewRootImpl.RunQueue is processed and emptied.
 *
 * Unfortunately, that means that as long as no view hierarchy is attached, ie as long as there
 * are no activities alive, the actions stay in ViewRootImpl.RunQueue. That means SpellChecker$1
 * ends up being kept in memory. It holds on to SpellChecker which in turns holds on
 * to the detached TextView and corresponding destroyed activity & view hierarchy.
 *
 * We have a fix for this! When the spell check session is closed, we replace
 * SpellCheckerSession.mSpellCheckerSessionListener (which normally is the SpellChecker) with a
 * no-op implementation. So even if callbacks are enqueued to the main thread handler, these
 * callbacks will call the no-op implementation and SpellChecker will not be scheduling a spell
 * check.
 *
 * Sources to corroborate:
 *
 * https://android.googlesource.com/platform/frameworks/base/+/marshmallow-release
 * /core/java/android/view/textservice/SpellCheckerSession.java
 * /core/java/android/view/textservice/TextServicesManager.java
 * /core/java/android/widget/SpellChecker.java
 * /core/java/android/view/ViewRootImpl.java
 */
private fun fixSpellCheckerLeak() {
  if (VERSION.SDK_INT != 23) {
    return
  }

  try {
    val textServiceClass = TextServicesManager::class.java
    val getInstanceMethod = textServiceClass.getDeclaredMethod("getInstance")

    val sServiceField = textServiceClass.getDeclaredField("sService")
    sServiceField.isAccessible = true

    val serviceStubInterface =
      Class.forName("com.android.internal.textservice.ITextServicesManager")

    val spellCheckSessionClass = Class.forName("android.view.textservice.SpellCheckerSession")
    val mSpellCheckerSessionListenerField =
      spellCheckSessionClass.getDeclaredField("mSpellCheckerSessionListener")
    mSpellCheckerSessionListenerField.isAccessible = true

    val spellCheckerSessionListenerImplClass =
      Class.forName("android.view.textservice.SpellCheckerSession\$SpellCheckerSessionListenerImpl")
    val listenerImplHandlerField = spellCheckerSessionListenerImplClass.getDeclaredField("mHandler")
    listenerImplHandlerField.isAccessible = true

    val spellCheckSessionHandlerClass =
      Class.forName("android.view.textservice.SpellCheckerSession\$1")
    val outerInstanceField = spellCheckSessionHandlerClass.getDeclaredField("this$0")
    outerInstanceField.isAccessible = true

    val listenerInterface =
      Class.forName("android.view.textservice.SpellCheckerSession\$SpellCheckerSessionListener")
    val noOpListener = Proxy.newProxyInstance(
        listenerInterface.classLoader, arrayOf(listenerInterface)
    ) { _: Any, _: Method, _: kotlin.Array<Any>? ->
      Timber.d("Received call to no-op SpellCheckerSessionListener after session closed")
    }

    // Ensure a TextServicesManager instance is created and TextServicesManager.sService set.
    getInstanceMethod
        .invoke(null)
    val realService = sServiceField[null]!!

    val spellCheckerListenerToSession = mutableMapOf<Any, Any>()

    val proxyService = Proxy.newProxyInstance(
        serviceStubInterface.classLoader, arrayOf(serviceStubInterface)
    ) { _: Any, method: Method, args: kotlin.Array<Any>? ->
      try {
        if (method.name == "getSpellCheckerService") {
          // getSpellCheckerService is called when the session is opened, which allows us to
          // capture the corresponding SpellCheckerSession instance via
          // SpellCheckerSessionListenerImpl.mHandler.this$0
          val spellCheckerSessionListener = args!![3]
          val handler = listenerImplHandlerField[spellCheckerSessionListener]!!
          val spellCheckerSession = outerInstanceField[handler]!!
          // We add to a map of SpellCheckerSessionListenerImpl to SpellCheckerSession
          spellCheckerListenerToSession[spellCheckerSessionListener] = spellCheckerSession
        } else if (method.name == "finishSpellCheckerService") {
          // finishSpellCheckerService is called when the session is open. After the session has been
          // closed, any pending work posted to SpellCheckerSession.mHandler should be ignored. We do
          // so by replacing mSpellCheckerSessionListener with a no-op implementation.
          val spellCheckerSessionListener = args!![0]
          val spellCheckerSession =
            spellCheckerListenerToSession.remove(spellCheckerSessionListener)!!
          // We use the SpellCheckerSessionListenerImpl to find the corresponding SpellCheckerSession
          // At this point in time the session was just closed to
          // SpellCheckerSessionListenerImpl.mHandler is null, which is why we had to capture
          // the SpellCheckerSession during the getSpellCheckerService call.
          mSpellCheckerSessionListenerField[spellCheckerSession] = noOpListener
        }
      } catch (ignored: Exception) {
        Timber.d(ignored, "Unable to fix SpellChecker leak")
      }
      // Standard delegation
      try {
        return@newProxyInstance if (args != null) {
          method.invoke(realService, *args)
        } else {
          method.invoke(realService)
        }
      } catch (invocationException: InvocationTargetException) {
        throw invocationException.targetException
      }
    }
    sServiceField[null] = proxyService
  } catch (ignored: Exception) {
    Timber.d(ignored, "Unable to fix SpellChecker leak")
  }
}

leaktrace.txt

┬───
│ GC Root: System class
│
├─ flow.Preconditions class
│    Leaking: NO (Thread↓ is not leaking and a class is never leaking)
│    ↓ static Preconditions.MAIN_THREAD
├─ java.lang.Thread instance
│    Leaking: NO (the main thread always runs)
│    Thread name: 'main'
│    ↓ Thread.localValues
│             ~~~~~~~~~~~
├─ java.lang.ThreadLocal$Values instance
│    Leaking: UNKNOWN
│    ↓ ThreadLocal$Values.table
│                         ~~~~~
├─ java.lang.Object[] array
│    Leaking: UNKNOWN
│    ↓ Object[].[104]
│               ~~~~~
├─ android.view.ViewRootImpl$RunQueue instance
│    Leaking: UNKNOWN
│    ↓ ViewRootImpl$RunQueue.mActions
│                            ~~~~~~~~
├─ java.util.ArrayList instance
│    Leaking: UNKNOWN
│    ↓ ArrayList.array
│                ~~~~~
├─ java.lang.Object[] array
│    Leaking: UNKNOWN
│    ↓ Object[].[0]
│               ~~~
├─ android.view.ViewRootImpl$RunQueue$HandlerAction instance
│    Leaking: UNKNOWN
│    ↓ ViewRootImpl$RunQueue$HandlerAction.action
│                                          ~~~~~~
├─ android.widget.SpellChecker$1 instance
│    Leaking: UNKNOWN
│    Anonymous class implementing java.lang.Runnable
│    ↓ SpellChecker$1.this$0
│                     ~~~~~~
├─ android.widget.SpellChecker instance
│    Leaking: UNKNOWN
│    ↓ SpellChecker.mTextView
│                   ~~~~~~~~~
├─ com.squareup.widgets.OnScreenRectangleEditText instance
│    Leaking: YES (View.mContext references a destroyed activity)
│    mContext instance of flow.path.PathContext, wrapping activity com.squareup.ui.main.MainActivity with mDestroyed = true
│    View#mParent is set
│    View#mAttachInfo is null (view detached)
│    View.mWindowAttachCount = 1
│    mListeners[0] = instance of com.squareup.ui.BaseXableEditText$ShowButtonTextWatcher
│    mListeners[1] = instance of anonymous subclass of com.squareup.text.EmptyTextWatcher
│    ↓ OnScreenRectangleEditText.mContext
├─ flow.path.PathContext instance
│    Leaking: YES (OnScreenRectangleEditText↑ is leaking and PathContext wraps an Activity with Activity.mDestroyed true)
│    ↓ PathContext.inflater
├─ com.android.internal.policy.PhoneLayoutInflater instance
│    Leaking: YES (PathContext↑ is leaking)
│    ↓ PhoneLayoutInflater.mPrivateFactory
╰→ com.squareup.ui.main.MainActivity instance
​     Leaking: YES (ObjectWatcher was watching this because com.squareup.ui.main.MainActivity received Activity#onDestroy() callback and Activity#mDestroyed is true)
​     key = 20fe26cf-048e-4601-9391-404b57599f38
​     watchDurationMillis = 5846
​     retainedDurationMillis = 845