Add more fine-grainted mechanism to understand if TreeFuture was cancelled by Policy

Summary:
This is a bit of a defensive approach to try to avoid the issues with the cancellation experiment.

The steps are the following:

1. Adding a flag in `TreeFuture` that marks it as "cancelled" whenever we call cancel due to our policy.
2. When we cancel the Future, it can happen that an exception is triggered. If so, we verify if it is one of the expected exceptions **and** if it was marked cancelled by us. If that is the case, we throw a specific `FutureCancelledByPolicyException` and this should be caught in the `CancellationHelper` and then a `TreeFutureResult.cancelled()` should be returned.

Reviewed By: adityasharat

Differential Revision: D44495245

fbshipit-source-id: 29d0d02d9d6433e3a58ace9e5e82d0334fa17e5a

Author: Fabio Carballo

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

@@ -107,7 +107,8 @@ public ResolveTreeFuture(
             resolveVersion,
             component.getId(),
             mComponentContext.getTreeProps(),
-            ExecutionModeKt.getExecutionMode(source));
+            ExecutionModeKt.getExecutionMode(source),
+            source);
 
     // Allow interrupt to happen during tryRegisterForResponse when config is enabled.
     mEnableEarlyInterrupt = ComponentsConfiguration.isInterruptEarlyWithSplitFuturesEnabled;

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

@@ -26,6 +26,7 @@
 import android.os.Process;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
+import com.facebook.litho.cancellation.FutureCancelledByPolicyException;
 import com.facebook.litho.config.ComponentsConfiguration;
 import com.facebook.rendercore.Systracer;
 import com.facebook.rendercore.instrumentation.FutureInstrumenter;
