Merge "Fixing exception when applying transaction while the surface is released" into ub-launcher3-rvc-dev
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java
index c18a0fd..e5782e7 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java
@@ -33,11 +33,11 @@
 import com.android.launcher3.statehandlers.DepthController;
 import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.quickstep.util.RemoteAnimationProvider;
+import com.android.quickstep.util.SurfaceTransactionApplier;
 import com.android.quickstep.util.TaskViewSimulator;
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.views.RecentsView;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
 
 /**
  * Provider for the atomic (for 3-button mode) remote window animation from the app to the overview.
@@ -132,8 +132,7 @@
 
         TransformParams params = new TransformParams()
                 .setTargetSet(targets)
-                .setSyncTransactionApplier(
-                        new SyncRtSurfaceTransactionApplierCompat(mActivity.getRootView()));
+                .setSyncTransactionApplier(new SurfaceTransactionApplier(mActivity.getRootView()));
 
         AnimatedFloat recentsAlpha = new AnimatedFloat(() -> { });
         params.setBaseBuilderProxy((builder, app, p)
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
index 045f459..6cbe794 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
@@ -40,13 +40,14 @@
 import com.android.quickstep.RecentsAnimationCallbacks.RecentsAnimationListener;
 import com.android.quickstep.util.ActiveGestureLog;
 import com.android.quickstep.util.ActivityInitListener;
+import com.android.quickstep.util.RectFSpringAnim;
+import com.android.quickstep.util.SurfaceTransactionApplier;
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.InputConsumerController;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
 
 import java.util.ArrayList;
 import java.util.function.Consumer;
@@ -114,10 +115,10 @@
     public abstract Intent getLaunchIntent();
 
     protected void linkRecentsViewScroll() {
-        SyncRtSurfaceTransactionApplierCompat.create(mRecentsView, applier -> {
+        SurfaceTransactionApplier.create(mRecentsView, applier -> {
             mTransformParams.setSyncTransactionApplier(applier);
             runOnRecentsAnimationStart(() ->
-                    mRecentsAnimationTargets.addDependentTransactionApplier(applier));
+                    mRecentsAnimationTargets.addReleaseCheck(applier));
         });
 
         mRecentsView.setOnScrollChangeListener((v, scrollX, scrollY, oldScrollX, oldScrollY) -> {
@@ -365,6 +366,17 @@
         }
     }
 
+    @Override
+    protected RectFSpringAnim createWindowAnimationToHome(float startProgress,
+            HomeAnimationFactory homeAnimationFactory) {
+        RectFSpringAnim anim =
+                super.createWindowAnimationToHome(startProgress, homeAnimationFactory);
+        if (mRecentsAnimationTargets != null) {
+            mRecentsAnimationTargets.addReleaseCheck(anim);
+        }
+        return anim;
+    }
+
     public interface Factory {
 
         BaseSwipeUpHandler newHandler(
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java
index c68d6e2..e2e25f3 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java
@@ -42,6 +42,7 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.statehandlers.DepthController;
 import com.android.launcher3.util.DefaultDisplay;
+import com.android.quickstep.util.SurfaceTransactionApplier;
 import com.android.quickstep.util.TaskViewSimulator;
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.views.RecentsView;
@@ -49,7 +50,6 @@
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
 
 /**
  * Utility class for helpful methods related to {@link TaskView} objects and their tasks.
@@ -128,11 +128,10 @@
             RemoteAnimationTargetCompat[] wallpaperTargets, DepthController depthController,
             PendingAnimation out) {
 
-        SyncRtSurfaceTransactionApplierCompat applier =
-                new SyncRtSurfaceTransactionApplierCompat(v);
+        SurfaceTransactionApplier applier = new SurfaceTransactionApplier(v);
         final RemoteAnimationTargets targets =
                 new RemoteAnimationTargets(appTargets, wallpaperTargets, MODE_OPENING);
-        targets.addDependentTransactionApplier(applier);
+        targets.addReleaseCheck(applier);
 
         TransformParams params = new TransformParams()
                     .setSyncTransactionApplier(applier)
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RectFSpringAnim.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RectFSpringAnim.java
index 8a6c4a1..e5d2c53 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RectFSpringAnim.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RectFSpringAnim.java
@@ -29,6 +29,7 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.FlingSpringAnim;
 import com.android.launcher3.util.DynamicResource;
+import com.android.quickstep.RemoteAnimationTargets.ReleaseCheck;
 import com.android.systemui.plugins.ResourceProvider;
 
 import java.util.ArrayList;
@@ -39,7 +40,7 @@
  * Applies spring forces to animate from a starting rect to a target rect,
  * while providing update callbacks to the caller.
  */
