Refactor to ComposeLifecycleOwner for Better Lifecycle Synchronization

workflow-ui/compose/src/androidTest/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwnerTest.kt

@@ -0,0 +1,122 @@
+package com.squareup.workflow1.ui.compose
+
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.platform.LocalLifecycleOwner
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.Lifecycle.State.CREATED
+import androidx.lifecycle.Lifecycle.State.RESUMED
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.LifecycleRegistry
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+
+class ComposeLifecycleOwnerTest {
+
+  @get:Rule
+  val composeTestRule = createComposeRule()
+
+  private var mParentLifecycle: LifecycleRegistry? = null
+
+  @Test
+  fun childLifecycleOwner_initialStateIsResumedWhenParentIsResumed() {
+    val parentLifecycle = ensureParentLifecycle()
+
+    lateinit var childLifecycleOwner: LifecycleOwner
+    composeTestRule.setContent {
+      parentLifecycle.currentState = RESUMED
+      childLifecycleOwner = rememberChildLifecycleOwner(parentLifecycle)
+      CompositionLocalProvider(LocalLifecycleOwner provides childLifecycleOwner) {
+        // let's assert right away as things are composing, because we want to ensure that
+        // the lifecycle is in the correct state as soon as possible & not just after composition
+        // has finished
+        assertThat(childLifecycleOwner.lifecycle.currentState).isEqualTo(Lifecycle.State.RESUMED)
+      }
+    }
+
+    // Allow the composition to complete
+    composeTestRule.waitForIdle()
+
+    // Outside the composition, assert the lifecycle state again
+    assertThat(childLifecycleOwner.lifecycle.currentState)
+      .isEqualTo(Lifecycle.State.RESUMED)
+  }
+
+  @Test
+  fun childLifecycleOwner_initialStateIsResumedAfterParentResumed() {
+    val parentLifecycle = ensureParentLifecycle()
+
+    lateinit var childLifecycleOwner: LifecycleOwner
+    composeTestRule.setContent {
+      childLifecycleOwner = rememberChildLifecycleOwner(parentLifecycle)
+      parentLifecycle.currentState = CREATED
+      CompositionLocalProvider(LocalLifecycleOwner provides childLifecycleOwner) {
+        // let's assert right away as things are composing, because we want to ensure that
+        // the lifecycle is in the correct state as soon as possible & not just after composition
+        // has finished
+        assertThat(childLifecycleOwner.lifecycle.currentState).isEqualTo(Lifecycle.State.CREATED)
+      }
+    }
+
+    // Allow the composition to complete
+    composeTestRule.waitForIdle()
+
+    // Outside the composition, assert the lifecycle state again
+    assertThat(childLifecycleOwner.lifecycle.currentState)
+      .isEqualTo(Lifecycle.State.CREATED)
+  }
+
+  @Test
+  fun childLifecycleOwner_initialStateRemainsSameAfterParentLifecycleChange() {
+    lateinit var updatedChildLifecycleOwner: LifecycleOwner
+    lateinit var tempChildLifecycleOwner: LifecycleOwner
+
+    val customParentLifecycleOwner: LifecycleOwner = object : LifecycleOwner {
+      private val registry = LifecycleRegistry(this)
+      override val lifecycle: Lifecycle
+        get() = registry
+    }
+
+    composeTestRule.setContent {
+      var seenRecomposition by remember { mutableStateOf(false) }
+      // after initial composition, change the parent lifecycle owner
+      LaunchedEffect(Unit) { seenRecomposition = true }
+      CompositionLocalProvider(
+        if (seenRecomposition) {
+          LocalLifecycleOwner provides customParentLifecycleOwner
+        } else {
+          LocalLifecycleOwner provides LocalLifecycleOwner.current
+        }
+      ) {
+
+        updatedChildLifecycleOwner = rememberChildLifecycleOwner()
+        // let's save the original reference to lifecycle owner on first pass
+        if (!seenRecomposition) {
+          tempChildLifecycleOwner = updatedChildLifecycleOwner
+        }
+      }
+    }
+
+    // Allow the composition to complete
+    composeTestRule.waitForIdle()
+    // assert that the [ComposeLifecycleOwner] is the same instance when the parent lifecycle owner
+    // is changed.
+    assertThat(updatedChildLifecycleOwner).isEqualTo(tempChildLifecycleOwner)
+  }
+
+  private fun ensureParentLifecycle(): LifecycleRegistry {
+    if (mParentLifecycle == null) {
+      val owner = object : LifecycleOwner {
+        override val lifecycle = LifecycleRegistry.createUnsafe(this)
+      }
+      mParentLifecycle = owner.lifecycle
+    }
+    return mParentLifecycle!!
+  }
+}

workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt

@@ -0,0 +1,118 @@
+package com.squareup.workflow1.ui.compose
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.RememberObserver
+import androidx.compose.runtime.remember
+import androidx.compose.ui.platform.LocalLifecycleOwner
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.Lifecycle.Event
+import androidx.lifecycle.LifecycleEventObserver
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.LifecycleRegistry
+
+/**
+ * Returns a [LifecycleOwner] that is a mirror of the current [LocalLifecycleOwner] until this
+ * function leaves the composition. Similar to [WorkflowLifecycleOwner] for views, but a
+ * bit simpler since we don't need to worry about attachment state.
+ */
+@Composable internal fun rememberChildLifecycleOwner(
+  parentLifecycle: Lifecycle = LocalLifecycleOwner.current.lifecycle
+): LifecycleOwner {
+  val owner = remember {
+    ComposeLifecycleOwner.installOn(
+      initialParentLifecycle = parentLifecycle
+    )
+  }
+  val lifecycleOwner = remember(parentLifecycle) {
+    owner.apply { updateParentLifecycle(parentLifecycle) }
+  }
+  return lifecycleOwner
+}
+
+/**
+ * A custom [LifecycleOwner] that synchronizes its lifecycle with a parent [Lifecycle] and
+ * integrates with Jetpack Compose's lifecycle through [RememberObserver].
+ *
+ * ## Purpose
+ *
+ * - Ensures that any lifecycle-aware components within a composable function have a lifecycle that
+ *   accurately reflects both the parent lifecycle and the composable's own lifecycle.
+ * - Manages lifecycle transitions and observer registration/removal to prevent memory leaks and
+ *   ensure proper cleanup when the composable leaves the composition.
+ *
+ * ## Key Features
+ *
+ * - Lifecycle Synchronization: Mirrors lifecycle events from the provided `parentLifecycle` to
+ *   its own [LifecycleRegistry], ensuring consistent state transitions.
+ * - Compose Integration: Implements [RememberObserver] to align with the composable's lifecycle
+ *   in the Compose memory model.
+ * - Automatic Observer Management: Adds and removes a [LifecycleEventObserver] to the parent
+ *   lifecycle, preventing leaks and ensuring proper disposal.
+ * - **State Transition Safety:** Carefully manages lifecycle state changes to avoid illegal
+ *   transitions, especially during destruction.
+ *
+ * ## Usage Notes
+ *
+ * - Should be used in conjunction with `remember` and provided the `parentLifecycle` as a key to
+ *   ensure it updates correctly when the parent lifecycle changes.
+ * - By integrating with Compose's lifecycle, it ensures that resources are properly released when
+ *   the composable leaves the composition.
+ *
+ * @param initialParentLifecycle The parent [Lifecycle] with which this lifecycle owner should
+ * synchronize with initially. If new parent lifecycles are provided, they should be passed to
+ * [updateParentLifecycle].
+ */
+private class ComposeLifecycleOwner(
+  initialParentLifecycle: Lifecycle
+) : LifecycleOwner, RememberObserver, LifecycleEventObserver {
+
+  private var parentLifecycle: Lifecycle = initialParentLifecycle
+
+  private val registry = LifecycleRegistry(this)
+  override val lifecycle: Lifecycle
+    get() = registry
+
+  override fun onRemembered() {
+  }
+
+  override fun onAbandoned() {
+    onForgotten()
+  }
+
+  override fun onForgotten() {
+    parentLifecycle.removeObserver(this)
+
+    // If we're leaving the composition, ensure the lifecycle is cleaned up
+    if (registry.currentState != Lifecycle.State.INITIALIZED) {
+      registry.currentState = Lifecycle.State.DESTROYED
+    }
+  }
+
+  fun updateParentLifecycle(lifecycle: Lifecycle) {
+    parentLifecycle.removeObserver(this)
+    parentLifecycle = lifecycle
+    parentLifecycle.addObserver(this)
+  }
+
+  override fun onStateChanged(
+    source: LifecycleOwner,
+    event: Event
+  ) {
+    registry.handleLifecycleEvent(event)
+  }
+
+  companion object {
+    fun installOn(initialParentLifecycle: Lifecycle): ComposeLifecycleOwner {
+      return ComposeLifecycleOwner(initialParentLifecycle).also {
+        // We need to synchronize the lifecycles before the child ever even sees the lifecycle
+        // because composes contract tries to guarantee that the lifecycle is in at least the
+        // CREATED state by the time composition is actually running. If we don't synchronize
+        // the lifecycles right away, then we break that invariant. One concrete case of this is
+        // that SavedStateRegistry requires its lifecycle to be CREATED before reading values
+        // from it, and consuming values from an SSR is a valid thing to do from composition
+        // directly, and in fact AndroidComposeView itself does this.
+        initialParentLifecycle.addObserver(it)
+      }
+    }
+  }
+}

workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/WorkflowRendering.kt

@@ -3,24 +3,16 @@ package com.squareup.workflow1.ui.compose
 import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
-import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.key
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.LocalLifecycleOwner
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.Lifecycle.State.DESTROYED
-import androidx.lifecycle.Lifecycle.State.INITIALIZED
-import androidx.lifecycle.LifecycleEventObserver
-import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.LifecycleRegistry
 import com.squareup.workflow1.ui.Compatible
 import com.squareup.workflow1.ui.Screen
 import com.squareup.workflow1.ui.ScreenViewHolder
 import com.squareup.workflow1.ui.ViewEnvironment
 import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
 import com.squareup.workflow1.ui.WorkflowViewStub
-import com.squareup.workflow1.ui.androidx.WorkflowLifecycleOwner
 
 /**
  * Renders [rendering] into the composition using the [ViewEnvironment] found in
@@ -94,41 +86,3 @@ public fun WorkflowRendering(
     }
   }
 }
-
-/**
- * Returns a [LifecycleOwner] that is a mirror of the current [LocalLifecycleOwner] until this
- * function leaves the composition. Similar to [WorkflowLifecycleOwner] for views, but a
- * bit simpler since we don't need to worry about attachment state.
- */
-@Composable private fun rememberChildLifecycleOwner(): LifecycleOwner {
-  val lifecycleOwner = remember {
-    object : LifecycleOwner {
-      val registry = LifecycleRegistry(this)
-      override val lifecycle: Lifecycle
-        get() = registry
-    }
-  }
-  val parentLifecycle = LocalLifecycleOwner.current.lifecycle
-
-  DisposableEffect(parentLifecycle) {
-    val parentObserver = LifecycleEventObserver { _, event ->
-      // Any time the parent lifecycle changes state, perform the same change on our lifecycle.
-      lifecycleOwner.registry.handleLifecycleEvent(event)
-    }
-
-    parentLifecycle.addObserver(parentObserver)
-    onDispose {
-      parentLifecycle.removeObserver(parentObserver)
-
-      // If we're leaving the composition it means the WorkflowRendering is either going away itself
-      // or about to switch to an incompatible rendering – either way, this lifecycle is dead. Note
-      // that we can't transition from INITIALIZED to DESTROYED – the LifecycleRegistry will throw.
-      // WorkflowLifecycleOwner has this same check.
-      if (lifecycleOwner.registry.currentState != INITIALIZED) {
-        lifecycleOwner.registry.currentState = DESTROYED
-      }
-    }
-  }
-
-  return lifecycleOwner
-}