View.backPressHandler memory leak fix.

workflow-ui/core-android/src/androidTest/AndroidManifest.xml

@@ -7,6 +7,6 @@
     <activity
         android:name="com.squareup.workflow1.ui.container.fixtures.BackStackContainerLifecycleActivity"
         android:theme="@style/Theme.AppCompat.NoActionBar"/>
-<activity android:name="androidx.activity.ComponentActivity"/>
+    <activity android:name="androidx.activity.ComponentActivity"/>
   </application>
 </manifest>

workflow-ui/core-android/src/androidTest/java/com/squareup/workflow1/ui/BackPressedHandlerTest.kt

@@ -0,0 +1,117 @@
+package com.squareup.workflow1.ui
+
+import android.view.View
+import android.view.ViewGroup
+import androidx.activity.ComponentActivity
+import androidx.activity.OnBackPressedCallback
+import androidx.activity.OnBackPressedDispatcherSpy
+import androidx.lifecycle.Lifecycle.State.DESTROYED
+import androidx.lifecycle.Lifecycle.State.RESUMED
+import androidx.lifecycle.LifecycleRegistry
+import androidx.lifecycle.ViewTreeLifecycleOwner
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import com.google.common.truth.Truth.assertThat
+import com.squareup.workflow1.ui.internal.test.DetectLeaksAfterTestSuccess
+import com.squareup.workflow1.ui.internal.test.IdlingDispatcherRule
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.RuleChain
+
+@OptIn(WorkflowUiExperimentalApi::class)
+internal class BackPressedHandlerTest {
+  private val scenarioRule = ActivityScenarioRule(ComponentActivity::class.java)
+  private val scenario get() = scenarioRule.scenario
+
+  @get:Rule val rules: RuleChain = RuleChain.outerRule(DetectLeaksAfterTestSuccess())
+    .around(scenarioRule)
+    .around(IdlingDispatcherRule)
+
+  private var viewHandlerCount = 0
+  private val viewBackHandler: BackPressHandler = {
+    viewHandlerCount++
+  }
+
+  @Test fun itWorksWhenHandlerIsAddedBeforeAttach() {
+    scenario.onActivity { activity ->
+      val view = View(activity)
+      view.backPressedHandler = viewBackHandler
+
+      activity.setContentView(view)
+      assertThat(viewHandlerCount).isEqualTo(0)
+
+      activity.onBackPressed()
+      assertThat(viewHandlerCount).isEqualTo(1)
+    }
+  }
+
+  @Test fun itWorksWhenHandlerIsAddedAfterAttach() {
+    scenario.onActivity { activity ->
+      val view = View(activity)
+      activity.setContentView(view)
+
+      view.backPressedHandler = viewBackHandler
+      assertThat(viewHandlerCount).isEqualTo(0)
+
+      activity.onBackPressed()
+      assertThat(viewHandlerCount).isEqualTo(1)
+    }
+  }
+
+  @Test fun onlyActiveWhileViewIsAttached() {
+    var fallbackCallCount = 0
+    val defaultBackHandler = object : OnBackPressedCallback(true) {
+      override fun handleOnBackPressed() {
+        fallbackCallCount++
+      }
+    }
+
+    scenario.onActivity { activity ->
+      activity.onBackPressedDispatcher.addCallback(defaultBackHandler)
+
+      val view = View(activity)
+      view.backPressedHandler = viewBackHandler
+
+      activity.onBackPressed()
+      assertThat(fallbackCallCount).isEqualTo(1)
+      assertThat(viewHandlerCount).isEqualTo(0)
+
+      activity.setContentView(view)
+      activity.onBackPressed()
+      assertThat(fallbackCallCount).isEqualTo(1)
+      assertThat(viewHandlerCount).isEqualTo(1)
+
+      (view.parent as ViewGroup).removeView(view)
+      activity.onBackPressed()
+      assertThat(fallbackCallCount).isEqualTo(2)
+      assertThat(viewHandlerCount).isEqualTo(1)
+
+      activity.setContentView(view)
+      activity.onBackPressed()
+      assertThat(fallbackCallCount).isEqualTo(2)
+      assertThat(viewHandlerCount).isEqualTo(2)
+    }
+  }
+
+  @Test fun callbackIsRemoved() {
+    scenario.onActivity { activity ->
+      val spy = OnBackPressedDispatcherSpy(activity.onBackPressedDispatcher)
+      assertThat(spy.callbacks()).isEmpty()
+
+      val lifecycle = LifecycleRegistry(activity)
+      lifecycle.currentState = RESUMED
+
+      val view = View(activity)
+      view.backPressedHandler = viewBackHandler
+      assertThat(spy.callbacks()).hasSize(1)
+
+      ViewTreeLifecycleOwner.set(view) { lifecycle }
+      activity.setContentView(view)
+
+      (view.parent as ViewGroup).removeView(view)
+      assertThat(spy.callbacks()).hasSize(1)
+
+      lifecycle.currentState = DESTROYED
+      assertThat(spy.callbacks()).isEmpty()
+    }
+  }
+}

