ensure to keep a singule compose lifecycle owner and attach new paren…

…t lifecycles if new ones are provided
square · Nov 20, 2024 · 953e2ab · 953e2ab
1 parent bf97796
commit 953e2ab
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt b/workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeLifecycleOwner.kt
@@ -18,8 +18,11 @@ import androidx.lifecycle.LifecycleRegistry
 @Composable internal fun rememberChildLifecycleOwner(
   parentLifecycle: Lifecycle = LocalLifecycleOwner.current.lifecycle
 ): LifecycleOwner {
-  val lifecycleOwner = remember {
-    ComposeLifecycleOwner.installOn(parentLifecycle)
+  val owner = remember { ComposeLifecycleOwner.installOn(parentLifecycle) }
+  val lifecycleOwner = remember(parentLifecycle) {
+    owner.apply {
+      updateParentLifecycle(parentLifecycle)
+    }
   }
   return lifecycleOwner
 }
@@ -56,9 +59,11 @@ import androidx.lifecycle.LifecycleRegistry
  * @param parentLifecycle The parent [Lifecycle] with which this lifecycle owner should synchronize.
  */
 private class ComposeLifecycleOwner(
-  private val parentLifecycle: Lifecycle
+  initialParentLifecycle: Lifecycle
 ) : LifecycleOwner, RememberObserver, LifecycleEventObserver {
 
+  private var parentLifecycle: Lifecycle = initialParentLifecycle
+
   private val registry = LifecycleRegistry(this)
   override val lifecycle: Lifecycle
     get() = registry
@@ -79,6 +84,12 @@ private class ComposeLifecycleOwner(
     }
   }
 
+  fun updateParentLifecycle(lifecycle: Lifecycle) {
+    parentLifecycle.removeObserver(this)
+    parentLifecycle = lifecycle
+    parentLifecycle.addObserver(this)
+  }
+
   override fun onStateChanged(
     source: LifecycleOwner,
     event: Event