Adds LifecycleOwnerWrapper for use-live-data

Summary:
This is the rework of D54364594. One of reasons we need to add a wrap is because we added default lifecycle owner to BaseMountingView(D56757176) and would like to use it for useLiveData, but useLiveData may be called before the view is attached. Adding a default wrapper could solve this issue.

In this case, we need to set LifecycleOwnerWrapper as a TreeProp to all ComponentTree rather than only root ComponentTree

Reviewed By: adityasharat

Differential Revision: D57863409

fbshipit-source-id: 1ba68f8508dca325bbc0ca7b3885f864d2740ec8

litho-core/src/main/java/com/facebook/litho/ComponentTree.java

@@ -78,6 +78,7 @@
 import com.facebook.litho.debug.DebugOverlay;
 import com.facebook.litho.debug.LithoDebugEvent;
 import com.facebook.litho.debug.LithoDebugEventAttributes;
+import com.facebook.litho.lifecycle.LifecycleOwnerWrapper;
 import com.facebook.litho.perfboost.LithoPerfBooster;
 import com.facebook.litho.stats.LithoStats;
 import com.facebook.rendercore.LogLevel;
@@ -129,6 +130,7 @@ public class ComponentTree
       "ComponentTree:CTContextIsDifferentFromRootBuilderContext";
   public static final int STATE_UPDATES_IN_LOOP_THRESHOLD = 50;
   private static boolean sBoostPerfLayoutStateFuture = false;
+
   @Nullable LithoVisibilityEventsController mLithoVisibilityEventsController;
 
   @GuardedBy("this")