workflow-ui/core-android/src/androidTest/java/com/squareup/workflow1/ui/OnBackPressedDispatcherSpy.java

@@ -0,0 +1,15 @@
+package androidx.activity;
+
+import java.util.ArrayDeque;
+
+public class OnBackPressedDispatcherSpy {
+  private final OnBackPressedDispatcher dispatcher;
+
+  public OnBackPressedDispatcherSpy(OnBackPressedDispatcher dispatcher) {
+    this.dispatcher = dispatcher;
+  }
+
+  public ArrayDeque<OnBackPressedCallback> callbacks() {
+    return dispatcher.mOnBackPressedCallbacks;
+  }
+}

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/BackPressHandler.kt

@@ -7,6 +7,7 @@ import android.view.View.OnAttachStateChangeListener
 import androidx.activity.OnBackPressedCallback
 import androidx.activity.OnBackPressedDispatcherOwner
 import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.ViewTreeLifecycleOwner
 
@@ -21,79 +22,123 @@ public typealias BackPressHandler = () -> Unit
  * A function to be called if the device back button is pressed while this
  * view is attached to a window.
  *
- * Implemented via [OnBackPressedDispatcher][androidx.activity.OnBackPressedDispatcher],
- * making this a last-registered-first-served mechanism.
+ * Implemented via [OnBackPressedDispatcher][androidx.activity.OnBackPressedDispatcher].
+ * That means that this is a last-registered-first-served mechanism, and that it is
+ * compatible with Compose back button handling.
  */
 @WorkflowUiExperimentalApi
 public var View.backPressedHandler: BackPressHandler?
