Fix DynamicValue unsubscription for Primitive Components

Summary:
When there is a PrimitiveComponent with a Style that uses dynamic values, then these dynamic values are not unsubscribed properly during unmount. It causes that the component instance is kept in a map inside of DynamicPropsExtension which may lead to memory leaks.

This diff fixes the issue.

Reviewed By: adityasharat

Differential Revision: D49862897

fbshipit-source-id: 01fa9a6bcf1f318c414364c384cd077f555b6cfd

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

@@ -16,18 +16,22 @@
 
 package com.facebook.litho;
 
+import android.graphics.Rect;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.facebook.infer.annotation.Nullsafe;
+import com.facebook.litho.config.ComponentsConfiguration;
 import com.facebook.rendercore.RenderTreeNode;
 import com.facebook.rendercore.RenderUnit;
 import com.facebook.rendercore.extensions.ExtensionState;
 import com.facebook.rendercore.extensions.MountExtension;
 import com.facebook.rendercore.extensions.OnItemCallbacks;
+import java.util.Map;
 
 @Nullsafe(Nullsafe.Mode.LOCAL)
 public class DynamicPropsExtension
-    extends MountExtension<Void, DynamicPropsExtension.DynamicPropsExtensionState>
+    extends MountExtension<
+        DynamicPropsExtensionInput, DynamicPropsExtension.DynamicPropsExtensionState>
     implements OnItemCallbacks<DynamicPropsExtension.DynamicPropsExtensionState> {
 
   private static final DynamicPropsExtension sInstance = new DynamicPropsExtension();
@@ -43,21 +47,61 @@ public static DynamicPropsExtension getInstance() {
     return sInstance;
   }
 
+  @Override
+  public void beforeMount(
+      ExtensionState<DynamicPropsExtensionState> extensionState,
+      @Nullable DynamicPropsExtensionInput dynamicPropsExtensionInput,
+      Rect localVisibleRect) {
+    extensionState.getState().mInput =
+        dynamicPropsExtensionInput != null
+            ? dynamicPropsExtensionInput.getDynamicValueOutputs()
+            : null;
+  }
+
+  @Override
+  public void onUnmount(ExtensionState<DynamicPropsExtensionState> extensionState) {
+    extensionState.releaseAllAcquiredReferences();
+    final DynamicPropsExtensionState state = extensionState.getState();
+    state.mInput = null;
+    state.mPreviousInput = null;
+  }
+
+  @Override
+  public void afterMount(ExtensionState<DynamicPropsExtensionState> extensionState) {
+    final DynamicPropsExtensionState state = extensionState.getState();
+    state.mPreviousInput = state.mInput;
+    state.mInput = null;
+  }
+
   @Override
   public void onBindItem(
       final ExtensionState<DynamicPropsExtensionState> extensionState,
       final RenderUnit<?> renderUnit,
       final Object content,
       final @Nullable Object layoutData) {
-    if (renderUnit instanceof LithoRenderUnit) {
-      final LithoRenderUnit lithoRenderUnit = (LithoRenderUnit) renderUnit;
-      final DynamicPropsExtensionState state = extensionState.getState();
-
-      state.mDynamicPropsManager.onBindComponentToContent(
-          lithoRenderUnit.getComponent(),
-          lithoRenderUnit.componentContext,
-          lithoRenderUnit.commonDynamicProps,
-          content);
+    final DynamicPropsExtensionState state = extensionState.getState();
+
+    if (ComponentsConfiguration.enablePrimitiveDynamicPropsExtensionFix) {
+      final DynamicValueOutput dynamicValueOutput =
+          state.mInput != null ? state.mInput.get(renderUnit.getId()) : null;
+
+      if (dynamicValueOutput != null) {
+        state.mDynamicPropsManager.onBindComponentToContent(
+            dynamicValueOutput.getComponent(),
+            dynamicValueOutput.getScopedContext(),
+            dynamicValueOutput.getCommonDynamicProps(),
+            content);
+      }
+    } else {
+      if (renderUnit instanceof LithoRenderUnit) {
+        final LithoRenderUnit lithoRenderUnit = (LithoRenderUnit) renderUnit;
+
+        state.mDynamicPropsManager.onBindComponentToContent(
+            lithoRenderUnit.getComponent(),
+            lithoRenderUnit.componentContext,
+            lithoRenderUnit.commonDynamicProps,
+            content);
+      }
     }
   }
 
@@ -67,12 +111,23 @@ public void onUnbindItem(
       final RenderUnit<?> renderUnit,
       final Object content,
       final @Nullable Object layoutData) {
-    if (renderUnit instanceof LithoRenderUnit) {
-      final LithoRenderUnit lithoRenderUnit = (LithoRenderUnit) renderUnit;
-      final DynamicPropsExtensionState state = extensionState.getState();
-
-      state.mDynamicPropsManager.onUnbindComponent(
-          lithoRenderUnit.getComponent(), lithoRenderUnit.commonDynamicProps, content);
+    final DynamicPropsExtensionState state = extensionState.getState();
+
+    if (ComponentsConfiguration.enablePrimitiveDynamicPropsExtensionFix) {
+      final DynamicValueOutput dynamicValueOutput =
+          state.mPreviousInput != null ? state.mPreviousInput.get(renderUnit.getId()) : null;
+
+      if (dynamicValueOutput != null) {
+        state.mDynamicPropsManager.onUnbindComponent(
+            dynamicValueOutput.getComponent(), dynamicValueOutput.getCommonDynamicProps(), content);
+      }
+    } else {
+      if (renderUnit instanceof LithoRenderUnit) {
+        final LithoRenderUnit lithoRenderUnit = (LithoRenderUnit) renderUnit;
+
+        state.mDynamicPropsManager.onUnbindComponent(
+            lithoRenderUnit.getComponent(), lithoRenderUnit.commonDynamicProps, content);
+      }
     }
   }
 
@@ -116,6 +171,10 @@ public void onBoundsAppliedToItem(
   static class DynamicPropsExtensionState {
     private final DynamicPropsManager mDynamicPropsManager = new DynamicPropsManager();
 
+    @Nullable private Map<Long, DynamicValueOutput> mInput;
+
+    @Nullable private Map<Long, DynamicValueOutput> mPreviousInput;
+
     @VisibleForTesting
     public DynamicPropsManager getDynamicPropsManager() {
       return mDynamicPropsManager;

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.litho
+
+interface DynamicPropsExtensionInput {
+  val dynamicValueOutputs: Map<Long, DynamicValueOutput>
+}

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

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.litho
+
+import android.util.SparseArray
+
+class DynamicValueOutput(
+    val component: Component,
+    val scopedContext: ComponentContext?,
+    val commonDynamicProps: SparseArray<out DynamicValue<*>>
+)

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

@@ -72,7 +72,8 @@ public class LayoutState
         TransitionsExtensionInput,
         EndToEndTestingExtensionInput,
         PotentiallyPartialResult,
-        ViewAttributesInput {
+        ViewAttributesInput,
+        DynamicPropsExtensionInput {
 
   @Nullable private Transition.RootBoundsTransition mRootWidthAnimation;
   @Nullable private Transition.RootBoundsTransition mRootHeightAnimation;
@@ -110,6 +111,7 @@ public static boolean isFromSyncLayout(@RenderSource int source) {
   final LongSparseArray<Integer> mOutputsIdToPositionMap = new LongSparseArray<>(8);
   final Map<Long, ViewAttributes> mRenderUnitsWithViewAttributes = new HashMap<>(8);
   final Map<Long, IncrementalMountOutput> mIncrementalMountOutputs = new LinkedHashMap<>(8);
+  final Map<Long, DynamicValueOutput> mDynamicValueOutputs = new LinkedHashMap<>(8);
   final ArrayList<IncrementalMountOutput> mMountableOutputTops = new ArrayList<>();
   final ArrayList<IncrementalMountOutput> mMountableOutputBottoms = new ArrayList<>();
   final LongSparseArray<AnimatableItem> mAnimatableItems = new LongSparseArray<>(8);
@@ -787,4 +789,9 @@ public Transition.RootBoundsTransition getRootWidthAnimation() {
   public VisibilityBoundsTransformer getVisibilityBoundsTransformer() {
     return getComponentContext().getVisibilityBoundsTransformer();
   }
+
+  @Override
+  public Map<Long, DynamicValueOutput> getDynamicValueOutputs() {
+    return mDynamicValueOutputs;
+  }
 }

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

@@ -1086,6 +1086,18 @@ private static void addRenderTreeNode(
       layoutState.mRenderUnitsWithViewAttributes.put(id, attrs);
     }
 
+    if (node.getRenderUnit() instanceof LithoRenderUnit) {
+      final LithoRenderUnit lithoRenderUnit = (LithoRenderUnit) node.getRenderUnit();
+      if (lithoRenderUnit.commonDynamicProps != null) {
+        layoutState.mDynamicValueOutputs.put(
+            lithoRenderUnit.getId(),
+            new DynamicValueOutput(
+                lithoRenderUnit.getComponent(),
+                lithoRenderUnit.componentContext,
+                lithoRenderUnit.commonDynamicProps));
+      }
+    }
+
     final AnimatableItem animatableItem =
         createAnimatableItem(unit, absoluteBounds, type, transitionId);
 

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

@@ -234,6 +234,8 @@ public static boolean isRenderInfoDebuggingEnabled() {
 
   public static int useCachedLruCacheSize = 0;
 
+  public static boolean enablePrimitiveDynamicPropsExtensionFix = false;
+
   public static boolean shouldOverrideHasTransientState = false;
 
   public static boolean shouldUseLruCacheForUseCached() {

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

@@ -434,6 +434,120 @@ class DynamicPropsTest {
     alphaDV.set(1f)
     assertThat(hostView.alpha).isEqualTo(1f)
   }
+
+  @Test
+  fun testPrimitiveWithDynamicValueIsCorrectlyUnsubscribed() {
+    val alphaDV1 = DynamicValue(1f)
+    var lithoView =
+        legacyLithoViewRule
+            .attachToWindow()
+            .setRoot(
+                SimpleTestPrimitiveComponent(
+                    style = Style.width(80.px).height(80.px).alpha(alphaDV1)))
+            .measure()
+            .layout()
+            .lithoView
+    legacyLithoViewRule.idle()
+
+    assertThat(alphaDV1.numberOfListeners).isEqualTo(1)
+
+    // simulate update
+    val alphaDV2 = DynamicValue(1f)
+    lithoView =
+        legacyLithoViewRule
+            .attachToWindow()
+            .setRoot(
+                SimpleTestPrimitiveComponent(
+                    style = Style.width(80.px).height(80.px).alpha(alphaDV2)))
+            .measure()
+            .layout()
+            .lithoView
+    legacyLithoViewRule.idle()
+
+    // should unsubscribe from the old DV and subscribe to the new DV
+    if (ComponentsConfiguration.enablePrimitiveDynamicPropsExtensionFix) {
+      // works correctly when fix is applied
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(0)
+      assertThat(alphaDV2.numberOfListeners).isEqualTo(1)
+    } else {
+      // doesn't work without the fix
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(1)
+      assertThat(alphaDV2.numberOfListeners).isEqualTo(0)
+    }
+
+    lithoView.unmountAllItems()
+    legacyLithoViewRule.idle()
+
+    // should unsubscribe from all DVs
+    if (ComponentsConfiguration.enablePrimitiveDynamicPropsExtensionFix) {
+      // works correctly when fix is applied
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(0)
+    } else {
+      // doesn't work without the fix
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(1)
+    }
+    assertThat(alphaDV2.numberOfListeners).isEqualTo(0)
+  }
+
+  @Test
+  fun testWrappedPrimitiveWithDynamicValueIsCorrectlyUnsubscribed() {
+    val alphaDV1 = DynamicValue(1f)
+    var lithoView =
+        legacyLithoViewRule
+            .attachToWindow()
+            .setRoot(
+                Wrapper.create(context)
+                    .delegate(
+                        SimpleTestPrimitiveComponent(style = Style.width(80.px).height(80.px)))
+                    .kotlinStyle(Style.alpha(alphaDV1))
+                    .build())
+            .measure()
+            .layout()
+            .lithoView
+    legacyLithoViewRule.idle()
+
+    assertThat(alphaDV1.numberOfListeners).isEqualTo(1)
+
+    // simulate update
+    val alphaDV2 = DynamicValue(1f)
+    lithoView =
+        legacyLithoViewRule
+            .attachToWindow()
+            .setRoot(
+                Wrapper.create(context)
+                    .delegate(
+                        SimpleTestPrimitiveComponent(style = Style.width(80.px).height(80.px)))
+                    .kotlinStyle(Style.alpha(alphaDV2))
+                    .build())
+            .measure()
+            .layout()
+            .lithoView
+    legacyLithoViewRule.idle()
+
+    // should unsubscribe from the old DV and subscribe to the new DV
+    if (ComponentsConfiguration.enablePrimitiveDynamicPropsExtensionFix) {
+      // works correctly when fix is applied
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(0)
+      assertThat(alphaDV2.numberOfListeners).isEqualTo(1)
+    } else {
+      // doesn't work without the fix
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(1)
+      assertThat(alphaDV2.numberOfListeners).isEqualTo(0)
+    }
+
+    lithoView.unmountAllItems()
+    legacyLithoViewRule.idle()
+
+    // should unsubscribe from all DVs
+    if (ComponentsConfiguration.enablePrimitiveDynamicPropsExtensionFix) {
+      // works correctly when fix is applied
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(0)
+    } else {
+      // doesn't work without the fix
+      assertThat(alphaDV1.numberOfListeners).isEqualTo(1)
+    }
+    assertThat(alphaDV2.numberOfListeners).isEqualTo(0)
+  }
 }
 
 private class SimpleTestPrimitiveComponent(