@@ -102,6 +103,13 @@ public TreeFutureResult<T> call() {
   /** Returns an integer that identifies uniquely the version of this {@link TreeFuture}. */
   public abstract int getVersion();
 
+  /**
+   * This indicates if this {@link TreeFuture} was canceleled by the framework. If so, we can
+   * consider that any {@link InterruptedException} or {@link CancellationException} was actively
+   * triggered by us.
+   */
+  private boolean mCancelled = false;
+
   /**
    * This method will forcefully cancel any possible execution or remaining execution of the task
    * associated to this {@link TreeFuture}.
@@ -118,8 +126,15 @@ public TreeFutureResult<T> call() {
    * @return false if the task could not be cancelled, typically because it has already completed
    *     normally; true otherwise
    */
-  public boolean forceCancellation() {
-    return mFutureTask.cancel(true);
+  public boolean cancel(boolean useInterruption) {
+    mCancelled = true;
+
+    if (useInterruption) {
+      return mFutureTask.cancel(true);
+    } else {
+      release();
+      return true;
+    }
   }
 
   /** Calculates a new result for this TreeFuture. */
@@ -341,7 +356,7 @@ TreeFutureResult<T> runAndGet(@LayoutState.CalculateLayoutSource final int sourc
       didRaiseThreadPriority = false;
     }
 
-    TreeFutureResult<T> treeFutureResult;
+    TreeFutureResult<T> treeFutureResult = null;
 
     final boolean shouldTrace = notRunningOnMyThread && ComponentsSystrace.isTracing();
     try {
@@ -404,24 +419,37 @@ TreeFutureResult<T> runAndGet(@LayoutState.CalculateLayoutSource final int sourc
     } catch (ExecutionException | InterruptedException | CancellationException e) {
       onWaitEnd(shouldTrace, true);
 
-      final Throwable cause = e.getCause();
-      if (cause instanceof RuntimeException) {
-        throw (RuntimeException) cause;
+      if (isExpectedInterruptionOrCancellation(e)) {
+        throw new FutureCancelledByPolicyException();
       } else {
-        throw new RuntimeException(e.getMessage(), e);
+        final Throwable cause = e.getCause();
+        if (cause instanceof RuntimeException) {
+          throw (RuntimeException) cause;
+        } else {
+          throw new RuntimeException(e.getMessage(), e);
+        }
       }
     } finally {
       onGetEnd(shouldTrace);
     }
 
     synchronized (TreeFuture.this) {
+      if (mCancelled) {
+        return TreeFutureResult.cancelled();
+      }
+
       if (mReleased) {
         return TreeFutureResult.interruptWithMessage(FUTURE_RESULT_NULL_REASON_RELEASED);
       }
+
       return treeFutureResult;
     }
   }
 
+  private boolean isExpectedInterruptionOrCancellation(Exception e) {
+    return (e instanceof CancellationException || e instanceof InterruptedException) && mCancelled;
+  }
+
   /**
    * Given a provided tree-future, this method will track it via a given list, and run it.
    *

litho-core/src/main/java/com/facebook/litho/cancellation/CancellationHelper.kt

@@ -26,7 +26,6 @@ import com.facebook.litho.TreeFuture.TreeFutureResult
 import com.facebook.litho.cancellation.CancellationPolicy.CancellationExecutionMode
 import com.facebook.litho.stats.LithoStats
 import java.util.LinkedList
-import java.util.concurrent.CancellationException
 
 /**
  * This is used to run either *Resolve* or *Layout* with a cancellation strategy.
@@ -79,10 +78,8 @@ F : RequestMetadataSupplier<M> {
 
   return try {
     TreeFuture.trackAndRunTreeFuture(treeFuture, futures, source, mutex, futureExecutionListener)
-  } catch (exception: RuntimeException) {
-    if (exception.cause is CancellationException) {
-      debugLog { "A cancellation exception on future ${treeFuture.version}" }
-    }
+  } catch (e: FutureCancelledByPolicyException) {
+    debugLog { "Future was cancelled by CancellationPolicy" }
     TreeFutureResult.cancelled()
   }
 }
@@ -109,11 +106,7 @@ T : PotentiallyPartialResult {
     if (futureToCancel != null) {
       incrementCancellationStats(attribution)
 
-      when (cancellationMode) {
-        CancellationExecutionMode.SHORT_CIRCUIT -> futureToCancel.release()
-        CancellationExecutionMode.INTERRUPT -> futureToCancel.forceCancellation()
-      }
-
+      futureToCancel.cancel(cancellationMode == CancellationExecutionMode.INTERRUPT)
       cancelledFutures.add(futureToCancel)
 
       debugLog { "Cancelled future (${futureToCancel.version}) with success." }

litho-core/src/main/java/com/facebook/litho/cancellation/FutureCancelledByPolicyException.kt

@@ -0,0 +1,24 @@
+/*
+ * 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.cancellation
+
+/**
+ * This exception should only be thrown whenever there is an [InterruptedException] or
+ * [CancellationException] thrown inside a [TreeFuture] while being marked with
+ * [TreeFuture.isCancelled]
+ */
+internal class FutureCancelledByPolicyException() : RuntimeException()

litho-core/src/main/java/com/facebook/litho/cancellation/LayoutCancellationPolicy.kt

@@ -72,7 +72,7 @@ interface LayoutCancellationPolicy : CancellationPolicy<LayoutMetadata> {
       }
 
       return if (cancellableRequests.isNotEmpty()) {
-        Result.CancelRunningRequests(ongoingRequests.map(LayoutMetadata::localVersion))
+        Result.CancelRunningRequests(cancellableRequests.map(LayoutMetadata::localVersion))
       } else {
         Result.ProcessIncomingRequest
       }

litho-core/src/main/java/com/facebook/litho/cancellation/ResolveCancellationPolicy.kt

@@ -18,6 +18,7 @@
 
 package com.facebook.litho.cancellation
 
+import com.facebook.litho.LayoutState
 import com.facebook.litho.cancellation.CancellationPolicy.Result
 
 /**
@@ -58,18 +59,25 @@ interface ResolveCancellationPolicy : CancellationPolicy<ResolveMetadata> {
 
       for (runningResolve in ongoingRequests) {
         if (runningResolve.isEquivalentTo(incomingRequest)) {
-          if (incomingRequest.executionMode == ExecutionMode.ASYNC) {
-            return Result.DropIncomingRequest
+          when {
+            incomingRequest.isAsync && !incomingRequest.isStateUpdate -> {
+              return Result.DropIncomingRequest
+            }
+            incomingRequest.isAsync &&
+                incomingRequest.isStateUpdate &&
+                !runningResolve.isStateUpdate -> {
+              return Result.DropIncomingRequest
+            }
           }
         } else {
-          if (runningResolve.executionMode == ExecutionMode.ASYNC) {
+          if (runningResolve.isAsync) {
             cancellableResolves.add(runningResolve)
           }
         }
       }
 
       return if (cancellableResolves.isNotEmpty()) {
-        Result.CancelRunningRequests(ongoingRequests.map { it.id })
+        Result.CancelRunningRequests(cancellableResolves.map { it.id })
       } else {
         Result.ProcessIncomingRequest
       }
@@ -100,23 +108,39 @@ interface ResolveCancellationPolicy : CancellationPolicy<ResolveMetadata> {
 
       for (runningResolve in ongoingRequests) {
         if (runningResolve.isEquivalentTo(incomingRequest)) {
-          if (incomingRequest.executionMode == ExecutionMode.ASYNC) {
-            return Result.DropIncomingRequest
-          } else if (runningResolve.executionMode == ExecutionMode.ASYNC) {
-            cancellableResolves.add(runningResolve)
+          when {
+            incomingRequest.isAsync && !incomingRequest.isStateUpdate -> {
+              return Result.DropIncomingRequest
+            }
+            incomingRequest.isAsync &&
+                incomingRequest.isStateUpdate &&
+                !runningResolve.isStateUpdate -> {
+              return Result.DropIncomingRequest
+            }
+            runningResolve.isAsync -> {
+              cancellableResolves.add(runningResolve)
+            }
           }
         } else {
-          if (runningResolve.executionMode == ExecutionMode.ASYNC) {
+          if (runningResolve.isAsync) {
             cancellableResolves.add(runningResolve)
           }
         }
       }
 
       return if (cancellableResolves.isNotEmpty()) {
-        Result.CancelRunningRequests(ongoingRequests.map { it.id })
+        Result.CancelRunningRequests(cancellableResolves.map { it.id })
       } else {
         Result.ProcessIncomingRequest
       }
     }
   }
 }
+
+private val ResolveMetadata.isAsync: Boolean
+  get() = executionMode == ExecutionMode.ASYNC
+
+private val ResolveMetadata.isStateUpdate: Boolean
+  get() =
+      source == LayoutState.CalculateLayoutSource.UPDATE_STATE_ASYNC ||
+          source == LayoutState.CalculateLayoutSource.UPDATE_STATE_SYNC

litho-core/src/main/java/com/facebook/litho/cancellation/ResolveMetadata.kt

@@ -18,14 +18,16 @@
 
 package com.facebook.litho.cancellation
 
+import com.facebook.litho.LayoutState
 import com.facebook.litho.TreeProps
 
 /** This metadata encompasses all characteristics associated to a specific `Resolve` request. */
 data class ResolveMetadata(
     val localVersion: Int,
     val componentId: Int,
     val treeProps: TreeProps?,
-    val executionMode: ExecutionMode
+    val executionMode: ExecutionMode,
+    @LayoutState.CalculateLayoutSource val source: Int
 ) {
 
   val id: Int = localVersion

litho-core/src/test/java/com/facebook/litho/cancellation/DefaultResolveCancellationPolicyTest.kt

@@ -18,6 +18,7 @@
 
 package com.facebook.litho.cancellation
 
+import com.facebook.litho.LayoutState
 import com.facebook.litho.TreeProps
 import com.facebook.litho.cancellation.CancellationPolicy.CancellationExecutionMode
 import com.facebook.litho.cancellation.CancellationPolicy.Result
@@ -112,18 +113,45 @@ class DefaultResolveCancellationPolicyTest {
   }
 
   @Test
-  fun `on async resolve request with equivalent async request running - drop incoming request`() {
+  fun `on async non-state update resolve request with equivalent async request running - drop incoming request`() {
     val runningResolves =
-        listOf(resolveMetadata.copy(localVersion = 0, executionMode = ExecutionMode.ASYNC))
+        listOf(
+            resolveMetadata.copy(
+                localVersion = 0,
+                executionMode = ExecutionMode.ASYNC,
+                source = LayoutState.CalculateLayoutSource.UPDATE_STATE_ASYNC))
 
     val incomingResolve =
-        resolveMetadata.copy(executionMode = ExecutionMode.ASYNC, localVersion = 1)
+        resolveMetadata.copy(
+            executionMode = ExecutionMode.ASYNC,
+            localVersion = 1,
+            source = LayoutState.CalculateLayoutSource.SET_ROOT_ASYNC)
 
     val result = cancellationEvaluator.evaluate(runningResolves, incomingResolve)
 
     assertThat(result).isEqualTo(Result.DropIncomingRequest)
   }
 
+  @Test
+  fun `on async state update resolve request with equivalent async state update request running - process incoming request`() {
+    val runningResolves =
+        listOf(
+            resolveMetadata.copy(
+                localVersion = 0,
+                executionMode = ExecutionMode.ASYNC,
+                source = LayoutState.CalculateLayoutSource.UPDATE_STATE_ASYNC))
+
+    val incomingResolve =
+        resolveMetadata.copy(
+            executionMode = ExecutionMode.ASYNC,
+            localVersion = 1,
+            source = LayoutState.CalculateLayoutSource.UPDATE_STATE_ASYNC)
+
+    val result = cancellationEvaluator.evaluate(runningResolves, incomingResolve)
+
+    assertThat(result).isEqualTo(Result.ProcessIncomingRequest)
+  }
+
   @Test
   fun `on sync resolve request with equivalent sync request running - drop incoming request`() {
     val runningResolves =
@@ -153,5 +181,6 @@ class DefaultResolveCancellationPolicyTest {
           componentId = 1,
           localVersion = 0,
           treeProps = TreeProps().apply { put(String::class.java, "a property") },
-          executionMode = ExecutionMode.SYNC)
+          executionMode = ExecutionMode.SYNC,
+          source = LayoutState.CalculateLayoutSource.SET_ROOT_SYNC)
 }