@@ -236,7 +238,10 @@ public void run() {
   private int mCommittedLayoutVersion = INVALID_LAYOUT_VERSION;
 
   @GuardedBy("this")
-  private @Nullable TreePropContainer mRootTreePropContainer;
+  private @Nullable TreePropContainer mTreePropContainer;
+
+  @GuardedBy("this")
+  private @Nullable TreePropContainer mDefaultTreeProps;
 
   @GuardedBy("this")
   private int mWidthSpec = SIZE_UNINITIALIZED;
@@ -375,6 +380,15 @@ protected ComponentTree(Builder builder) {
                 : getLithoVisibilityEventsController(),
             null,
             builder.parentTreePropContainer);
+    if (ComponentsConfiguration.defaultInstance.enableLifecycleOwnerWrapper) {
+      initializeDefaultTreeProps();
+
+      TreePropContainer treePropContainer =
+          TreePropContainer.acquire(mContext.getTreePropContainer());
+      treePropContainer.putAll(mDefaultTreeProps);
+      mContext.setTreePropContainer(treePropContainer);
+      mContext.setParentTreePropContainerCloned(true);
+    }
 
     PreAllocationHandler preAllocationHandler = builder.config.preAllocationHandler;
     if (preAllocationHandler != null) {
@@ -1316,14 +1330,14 @@ private void dispatchStateUpdateEnqueuedEvent(String attribution, boolean isSync
   }
 
   void updateStateInternal(boolean isAsync, String attribution, boolean isCreateLayoutInProgress) {
-    final @Nullable TreePropContainer rootTreePropContainer;
+    final @Nullable TreePropContainer treePropContainer;
 
     synchronized (this) {
       if (mRoot == null) {
         return;
       }
 
-      rootTreePropContainer = TreePropContainer.copy(mRootTreePropContainer);
+      treePropContainer = TreePropContainer.copy(mTreePropContainer);
 
       if (isCreateLayoutInProgress) {
         logStateUpdatesFromCreateLayout(attribution);
@@ -1343,7 +1357,7 @@ void updateStateInternal(boolean isAsync, String attribution, boolean isCreateLa
         isAsync ? RenderSource.UPDATE_STATE_ASYNC : RenderSource.UPDATE_STATE_SYNC,
         INVALID_LAYOUT_VERSION,
         attribution,
-        rootTreePropContainer,
+        treePropContainer,
         isCreateLayoutInProgress,
         false);
   }
@@ -1377,9 +1391,11 @@ private void logStateUpdatesFromCreateLayout(@Nullable String attribution) {
    * <p>It will make sure that the tree properties are properly cloned and stored.
    */
   private void setInternalTreeProp(TreeProp<?> key, @Nullable Object value) {
-    if (!mContext.isParentTreePropContainerCloned()) {
-      mContext.setTreePropContainer(TreePropContainer.acquire(mContext.getTreePropContainer()));
-      mContext.setParentTreePropContainerCloned(true);
+    if (!ComponentsConfiguration.defaultInstance.enableLifecycleOwnerWrapper) {
+      if (!mContext.isParentTreePropContainerCloned()) {
+        mContext.setTreePropContainer(TreePropContainer.acquire(mContext.getTreePropContainer()));
+        mContext.setParentTreePropContainerCloned(true);
+      }
     }
 
     TreePropContainer treePropContainer = mContext.getTreePropContainer();
@@ -1840,15 +1856,21 @@ private void setRootAndSizeSpecInternal(
       }
 
       if (treePropContainerInitialized) {
-        mRootTreePropContainer = treePropContainer;
-      } else {
-        treePropContainer = mRootTreePropContainer;
+        if (ComponentsConfiguration.defaultInstance.enableLifecycleOwnerWrapper) {
+          TreePropContainer newTreeProps = TreePropContainer.acquire(treePropContainer);
+          newTreeProps.putAll(mDefaultTreeProps);
+          if (!newTreeProps.equals(mTreePropContainer)) {
+            mTreePropContainer = newTreeProps;
+          }
+        } else {
+          mTreePropContainer = treePropContainer;
+        }
       }
 
       requestedWidthSpec = mWidthSpec;
       requestedHeightSpec = mHeightSpec;
       requestedRoot = mRoot;
-      requestedTreePropContainer = mRootTreePropContainer;
+      requestedTreePropContainer = mTreePropContainer;
 
       mLastLayoutSource = source;
     }
@@ -1925,9 +1947,12 @@ private void requestRenderWithSplitFutures(
     // there is no need to calculate the resolved result again and we can proceed straight to
     // layout.
     if (currentResolveResult != null) {
-      boolean canLayoutWithoutResolve =
-          (currentResolveResult.component == root
-              && currentResolveResult.context.getTreePropContainer() == treePropContainer);
+      boolean isSameTreeProps =
+          currentResolveResult.context.getTreePropContainer() == treePropContainer
+              || (ComponentsConfiguration.defaultInstance.enableLifecycleOwnerWrapper
+                  && treePropContainer == null);
+
+      boolean canLayoutWithoutResolve = (currentResolveResult.component == root && isSameTreeProps);
       if (canLayoutWithoutResolve) {
         requestLayoutWithSplitFutures(
             currentResolveResult,
@@ -2602,6 +2627,11 @@ private static synchronized Looper getDefaultResolveThreadLooper() {
     return sDefaultResolveThreadLooper;
   }
 
+  private void initializeDefaultTreeProps() {
+    mDefaultTreeProps = new TreePropContainer();
+    mDefaultTreeProps.put(LifecycleOwnerTreeProp, new LifecycleOwnerWrapper(null));
+  }
+
   private boolean isCompatibleSpec(
       final @Nullable LayoutState layoutState, final int widthSpec, final int heightSpec) {
     return layoutState != null
@@ -2747,7 +2777,14 @@ public synchronized void subscribeToVisibilityEventsController(
   }
 
   synchronized void setLifecycleOwnerTreeProp(LifecycleOwner owner) {
-    setInternalTreeProp(LifecycleOwnerTreeProp, owner);
+    if (ComponentsConfiguration.defaultInstance.enableLifecycleOwnerWrapper) {
+      TreePropContainer treePropContainer = mContext.getTreePropContainer();
+      if (treePropContainer != null) {
+        ((LifecycleOwnerWrapper) treePropContainer.get(LifecycleOwnerTreeProp)).setDelegate(owner);
+      }
+    } else {
+      setInternalTreeProp(LifecycleOwnerTreeProp, owner);
+    }
   }
 
   public synchronized boolean isSubscribedToVisibilityEventsController() {

litho-core/src/main/java/com/facebook/litho/TreePropContainer.kt

@@ -41,6 +41,13 @@ class TreePropContainer {
     map[treeProp] = value
   }
 
+  @JvmName("putAll")
+  internal fun putAll(treeProps: TreePropContainer?) {
+    if (treeProps != null) {
+      synchronized(treeProps.map) { map.putAll(treeProps.map) }
+    }
+  }
+
   operator fun <T : Any> get(key: Class<T>): T? {
     val treeProp = legacyTreePropOf(key)
     return map[treeProp] as T?

litho-core/src/main/java/com/facebook/litho/config/ComponentsConfiguration.kt

@@ -158,6 +158,7 @@ internal constructor(
      * size do not get unmounted when they go out of the viewport.
      */
     @JvmField val enableFixForIM: Boolean = false,
+    @JvmField val enableLifecycleOwnerWrapper: Boolean = false
 ) {
 
   val shouldAddRootHostViewOrDisableBgFgOutputs: Boolean =
@@ -332,6 +333,7 @@ internal constructor(
         baseConfig.useDefaultItemAnimatorInLazyCollections
     private var primitiveRecyclerBinderStrategy = baseConfig.primitiveRecyclerBinderStrategy
     private var enableFixForIM = baseConfig.enableFixForIM
+    private var enableLifecycleOwnerWrapper = baseConfig.enableLifecycleOwnerWrapper
 
     fun shouldNotifyVisibleBoundsChangeWhenNestedLithoViewBecomesInvisible(
         enabled: Boolean
@@ -436,6 +438,10 @@ internal constructor(
 
     fun enableFixForIM(enabled: Boolean): Builder = also { enableFixForIM = enabled }
 
+    fun enableLifecycleOwnerWrapper(enabled: Boolean): Builder = also {
+      enableLifecycleOwnerWrapper = enabled
+    }
+
     fun build(): ComponentsConfiguration {
       return baseConfig.copy(
           specsApiStateUpdateDuplicateDetectionEnabled =
@@ -471,7 +477,8 @@ internal constructor(
           skipSecondIsInWorkingRangeCheck = skipSecondIsInWorkingRangeCheck,
           enableVisibilityFixForNestedLithoView = enableVisibilityFixForNestedLithoView,
           useDefaultItemAnimatorInLazyCollections = useDefaultItemAnimatorInLazyCollections,
-          enableFixForIM = enableFixForIM)
+          enableFixForIM = enableFixForIM,
+          enableLifecycleOwnerWrapper = enableLifecycleOwnerWrapper)
     }
   }
 }

litho-core/src/main/java/com/facebook/litho/lifecycle/LithoLifecycleOwner.kt

@@ -18,8 +18,95 @@
 
 package com.facebook.litho.lifecycle
 
+import androidx.annotation.UiThread
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleObserver
 import androidx.lifecycle.LifecycleOwner
 import com.facebook.litho.TreeProp
 import com.facebook.litho.treePropOf
+import com.facebook.rendercore.utils.ThreadUtils.assertMainThread
 
 @JvmField val LifecycleOwnerTreeProp: TreeProp<LifecycleOwner?> = treePropOf { null }
+
+/**
+ * The [LifecycleOwnerWrapper] is an abstraction that will allow [LiveData] to observe a [Lifecycle]
+ * before the [LifecycleOwner] is determined. Once the [LifecycleOwner] is set on the
+ * [LifecycleOwnerWrapper] all the [LifecycleObserver] in [observers] will be added to the delegate
+ * [Lifecycle].
+ */
+class LifecycleOwnerWrapper
+internal constructor(
+    private var delegate: LifecycleOwner?,
+) : LifecycleOwner, Lifecycle() {
+
+  private val observers: MutableSet<LifecycleObserver> = HashSet()
+  private var latestState = State.RESUMED
+
+  override val currentState: State
+    get() = synchronized(this) { delegate?.lifecycle?.currentState ?: latestState }
+
+  @UiThread
+  @Synchronized
+  override fun addObserver(observer: LifecycleObserver) {
+    // add new observer
+    observers.add(observer)
+
+    val delegate = delegate
+    delegate?.lifecycle?.addObserver(observer)
+
+    // If the delegate is null, no events are dispatched to the observer at this moment. It will be
+    // dispatched when the delegate is set
+  }
+
+  @UiThread
+  @Synchronized
+  override fun removeObserver(observer: LifecycleObserver) {
+    val delegate = delegate
+    delegate?.lifecycle?.removeObserver(observer)
+
+    observers.remove(observer)
+  }
+
+  override val lifecycle: Lifecycle = this
+
+  @UiThread
+  @Synchronized
+  fun setDelegate(value: LifecycleOwner?) {
+    assertMainThread()
+
+    if (value == this) {
+      throw IllegalArgumentException("Cannot set a LifecycleOwnerWrapper as its own delegate")
+    }
+
+    // if the value is already the delegate then do nothing
+    if (value === delegate) {
+      return
+    }
+
+    // if a delegate is present then all observers need
+    // to be removed from it and added to the new value
+    delegate?.let {
+      for (observer in observers) {
+        it.lifecycle.removeObserver(observer)
+      }
+    }
+
+    // add all observers to value
+    value?.lifecycle?.let { lifecycle ->
+      for (observer in observers) {
+        lifecycle.addObserver(observer)
+      }
+    }
+
+    // store the latest state if value is null
+    delegate?.let {
+      if (value == null) {
+        latestState = it.lifecycle.currentState
+      }
+    }
+    // set value as the new delegate
+    delegate = value
+  }
+
+  fun hasObservers(): Boolean = observers.isNotEmpty()
+}

litho-it/src/test/java/com/facebook/litho/ComponentTreeTest.kt

@@ -398,7 +398,11 @@ class ComponentTreeTest {
         size,
         treePropContainer)
     val c = componentTree.mainThreadLayoutState!!.componentContext
-    assertThat(c.treePropContainer).isSameAs(treePropContainer)
+    if (ComponentsConfiguration.defaultInstance.enableLifecycleOwnerWrapper) {
+      assertThat(c.treePropContainer?.get(Any::class.java)).isEqualTo("hello world")
+    } else {
+      assertThat(c.treePropContainer).isSameAs(treePropContainer)
+    }
   }
 
   @Test

litho-live-data/src/test/kotlin/com/facebook/litho/livedata/UseLiveDataTest.kt

@@ -30,7 +30,11 @@ import com.facebook.litho.Component
 import com.facebook.litho.ComponentScope
 import com.facebook.litho.ComponentTree
 import com.facebook.litho.KComponent
+import com.facebook.litho.LithoView
+import com.facebook.litho.config.ComponentsConfiguration
 import com.facebook.litho.kotlin.widget.Text
+import com.facebook.litho.lifecycle.LifecycleOwnerTreeProp
+import com.facebook.litho.lifecycle.LifecycleOwnerWrapper
 import com.facebook.litho.testing.LithoViewRule
 import com.facebook.litho.testing.assertj.LithoAssertions.assertThat
 import com.facebook.litho.testing.testrunner.LithoTestRunner
@@ -51,6 +55,7 @@ class UseLiveDataTest {
           lithoVisibilityEventsController = {
             AOSPLithoVisibilityEventsController(fakeLifecycleOwner)
           })
+  @get:Rule val ruleWithoutVisibilityEventsController: LithoViewRule = LithoViewRule()
 
   @Test
   fun `should observe initial live data value`() {
@@ -187,6 +192,50 @@ class UseLiveDataTest {
     Assertions.assertThat(liveData.hasObservers()).isFalse
   }
 
+  @Test
+  fun `should observe lifecycle change when lifecycle owner is set after render`() {
+    ComponentsConfiguration.defaultInstance =
+        ComponentsConfiguration.defaultInstance.copy(enableLifecycleOwnerWrapper = true)
+
+    // Override new lifecycle owner and provider with local
+    val rule = ruleWithoutVisibilityEventsController
+    val fakeLifecycleOwner = FakeLifecycleOwner(Lifecycle.State.INITIALIZED)
+    val controller = AOSPLithoVisibilityEventsController(fakeLifecycleOwner)
+
+    val liveData = MutableLiveData("hello")
+    val component = MyComponent(liveData)
+
+    fakeLifecycleOwner.onEvent(Lifecycle.Event.ON_START)
+
+    // 1. the litho view should reflect the initial live data value
+    val lithoView = LithoView(rule.context)
+    val testLithoView = rule.render(lithoView = lithoView) { component }
+    assertThat(testLithoView).hasVisibleText("hello")
+
+    lithoView.subscribeComponentTreeToVisibilityEventsController(controller)
+
+    // 2. we change the lifecycle to an inactive state, and update the live data.
+    // in this case, the live data won't emit as there is no active observer.
+    fakeLifecycleOwner.onEvent(Lifecycle.Event.ON_STOP)
+    liveData.value = "world"
+    rule.idle()
+    assertThat(testLithoView).doesNotHaveVisibleText("world")
+
+    // 3. we resume the lifecycle, and therefore it should show the latest store live data change
+    fakeLifecycleOwner.onEvent(Lifecycle.Event.ON_RESUME)
+    rule.idle()
+    assertThat(testLithoView).hasVisibleText("world")
+
+    val tree = lithoView.componentTree!!
+    tree.release()
+    Assertions.assertThat(liveData.hasObservers()).isFalse
+    val owner = tree.getContext().getTreePropContainer()?.get(LifecycleOwnerTreeProp)
+    Assertions.assertThat((owner as LifecycleOwnerWrapper)?.hasObservers()).isFalse
+
+    ComponentsConfiguration.defaultInstance =
+        ComponentsConfiguration.defaultInstance.copy(enableLifecycleOwnerWrapper = false)
+  }
+
   private class MyComponent(private val liveData: LiveData<String>) : KComponent() {
 
     override fun ComponentScope.render(): Component {
@@ -247,5 +296,7 @@ class UseLiveDataTest {
     override fun removeObserver(observer: LifecycleObserver) {
       observers.remove(observer)
     }
+
+    fun hasObservers(): Boolean = observers.isNotEmpty()
   }
 }