diff --git a/services/core/java/com/android/server/wm/Dimmer.java b/services/core/java/com/android/server/wm/Dimmer.java
index ee280846ff40a..4bd8cab057000 100644
--- a/services/core/java/com/android/server/wm/Dimmer.java
+++ b/services/core/java/com/android/server/wm/Dimmer.java
@@ -130,6 +130,7 @@ class Dimmer {
             mSurfaceAnimator = new SurfaceAnimator(dimAnimatable, () -> {
                 if (!mDimming) {
                     dimAnimatable.getPendingTransaction().remove(mDimLayer);
+                    mDimLayer = null;
                 }
             }, mHost.mWmService);
         }
diff --git a/services/core/java/com/android/server/wm/SurfaceAnimator.java b/services/core/java/com/android/server/wm/SurfaceAnimator.java
index 67686a586a651..33d952e80a6b8 100644
--- a/services/core/java/com/android/server/wm/SurfaceAnimator.java
+++ b/services/core/java/com/android/server/wm/SurfaceAnimator.java
@@ -410,11 +410,13 @@ class SurfaceAnimator {
 
         /**
          * @return The surface of the object to be animated.
+         *         This SurfaceControl must be valid if non-null.
          */
         @Nullable SurfaceControl getSurfaceControl();
 
         /**
          * @return The parent of the surface object to be animated.
+         *         This SurfaceControl must be valid if non-null.
          */
         @Nullable SurfaceControl getParentSurfaceControl();
 
diff --git a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
index 4379b7cdace6b..f47e794418812 100644
--- a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
+++ b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
@@ -93,6 +93,7 @@ class TaskScreenshotAnimatable implements SurfaceAnimator.Animatable {
     @Override
     public void onAnimationLeashDestroyed(SurfaceControl.Transaction t) {
         t.remove(mSurfaceControl);
+        mSurfaceControl = null;
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 05d47609de4c3..f3f6397a62ebc 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -1003,7 +1003,10 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
         return p.makeChildSurface(child)
                 .setParent(mSurfaceControl);
     }
-
+    /*
+     * @return The SurfaceControl parent for this containers SurfaceControl.
+     *         The SurfaceControl must be valid if non-null.
+     */
     @Override
     public SurfaceControl getParentSurfaceControl() {
         final WindowContainer parent = getParent();
@@ -1210,6 +1213,10 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
         }
     }
 
+    /**
+     * @return The SurfaceControl for this container.
+     *         The SurfaceControl must be valid if non-null.
+     */
     @Override
     public SurfaceControl getSurfaceControl() {
         return mSurfaceControl;