-public class RectFSpringAnim {
+public class RectFSpringAnim extends ReleaseCheck {
 
     private static final FloatPropertyCompat<RectFSpringAnim> RECT_CENTER_X =
             new FloatPropertyCompat<RectFSpringAnim>("rectCenterXSpring") {
@@ -116,6 +117,7 @@
         ResourceProvider rp = DynamicResource.provider(context);
         mMinVisChange = rp.getDimension(R.dimen.swipe_up_fling_min_visible_change);
         mYOvershoot = rp.getDimension(R.dimen.swipe_up_y_overshoot);
+        setCanRelease(true);
     }
 
     public void onTargetPositionChanged() {
@@ -190,10 +192,12 @@
                     maybeOnEnd();
                 });
 
+        setCanRelease(false);
+        mAnimsStarted = true;
+
         mRectXAnim.start();
         mRectYAnim.start();
         mRectScaleAnim.start();
-        mAnimsStarted = true;
         for (Animator.AnimatorListener animatorListener : mAnimatorListeners) {
             animatorListener.onAnimationStart(null);
         }
@@ -245,6 +249,7 @@
     private void maybeOnEnd() {
         if (mAnimsStarted && isEnded()) {
             mAnimsStarted = false;
+            setCanRelease(true);
             for (Animator.AnimatorListener animatorListener : mAnimatorListeners) {
                 animatorListener.onAnimationEnd(null);
             }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SurfaceTransactionApplier.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SurfaceTransactionApplier.java
new file mode 100644
index 0000000..0a3e3ec
--- /dev/null
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SurfaceTransactionApplier.java
@@ -0,0 +1,136 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * 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.android.quickstep.util;
+
+import static com.android.systemui.shared.system.TransactionCompat.deferTransactionUntil;
+import static com.android.systemui.shared.system.TransactionCompat.setEarlyWakeup;
+
+import android.annotation.TargetApi;
+import android.os.Build;
+import android.os.Handler;
+import android.os.Message;
+import android.view.SurfaceControl;
+import android.view.SurfaceControl.Transaction;
+import android.view.View;
+
+import com.android.quickstep.RemoteAnimationTargets.ReleaseCheck;
+import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
+import com.android.systemui.shared.system.ViewRootImplCompat;
+
+import java.util.function.Consumer;
+
+
+/**
+ * Helper class to apply surface transactions in sync with RenderThread similar to
+ *   android.view.SyncRtSurfaceTransactionApplier
+ * with some Launcher specific utility methods
+ */
+@TargetApi(Build.VERSION_CODES.R)
+public class SurfaceTransactionApplier extends ReleaseCheck {
+
+    private static final int MSG_UPDATE_SEQUENCE_NUMBER = 0;
+
+    private final SurfaceControl mBarrierSurfaceControl;
+    private final ViewRootImplCompat mTargetViewRootImpl;
+    private final Handler mApplyHandler;
+
+    private int mLastSequenceNumber = 0;
+
+    /**
+     * @param targetView The view in the surface that acts as synchronization anchor.
+     */
+    public SurfaceTransactionApplier(View targetView) {
+        mTargetViewRootImpl = new ViewRootImplCompat(targetView);
+        mBarrierSurfaceControl = mTargetViewRootImpl.getRenderSurfaceControl();
+        mApplyHandler = new Handler(this::onApplyMessage);
+    }
+
+    protected boolean onApplyMessage(Message msg) {
+        if (msg.what == MSG_UPDATE_SEQUENCE_NUMBER) {
+            setCanRelease(msg.arg1 == mLastSequenceNumber);
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Schedules applying surface parameters on the next frame.
+     *
+     * @param params The surface parameters to apply. DO NOT MODIFY the list after passing into
+     *               this method to avoid synchronization issues.
+     */
+    public void scheduleApply(final SurfaceParams... params) {
+        View view = mTargetViewRootImpl.getView();
+        if (view == null) {
+            return;
+        }
+
+        mLastSequenceNumber++;
+        final int toApplySeqNo = mLastSequenceNumber;
+        setCanRelease(false);
+        mTargetViewRootImpl.registerRtFrameCallback(frame -> {
+            if (mBarrierSurfaceControl == null || !mBarrierSurfaceControl.isValid()) {
+                Message.obtain(mApplyHandler, MSG_UPDATE_SEQUENCE_NUMBER, toApplySeqNo, 0)
+                        .sendToTarget();
+                return;
+            }
+            Transaction t = new Transaction();
+            for (int i = params.length - 1; i >= 0; i--) {
+                SurfaceParams surfaceParams = params[i];
+                if (surfaceParams.surface.isValid()) {
+                    deferTransactionUntil(t, surfaceParams.surface, mBarrierSurfaceControl, frame);
+                    surfaceParams.applyTo(t);
+                }
+            }
+            setEarlyWakeup(t);
+            t.apply();
+            Message.obtain(mApplyHandler, MSG_UPDATE_SEQUENCE_NUMBER, toApplySeqNo, 0)
+                    .sendToTarget();
+        });
+
+        // Make sure a frame gets scheduled.
+        view.invalidate();
+    }
+
+    /**
+     * Creates an instance of SyncRtSurfaceTransactionApplier, deferring until the target view is
+     * attached if necessary.
+     */
+    public static void create(
+            final View targetView, final Consumer<SurfaceTransactionApplier> callback) {
+        if (targetView == null) {
+            // No target view, no applier
+            callback.accept(null);
+        } else if (new ViewRootImplCompat(targetView).isValid()) {
+            // Already attached, we're good to go
+            callback.accept(new SurfaceTransactionApplier(targetView));
+        } else {
+            // Haven't been attached before we can get the view root
+            targetView.addOnAttachStateChangeListener(new View.OnAttachStateChangeListener() {
+                @Override
+                public void onViewAttachedToWindow(View v) {
+                    targetView.removeOnAttachStateChangeListener(this);
+                    callback.accept(new SurfaceTransactionApplier(targetView));
+                }
+
+                @Override
+                public void onViewDetachedFromWindow(View v) {
+                    // Do nothing
+                }
+            });
+        }
+    }
+}
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/TransformParams.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/TransformParams.java
index 9bb508e..7a62e83 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/TransformParams.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/TransformParams.java
@@ -57,7 +57,7 @@
     private float mTargetAlpha;
     private float mCornerRadius;
     private RemoteAnimationTargets mTargetSet;
-    private SyncRtSurfaceTransactionApplierCompat mSyncTransactionApplier;
+    private SurfaceTransactionApplier mSyncTransactionApplier;
 
     private BuilderProxy mHomeBuilderProxy = BuilderProxy.ALWAYS_VISIBLE;
     private BuilderProxy mBaseBuilderProxy = BuilderProxy.ALWAYS_VISIBLE;
@@ -112,7 +112,7 @@
      * are computed based on these TransformParams.
      */
     public TransformParams setSyncTransactionApplier(
-            SyncRtSurfaceTransactionApplierCompat applier) {
+            SurfaceTransactionApplier applier) {
         mSyncTransactionApplier = applier;
         return this;
     }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
index 5e90e0b..f27bedb 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
@@ -127,6 +127,7 @@
 import com.android.quickstep.util.LayoutUtils;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.quickstep.util.SplitScreenBounds;
+import com.android.quickstep.util.SurfaceTransactionApplier;
 import com.android.quickstep.util.TransformParams;
 import com.android.systemui.plugins.ResourceProvider;
 import com.android.systemui.shared.recents.IPinnedStackAnimationListener;
@@ -135,7 +136,6 @@
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.LauncherEventUtil;
 import com.android.systemui.shared.system.PackageManagerWrapper;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 
 import java.util.ArrayList;
@@ -211,7 +211,7 @@
     protected final BaseActivityInterface mSizeStrategy;
     protected RecentsAnimationController mRecentsAnimationController;
     protected RecentsAnimationTargets mRecentsAnimationTargets;
-    protected SyncRtSurfaceTransactionApplierCompat mSyncTransactionApplier;
+    protected SurfaceTransactionApplier mSyncTransactionApplier;
     protected int mTaskWidth;
     protected int mTaskHeight;
     protected boolean mEnableDrawingLiveTile = false;
@@ -502,7 +502,7 @@
         mModel.getThumbnailCache().getHighResLoadingState().addCallback(this);
         mActivity.addMultiWindowModeChangedListener(mMultiWindowModeChangedListener);
         ActivityManagerWrapper.getInstance().registerTaskStackListener(mTaskStackListener);
-        mSyncTransactionApplier = new SyncRtSurfaceTransactionApplierCompat(this);
+        mSyncTransactionApplier = new SurfaceTransactionApplier(this);
         RecentsModel.INSTANCE.get(getContext()).addThumbnailChangeListener(this);
         mIdp.addOnChangeListener(this);
         mIPinnedStackAnimationListener.setActivity(mActivity);
diff --git a/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java b/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
index 25c07f1..fe7b946 100644
--- a/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
+++ b/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
@@ -82,6 +82,7 @@
 import com.android.quickstep.util.MultiValueUpdateListener;
 import com.android.quickstep.util.RemoteAnimationProvider;
 import com.android.quickstep.util.StaggeredWorkspaceAnim;
+import com.android.quickstep.util.SurfaceTransactionApplier;
 import com.android.systemui.shared.system.ActivityCompat;
 import com.android.systemui.shared.system.ActivityOptionsCompat;
 import com.android.systemui.shared.system.QuickStepContract;
@@ -89,7 +90,6 @@
 import com.android.systemui.shared.system.RemoteAnimationDefinitionCompat;
 import com.android.systemui.shared.system.RemoteAnimationRunnerCompat;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
 import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
 import com.android.systemui.shared.system.WindowManagerWrapper;
 
@@ -455,9 +455,9 @@
 
         RemoteAnimationTargets openingTargets = new RemoteAnimationTargets(appTargets,
                 wallpaperTargets, MODE_OPENING);
-        SyncRtSurfaceTransactionApplierCompat surfaceApplier =
-                new SyncRtSurfaceTransactionApplierCompat(floatingView);
-        openingTargets.addDependentTransactionApplier(surfaceApplier);
+        SurfaceTransactionApplier surfaceApplier =
+                new SurfaceTransactionApplier(floatingView);
+        openingTargets.addReleaseCheck(surfaceApplier);
 
         // Scale the app icon to take up the entire screen. This simplifies the math when
         // animating the app window position / scale.
@@ -712,8 +712,7 @@
      */
     private Animator getUnlockWindowAnimator(RemoteAnimationTargetCompat[] appTargets,
             RemoteAnimationTargetCompat[] wallpaperTargets) {
-        SyncRtSurfaceTransactionApplierCompat surfaceApplier =
-                new SyncRtSurfaceTransactionApplierCompat(mDragLayer);
+        SurfaceTransactionApplier surfaceApplier = new SurfaceTransactionApplier(mDragLayer);
         ValueAnimator unlockAnimator = ValueAnimator.ofFloat(0, 1);
         unlockAnimator.setDuration(CLOSING_TRANSITION_DURATION_MS);
         float cornerRadius = mDeviceProfile.isMultiWindowMode ? 0 :
@@ -741,8 +740,7 @@
      */
     private Animator getClosingWindowAnimators(RemoteAnimationTargetCompat[] appTargets,
             RemoteAnimationTargetCompat[] wallpaperTargets) {
-        SyncRtSurfaceTransactionApplierCompat surfaceApplier =
-                new SyncRtSurfaceTransactionApplierCompat(mDragLayer);
+        SurfaceTransactionApplier surfaceApplier = new SurfaceTransactionApplier(mDragLayer);
         Matrix matrix = new Matrix();
         Point tmpPos = new Point();
         ValueAnimator closingAnimator = ValueAnimator.ofFloat(0, 1);
diff --git a/quickstep/src/com/android/quickstep/RemoteAnimationTargets.java b/quickstep/src/com/android/quickstep/RemoteAnimationTargets.java
index f90df45..ab5e3ba 100644
--- a/quickstep/src/com/android/quickstep/RemoteAnimationTargets.java
+++ b/quickstep/src/com/android/quickstep/RemoteAnimationTargets.java
@@ -16,19 +16,16 @@
 package com.android.quickstep;
 
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
 
-import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Queue;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 /**
  * Holds a collection of RemoteAnimationTargets, filtered by different properties.
  */
 public class RemoteAnimationTargets {
 
-    private final Queue<SyncRtSurfaceTransactionApplierCompat> mDependentTransactionAppliers =
-            new ArrayDeque<>(1);
+    private final CopyOnWriteArrayList<ReleaseCheck> mReleaseChecks = new CopyOnWriteArrayList<>();
 
     public final RemoteAnimationTargetCompat[] unfilteredApps;
     public final RemoteAnimationTargetCompat[] apps;
@@ -36,6 +33,8 @@
     public final int targetMode;
     public final boolean hasRecents;
 
+    private boolean mReleased = false;
+
     public RemoteAnimationTargets(RemoteAnimationTargetCompat[] apps,
             RemoteAnimationTargetCompat[] wallpapers, int targetMode) {
         ArrayList<RemoteAnimationTargetCompat> filteredApps = new ArrayList<>();
@@ -76,21 +75,65 @@
         return false;
     }
 
-    public void addDependentTransactionApplier(SyncRtSurfaceTransactionApplierCompat delay) {
-        mDependentTransactionAppliers.add(delay);
+    public void addReleaseCheck(ReleaseCheck check) {
+        mReleaseChecks.add(check);
     }
 
     public void release() {
-        SyncRtSurfaceTransactionApplierCompat applier = mDependentTransactionAppliers.poll();
-        if (applier == null) {
-            for (RemoteAnimationTargetCompat target : unfilteredApps) {
-                target.release();
+        if (mReleased) {
+            return;
+        }
+        for (ReleaseCheck check : mReleaseChecks) {
+            if (!check.mCanRelease) {
+                check.addOnSafeToReleaseCallback(this::release);
+                return;
             }
-            for (RemoteAnimationTargetCompat target : wallpapers) {
-                target.release();
+        }
+        mReleaseChecks.clear();
+        mReleased = true;
+
+        for (RemoteAnimationTargetCompat target : unfilteredApps) {
+            target.release();
+        }
+        for (RemoteAnimationTargetCompat target : wallpapers) {
+            target.release();
+        }
+    }
+
+    /**
+     * Interface for intercepting surface release method
+     */
+    public static class ReleaseCheck {
+
+        boolean mCanRelease = false;
+        private Runnable mAfterApplyCallback;
+
+        protected void setCanRelease(boolean canRelease) {
+            mCanRelease = canRelease;
+            if (mCanRelease && mAfterApplyCallback != null) {
+                Runnable r = mAfterApplyCallback;
+                mAfterApplyCallback = null;
+                r.run();
             }
-        } else {
-            applier.addAfterApplyCallback(this::release);
+        }
+
+        /**
+         * Adds a callback to notify when the surface can safely be released
+         */
+        void addOnSafeToReleaseCallback(Runnable callback) {
+            if (mCanRelease) {
+                callback.run();
+            } else {
+                if (mAfterApplyCallback == null) {
+                    mAfterApplyCallback = callback;
+                } else {
+                    final Runnable oldCallback = mAfterApplyCallback;
+                    mAfterApplyCallback = () -> {
+                        callback.run();
+                        oldCallback.run();
+                    };
+                }
+            }
         }
     }
 }
diff --git a/robolectric_tests/config/robolectric.properties b/robolectric_tests/config/robolectric.properties
index c162255..b171712 100644
--- a/robolectric_tests/config/robolectric.properties
+++ b/robolectric_tests/config/robolectric.properties
@@ -13,6 +13,6 @@
     com.android.launcher3.shadows.ShadowLooperExecutor \
     com.android.launcher3.shadows.ShadowMainThreadInitializedObject \
     com.android.launcher3.shadows.ShadowOverrides \
-    com.android.launcher3.shadows.ShadowSyncRtSurfaceTransactionApplierCompat \
+    com.android.launcher3.shadows.ShadowSurfaceTransactionApplier \
 
 application=com.android.launcher3.util.LauncherTestApplication
\ No newline at end of file
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/ShadowSyncRtSurfaceTransactionApplierCompat.java b/robolectric_tests/src/com/android/launcher3/shadows/ShadowSurfaceTransactionApplier.java
similarity index 83%
rename from robolectric_tests/src/com/android/launcher3/shadows/ShadowSyncRtSurfaceTransactionApplierCompat.java
rename to robolectric_tests/src/com/android/launcher3/shadows/ShadowSurfaceTransactionApplier.java
index 238926c..a9f2f27 100644
--- a/robolectric_tests/src/com/android/launcher3/shadows/ShadowSyncRtSurfaceTransactionApplierCompat.java
+++ b/robolectric_tests/src/com/android/launcher3/shadows/ShadowSurfaceTransactionApplier.java
@@ -26,11 +26,11 @@
 import org.robolectric.annotation.RealObject;
 
 /**
- * Shadow for SyncRtSurfaceTransactionApplierCompat to override default functionality
+ * Shadow for SurfaceTransactionApplier to override default functionality
  */
-@Implements(className = "com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat",
+@Implements(className = "com.android.quickstep.util.SurfaceTransactionApplier",
         isInAndroidSdk = false)
-public class ShadowSyncRtSurfaceTransactionApplierCompat {
+public class ShadowSurfaceTransactionApplier {
 
     @RealObject
     private Object mRealObject;