Don't copy null surface control and remove lingering references to leashes

- If a task that should otherwise be organized but does not yet have a
  surface when the task organizer is registered, then skip trying to
  return that task as we can not leash it.  Existing logic will send the
  task to the organizer once a surface control is set on it.
- In the entry points to the shell where leashes are given from WM core
  (ie. the organizers), ensure that we release them when we are notified
  that they are no longer organized

Bug: 235858985
Test: atest ShellTaskOrganizerTests
Test: atest WindowOrganizerTests

Change-Id: I863ab057a4a566060a1efbe8a20f2dd854fa8c74
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/RootDisplayAreaOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/RootDisplayAreaOrganizer.java
index 14ba9df..764e650 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/RootDisplayAreaOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/RootDisplayAreaOrganizer.java
@@ -85,6 +85,8 @@
         }
 
         mDisplayAreasInfo.remove(displayId);
+        mLeashes.get(displayId).release();
+        mLeashes.remove(displayId);
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index e9d24fb..d5d9d99 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -529,7 +529,11 @@
             }
 
             final int taskId = taskInfo.taskId;
-            final TaskListener listener = getTaskListener(mTasks.get(taskId).getTaskInfo());
+            final TaskAppearedInfo appearedInfo = mTasks.get(taskId);
+            final TaskListener listener = getTaskListener(appearedInfo.getTaskInfo());
+            if (appearedInfo.getLeash() != null) {
+                appearedInfo.getLeash().release();
+            }
             mTasks.remove(taskId);
             if (listener != null) {
                 listener.onTaskVanished(taskInfo);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/hidedisplaycutout/HideDisplayCutoutOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/hidedisplaycutout/HideDisplayCutoutOrganizer.java
index 3f7d78d..9478b34 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/hidedisplaycutout/HideDisplayCutoutOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/hidedisplaycutout/HideDisplayCutoutOrganizer.java
@@ -128,9 +128,10 @@
 
             final WindowContainerTransaction wct = new WindowContainerTransaction();
             final SurfaceControl.Transaction t = new SurfaceControl.Transaction();
-            applyBoundsAndOffsets(
-                    displayAreaInfo.token, mDisplayAreaMap.get(displayAreaInfo.token), wct, t);
+            final SurfaceControl leash = mDisplayAreaMap.get(displayAreaInfo.token);
+            applyBoundsAndOffsets(displayAreaInfo.token, leash, wct, t);
             applyTransaction(wct, t);
+            leash.release();
             mDisplayAreaMap.remove(displayAreaInfo.token);
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
index f61d1b9..451afa0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
@@ -159,6 +159,10 @@
 
     @Override
     public void onDisplayAreaVanished(@NonNull DisplayAreaInfo displayAreaInfo) {
+        final SurfaceControl leash = mDisplayAreaTokenMap.get(displayAreaInfo.token);
+        if (leash != null) {
+            leash.release();
+        }
         mDisplayAreaTokenMap.remove(displayAreaInfo.token);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
index 22b0ccb..da88c2d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
@@ -943,6 +943,7 @@
         mPipBoundsState.setBounds(new Rect());
         mPipUiEventLoggerLogger.setTaskInfo(null);
         mMainExecutor.executeDelayed(() -> mPipMenuController.detach(), 0);
+        mLeash = null;
 
         if (info.displayId != Display.DEFAULT_DISPLAY && mOnDisplayIdChangeCallback != null) {
             mOnDisplayIdChangeCallback.accept(Display.DEFAULT_DISPLAY);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index ee49366..2799c64 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -434,7 +434,15 @@
     }
 
     RemoteAnimationTarget[] onStartingSplitLegacy(RemoteAnimationTarget[] apps) {
-        return reparentSplitTasksForAnimation(apps, false /*splitExpectedToBeVisible*/);
+        try {
+            return reparentSplitTasksForAnimation(apps, false /*splitExpectedToBeVisible*/);
+        } finally {
+            for (RemoteAnimationTarget appTarget : apps) {
+                if (appTarget.leash != null) {
+                    appTarget.leash.release();
+                }
+            }
+        }
     }
 
     private RemoteAnimationTarget[] reparentSplitTasksForAnimation(RemoteAnimationTarget[] apps,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 59b0afe..ea1f106 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -1022,6 +1022,7 @@
         }
 
         mRootTaskInfo = null;
+        mRootTaskLeash = null;
     }
 
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index 23eec96a..1e01c4b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -242,6 +242,7 @@
         if (mRootTaskInfo.taskId == taskId) {
             mCallbacks.onRootTaskVanished();
             mRootTaskInfo = null;
+            mRootLeash = null;
             mSyncQueue.runInSync(t -> {
                 t.remove(mDimLayer);
                 mSplitDecorManager.release(t);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
index 0b53c40..4948158 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
@@ -46,6 +46,7 @@
 import android.os.RemoteException;
 import android.util.SparseArray;
 import android.view.SurfaceControl;
+import android.view.SurfaceSession;
 import android.window.ITaskOrganizer;
 import android.window.ITaskOrganizerController;
 import android.window.TaskAppearedInfo;
@@ -137,13 +138,25 @@
     }
 
     @Test
-    public void registerOrganizer_sendRegisterTaskOrganizer() throws RemoteException {
+    public void testRegisterOrganizer_sendRegisterTaskOrganizer() throws RemoteException {
         mOrganizer.registerOrganizer();
 
         verify(mTaskOrganizerController).registerTaskOrganizer(any(ITaskOrganizer.class));
     }
 
     @Test
+    public void testTaskLeashReleasedAfterVanished() throws RemoteException {
+        RunningTaskInfo taskInfo = createTaskInfo(1, WINDOWING_MODE_MULTI_WINDOW);
+        SurfaceControl taskLeash = new SurfaceControl.Builder(new SurfaceSession())
+                .setName("task").build();
+        mOrganizer.registerOrganizer();
+        mOrganizer.onTaskAppeared(taskInfo, taskLeash);
+        assertTrue(taskLeash.isValid());
+        mOrganizer.onTaskVanished(taskInfo);
+        assertTrue(!taskLeash.isValid());
+    }
+
+    @Test
     public void testOneListenerPerType() {
         mOrganizer.addListenerForType(new TrackingTaskListener(), TASK_LISTENER_TYPE_MULTI_WINDOW);
         try {
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 6e84681..dc56dc9 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -382,11 +382,12 @@
                 mService.mRootWindowContainer.forAllTasks((task) -> {
                     boolean returnTask = !task.mCreatedByOrganizer;
                     task.updateTaskOrganizerState(returnTask /* skipTaskAppeared */);
-                    if (returnTask) {
-                        SurfaceControl outSurfaceControl = state.addTaskWithoutCallback(task,
+                    // It is possible for the task to not yet have a surface control, so ensure that
+                    // the update succeeded in setting the organizer for the task before returning
+                    if (task.isOrganized() && returnTask) {
+                        SurfaceControl taskLeash = state.addTaskWithoutCallback(task,
                                 "TaskOrganizerController.registerTaskOrganizer");
-                        taskInfos.add(
-                                new TaskAppearedInfo(task.getTaskInfo(), outSurfaceControl));
+                        taskInfos.add(new TaskAppearedInfo(task.getTaskInfo(), taskLeash));
                     }
                 });
             };
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index 08bad70..0f471d6 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -458,6 +458,23 @@
     }
 
     @Test
+    public void testRegisterTaskOrganizerWithExistingTasks_noSurfaceControl()
+            throws RemoteException {
+        final Task rootTask = createRootTask();
+        final Task task = createTask(rootTask);
+        final Task rootTask2 = createRootTask();
+        final Task task2 = createTask(rootTask2);
+        rootTask2.setSurfaceControl(null);
+        ArrayList<TaskAppearedInfo> existingTasks = new ArrayList<>();
+        final ITaskOrganizer organizer = registerMockOrganizer(existingTasks);
+        assertContainsTasks(existingTasks, rootTask);
+
+        // Verify we don't get onTaskAppeared if we are returned the tasks
+        verify(organizer, never())
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
+    }
+
+    @Test
     public void testTaskTransaction() {
         removeGlobalMinSizeRestriction();
         final Task rootTask = new TaskBuilder(mSupervisor)