-  get() = handlerWrapperOrNull?.handler
+  get() = observerOrNull?.handler
   set(value) {
-    handlerWrapperOrNull?.stop()
+    observerOrNull?.stop()
 
-    val wrapper = value?.let {
-      HandleBackPressWhenAttached(this, it).apply { start() }
+    observerOrNull = value?.let {
+      AttachStateAndLifecycleObserver(this, it).apply { start() }
     }
-    setTag(R.id.view_back_handler, wrapper)
   }
 
 @WorkflowUiExperimentalApi
-private val View.handlerWrapperOrNull
-  get() = getTag(R.id.view_back_handler) as HandleBackPressWhenAttached?
+private var View.observerOrNull: AttachStateAndLifecycleObserver?
+  get() = getTag(R.id.view_back_handler) as AttachStateAndLifecycleObserver?
+  set(value) {
+    setTag(R.id.view_back_handler, value)
+  }
 
 /**
- * Uses the [androidx.activity.OnBackPressedDispatcher] to associate a [BackPressHandler]
- * with a [View].
+ * This is more complicated than one would hope because [Lifecycle] and memory leaks.
+ *
+ * - We need to claim our spot in the
+ *   [OnBackPressedDispatcher][androidx.activity.OnBackPressedDispatcher] immediately,
+ *   to ensure our [onBackPressedCallback] shadows upstream ones, and can be shadowed
+ *   appropriately itself
+ * - The whole point of this mechanism is to be active only while the view is active
+ * - That's what [ViewTreeLifecycleOwner] is for, but we can't really find that until
+ *   we're attached -- which often does not happen in the order required for registering
+ *   with the dispatcher
+ *
+ * So, our [start] is called immediately, to get [onBackPressedCallback] listed at the right
+ * spot in the dispatcher's stack. But the [onBackPressedCallback]'s enabled / disabled state
+ * is tied to whether the [view] is attached or not.
  *
- * - Registers [handler] on [start]
- * - Enables and disables it as [view] is attached to or detached from its window
- * - De-registers it on [stop], or when its [lifecycle][ViewTreeLifecycleOwner] is destroyed
+ * Also note that we expect to find a [ViewTreeLifecycleOwner] at attach time,
+ * so that we can know when it's time to remove the [onBackPressedCallback] from
+ * the dispatch stack
+ * ([no memory leaks please](https://github.com/square/workflow-kotlin/issues/889)).
+ *
+ * Why is it okay to wait for the [ViewTreeLifecycleOwner] to be destroyed before we
+ * remove [onBackPressedCallback] from the dispatcher? In normal apps that's
+ * the `Activity` or a `Fragment`, which will live a very long time, but Workflow UI
+ * is more controlling than that. `WorkflowViewStub` and the rest of the stock container
+ * classes use `WorkflowLifecycleOwner` to provide a short lived [ViewTreeLifecycleOwner]
+ * for each [View] they create, and tear it down before moving to the next one.
+ *
+ * None the less, as a belt-and-suspenders guard against leaking,
+ * we also take care to null out the pointer from the [onBackPressedCallback] to the
+ * actual [handler] while the [view] is detached. We can't be confident that the
+ * [ViewTreeLifecycleOwner] we find will be a well behaved one that was put in place
+ * by `WorkflowLifecycleOwner`. Who knows what adventures our clients will get up to.
  */
 @WorkflowUiExperimentalApi
-private class HandleBackPressWhenAttached(
+private class AttachStateAndLifecycleObserver(
   private val view: View,
   val handler: BackPressHandler
 ) : OnAttachStateChangeListener, DefaultLifecycleObserver {
-  private val onBackPressedCallback = object : OnBackPressedCallback(false) {
-    override fun handleOnBackPressed() {
-      handler.invoke()
-    }
-  }
+  private val onBackPressedCallback = NullableOnBackPressedCallback()
+  private var lifecycleOrNull: Lifecycle? = null
 
   fun start() {
     view.context.onBackPressedDispatcherOwnerOrNull()
       ?.let { owner ->
         owner.onBackPressedDispatcher.addCallback(owner, onBackPressedCallback)
         view.addOnAttachStateChangeListener(this)
         if (view.isAttachedToWindow) onViewAttachedToWindow(view)
-
-        // We enable the handler only while its view is attached to a window.
-        // This ensures that a temporarily removed view (e.g. for caching)
-        // does not participate in back button handling.
-        ViewTreeLifecycleOwner.get(view)?.lifecycle?.addObserver(this)
       }
   }
 
   fun stop() {
     onBackPressedCallback.remove()
     view.removeOnAttachStateChangeListener(this)
-    ViewTreeLifecycleOwner.get(view)?.lifecycle?.removeObserver(this)
+    lifecycleOrNull?.removeObserver(this)
   }
 
   override fun onViewAttachedToWindow(attachedView: View) {
     require(view === attachedView)
-    onBackPressedCallback.isEnabled = true
+    lifecycleOrNull?.let { lifecycle ->
+      lifecycle.removeObserver(this)
+      lifecycleOrNull = null
+    }
+    ViewTreeLifecycleOwner.get(view)?.lifecycle?.let { lifecycle ->
+      lifecycleOrNull = lifecycle
+      onBackPressedCallback.handlerOrNull = handler
+      onBackPressedCallback.isEnabled = true
+      lifecycle.addObserver(this)
+    }
+      ?: error(
+        "Expected to find a ViewTreeLifecycleOwner to manage the " +
+          "backPressedHandler ($handler) for $view"
+      )
   }
 
   override fun onViewDetachedFromWindow(detachedView: View) {
     require(view === detachedView)
     onBackPressedCallback.isEnabled = false
+    onBackPressedCallback.handlerOrNull = null
   }
 
   override fun onDestroy(owner: LifecycleOwner) {
     stop()
   }
 }
 
+@WorkflowUiExperimentalApi
+internal class NullableOnBackPressedCallback : OnBackPressedCallback(false) {
+  var handlerOrNull: BackPressHandler? = null
+
+  override fun handleOnBackPressed() {
+    handlerOrNull?.invoke()
+  }
+}
+
 @WorkflowUiExperimentalApi
 public tailrec fun Context.onBackPressedDispatcherOwnerOrNull(): OnBackPressedDispatcherOwner? =
   when (this) {

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/androidx/WorkflowLifecycleOwner.kt

@@ -112,9 +112,6 @@ internal class RealWorkflowLifecycleOwner(
   OnAttachStateChangeListener,
   LifecycleEventObserver {
 
-  /**
-   * Weak reference ensures that we don't leak the view.
-   */
   private var view: View? = null
 
   private val localLifecycle =