Fix back callback ordering for quick back gestures in succession

Bug: 324036420
Flag: Flag: ACONFIG com.android.systemui.predictive_back_system_animations TEAMFOOD
Test: atest BackProgressAnimatorTest
Test: Manual, i.e. verifying that starting a back gesture quickly after another one is cancelled doesn't invoke callback functions in wrong order
Change-Id: Ie5ebef743f8de07032b304f8a89dff03f318426a

3 files changed, 57 insertions, 14 deletions

diff --git a/core/java/android/window/BackProgressAnimator.java b/core/java/android/window/BackProgressAnimator.java
index e7280d0e7867..40e28cbbbd05 100644
--- a/core/java/android/window/BackProgressAnimator.java
+++ b/core/java/android/window/BackProgressAnimator.java
@@ -17,6 +17,7 @@
 package android.window;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.util.FloatProperty;
 
 import com.android.internal.dynamicanimation.animation.DynamicAnimation;
@@ -44,6 +45,14 @@ public class BackProgressAnimator {
     private float mProgress = 0;
     private BackMotionEvent mLastBackEvent;
     private boolean mBackAnimationInProgress = false;
+    @Nullable
+    private Runnable mBackCancelledFinishRunnable;
+    private final DynamicAnimation.OnAnimationEndListener mOnAnimationEndListener =
+            (animation, canceled, value, velocity) -> {
+                invokeBackCancelledRunnable();
+                reset();
+            };
+
 
     private void setProgress(float progress) {
         mProgress = progress;
@@ -116,6 +125,11 @@ public class BackProgressAnimator {
      * Resets the back progress animation. This should be called when back is invoked or cancelled.
      */
     public void reset() {
+        if (mBackCancelledFinishRunnable != null) {
+            // Ensure that last progress value that apps see is 0
+            updateProgressValue(0);
+            invokeBackCancelledRunnable();
+        }
         mSpring.animateToFinalPosition(0);
         if (mSpring.canSkipToEnd()) {
             mSpring.skipToEnd();
@@ -136,17 +150,8 @@ public class BackProgressAnimator {
      * @param finishCallback the callback to be invoked when the progress is reach to 0.
      */
     public void onBackCancelled(@NonNull Runnable finishCallback) {
-        final DynamicAnimation.OnAnimationEndListener listener =
-                new DynamicAnimation.OnAnimationEndListener() {
-            @Override
-            public void onAnimationEnd(DynamicAnimation animation, boolean canceled, float value,
-                    float velocity) {
-                mSpring.removeEndListener(this);
-                finishCallback.run();
-                reset();
-            }
-        };
-        mSpring.addEndListener(listener);
+        mBackCancelledFinishRunnable = finishCallback;
+        mSpring.addEndListener(mOnAnimationEndListener);
         mSpring.animateToFinalPosition(0);
     }
 
@@ -164,4 +169,10 @@ public class BackProgressAnimator {
                         progress / SCALE_FACTOR, mLastBackEvent.getSwipeEdge()));
     }
 
-}
+    private void invokeBackCancelledRunnable() {
+        mSpring.removeEndListener(mOnAnimationEndListener);
+        mBackCancelledFinishRunnable.run();
+        mBackCancelledFinishRunnable = null;
+    }
+
+}
\ No newline at end of file
diff --git a/core/java/android/window/WindowOnBackInvokedDispatcher.java b/core/java/android/window/WindowOnBackInvokedDispatcher.java
index 5c911f4a632a..45d7767380a1 100644
--- a/core/java/android/window/WindowOnBackInvokedDispatcher.java
+++ b/core/java/android/window/WindowOnBackInvokedDispatcher.java
@@ -371,11 +371,11 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
                 }
                 final OnBackAnimationCallback callback = getBackAnimationCallback();
                 if (callback != null) {
+                    mProgressAnimator.reset();
                     callback.onBackStarted(new BackEvent(
                             backEvent.getTouchX(), backEvent.getTouchY(),
                             backEvent.getProgress(), backEvent.getSwipeEdge()));
-                    mProgressAnimator.onBackStarted(backEvent, event ->
-                            callback.onBackProgressed(event));
+                    mProgressAnimator.onBackStarted(backEvent, callback::onBackProgressed);
                 }
             });
         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackProgressAnimatorTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackProgressAnimatorTest.java
index 91503b1c3619..7e26577e96d4 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackProgressAnimatorTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackProgressAnimatorTest.java
@@ -18,6 +18,7 @@ package com.android.wm.shell.back;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 
 import android.os.Handler;
 import android.os.Looper;
@@ -28,6 +29,7 @@ import android.window.BackMotionEvent;
 import android.window.BackProgressAnimator;
 
 import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -102,6 +104,36 @@ public class BackProgressAnimatorTest {
         assertEquals(mReceivedBackEvent.getProgress(), mTargetProgress, 0 /* delta */);
     }
 
+    @Test
+    public void testResetCallsCancelCallbackImmediately() throws InterruptedException {
+        // Give the animator some progress.
+        final BackMotionEvent backEvent = backMotionEventFrom(100, mTargetProgress);
+        mMainThreadHandler.post(
+                () -> mProgressAnimator.onBackProgressed(backEvent));
+        mTargetProgressCalled.await(1, TimeUnit.SECONDS);
+        assertNotNull(mReceivedBackEvent);
+
+        mTargetProgress = 0;
+        mReceivedBackEvent = null;
+        mTargetProgressCalled = new CountDownLatch(1);
+
+        CountDownLatch cancelCallbackCalled = new CountDownLatch(1);
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(
+                () -> mProgressAnimator.onBackCancelled(cancelCallbackCalled::countDown));
+
+        // verify onBackProgressed and onBackCancelled not yet called
+        assertNull(mReceivedBackEvent);
+        assertEquals(1, cancelCallbackCalled.getCount());
+
+        // call reset
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> mProgressAnimator.reset());
+
+        // verify that back event with progress 0 is sent and cancel callback is invoked
+        assertNotNull(mReceivedBackEvent);
+        assertEquals(mReceivedBackEvent.getProgress(), mTargetProgress, 0 /* delta */);
+        assertEquals(0, cancelCallbackCalled.getCount());
+    }
+
     private void onGestureProgress(BackEvent backEvent) {
         if (mTargetProgress == backEvent.getProgress()) {
             mReceivedBackEvent = backEvent;