diff --git a/core/java/android/view/InsetsSourceControl.java b/core/java/android/view/InsetsSourceControl.java
index 85ff93b..1506ee4 100644
--- a/core/java/android/view/InsetsSourceControl.java
+++ b/core/java/android/view/InsetsSourceControl.java
@@ -48,6 +48,7 @@
     private Insets mInsetsHint;
 
     private boolean mSkipAnimationOnce;
+    private int mParcelableFlags;
 
     public InsetsSourceControl(@InternalInsetsType int type, @Nullable SurfaceControl leash,
             Point surfacePosition, Insets insetsHint) {
@@ -132,6 +133,10 @@
         return result;
     }
 
+    public void setParcelableFlags(int parcelableFlags) {
+        mParcelableFlags = parcelableFlags;
+    }
+
     @Override
     public int describeContents() {
         return 0;
@@ -140,9 +145,9 @@
     @Override
     public void writeToParcel(Parcel dest, int flags) {
         dest.writeInt(mType);
-        dest.writeTypedObject(mLeash, 0 /* parcelableFlags */);
-        dest.writeTypedObject(mSurfacePosition, 0 /* parcelableFlags */);
-        dest.writeTypedObject(mInsetsHint, 0 /* parcelableFlags */);
+        dest.writeTypedObject(mLeash, mParcelableFlags);
+        dest.writeTypedObject(mSurfacePosition, mParcelableFlags);
+        dest.writeTypedObject(mInsetsHint, mParcelableFlags);
         dest.writeBoolean(mSkipAnimationOnce);
     }
 
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index d42e0c3..a6464e7 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -5269,7 +5269,16 @@
                     // b) When loosing control, controller can restore server state by taking last
                     // dispatched state as truth.
                     mInsetsController.onStateChanged((InsetsState) args.arg1);
-                    mInsetsController.onControlsChanged((InsetsSourceControl[]) args.arg2);
+                    InsetsSourceControl[] controls = (InsetsSourceControl[]) args.arg2;
+                    if (mAdded) {
+                        mInsetsController.onControlsChanged(controls);
+                    } else if (controls != null) {
+                        for (InsetsSourceControl control : controls) {
+                            if (control != null) {
+                                control.release(SurfaceControl::release);
+                            }
+                        }
+                    }
                     args.recycle();
                     break;
                 }
@@ -8107,6 +8116,10 @@
                 }
             }
 
+            // If our window is removed, we might not get notified about losing control.
+            // Invoking this can release the leashes as soon as possible instead of relying on GC.
+            mInsetsController.onControlsChanged(null);
+
             mAdded = false;
         }
         WindowManagerGlobal.getInstance().doRemoveView(this);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
index b2ac61c..cb27ad9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
@@ -249,6 +249,7 @@
                                 !activeControl.getSurfacePosition().equals(lastSurfacePosition);
                         final boolean leashChanged =
                                 !haveSameLeash(mImeSourceControl, activeControl);
+                        final InsetsSourceControl lastImeControl = mImeSourceControl;
                         mImeSourceControl = activeControl;
                         if (mAnimation != null) {
                             if (positionChanged) {
@@ -262,6 +263,9 @@
                                 removeImeSurface();
                             }
                         }
+                        if (lastImeControl != null) {
+                            lastImeControl.release(SurfaceControl::release);
+                        }
                     }
                 }
             }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 6d51849..89bfff4 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -34,6 +34,7 @@
 import static android.content.pm.PackageManager.FEATURE_PC;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.os.InputConstants.DEFAULT_DISPATCHING_TIMEOUT_MILLIS;
+import static android.os.Parcelable.PARCELABLE_WRITE_RETURN_VALUE;
 import static android.os.Process.SYSTEM_UID;
 import static android.os.Process.myPid;
 import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
@@ -2534,12 +2535,13 @@
                 // We will leave the critical section before returning the leash to the client,
                 // so we need to copy the leash to prevent others release the one that we are
                 // about to return.
-                // TODO: We will have an extra copy if the client is not local.
-                //       For now, we rely on GC to release it.
-                //       Maybe we can modify InsetsSourceControl.writeToParcel so it can release
-                //       the extra leash as soon as possible.
-                outControls[i] = controls[i] != null
-                        ? new InsetsSourceControl(controls[i]) : null;
+                if (controls[i] != null) {
+                    // This source control is an extra copy if the client is not local. By setting
+                    // PARCELABLE_WRITE_RETURN_VALUE, the leash will be released at the end of
+                    // SurfaceControl.writeToParcel.
+                    outControls[i] = new InsetsSourceControl(controls[i]);
+                    outControls[i].setParcelableFlags(PARCELABLE_WRITE_RETURN_VALUE);
+                }
             }
         }
     }