Register remote animation per Task

Before, we register remote animation per organizer, and assume it only
organizes in one Task. Now, we register per organizer Task if there is a
rule that supports split in the Task width.

Bug: 207720388
Test: atest WMJetpackUnitTests
Change-Id: Iac4e10af9aa0e7846c81b036b2cf708b6cbe627e

core/java/android/window/ITaskFragmentOrganizerController.aidl

@@ -35,15 +35,15 @@ interface ITaskFragmentOrganizerController {
     /**
      * Registers remote animations per transition type for the organizer. It will override the
      * animations if the transition only contains windows that belong to the organized
-     * TaskFragments.
+     * TaskFragments in the given Task.
      */
-    void registerRemoteAnimations(in ITaskFragmentOrganizer organizer,
+    void registerRemoteAnimations(in ITaskFragmentOrganizer organizer, int taskId,
         in RemoteAnimationDefinition definition);
 
     /**
      * Unregisters remote animations per transition type for the organizer.
      */
-    void unregisterRemoteAnimations(in ITaskFragmentOrganizer organizer);
+    void unregisterRemoteAnimations(in ITaskFragmentOrganizer organizer, int taskId);
 
     /**
       * Checks if an activity organized by a {@link android.window.TaskFragmentOrganizer} and

core/java/android/window/TaskFragmentOrganizer.java

@@ -94,13 +94,16 @@ public class TaskFragmentOrganizer extends WindowOrganizer {
     /**
      * Registers remote animations per transition type for the organizer. It will override the
      * animations if the transition only contains windows that belong to the organized
-     * TaskFragments.
+     * TaskFragments in the given Task.
+     *
+     * @param taskId overrides if the transition only contains windows belonging to this Task.
      * @hide
      */
     @CallSuper
-    public void registerRemoteAnimations(@NonNull RemoteAnimationDefinition definition) {
+    public void registerRemoteAnimations(int taskId,
+            @NonNull RemoteAnimationDefinition definition) {
         try {
-            getController().registerRemoteAnimations(mInterface, definition);
+            getController().registerRemoteAnimations(mInterface, taskId, definition);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
@@ -111,9 +114,9 @@ public class TaskFragmentOrganizer extends WindowOrganizer {
      * @hide
      */
     @CallSuper
-    public void unregisterRemoteAnimations() {
+    public void unregisterRemoteAnimations(int taskId) {
         try {
-            getController().unregisterRemoteAnimations(mInterface);
+            getController().unregisterRemoteAnimations(mInterface, taskId);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }

libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java

@@ -35,6 +35,8 @@ import android.window.WindowContainerTransaction;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.util.Map;
 import java.util.concurrent.Executor;
 
@@ -56,7 +58,8 @@ class JetpackTaskFragmentOrganizer extends TaskFragmentOrganizer {
     final Map<IBinder, Configuration> mFragmentParentConfigs = new ArrayMap<>();
 
     private final TaskFragmentCallback mCallback;
-    private TaskFragmentAnimationController mAnimationController;
+    @VisibleForTesting
+    TaskFragmentAnimationController mAnimationController;
 
     /**
      * Callback that notifies the controller about changes to task fragments.
@@ -80,21 +83,25 @@ class JetpackTaskFragmentOrganizer extends TaskFragmentOrganizer {
 
     @Override
     public void unregisterOrganizer() {
-        stopOverrideSplitAnimation();
-        mAnimationController = null;
+        if (mAnimationController != null) {
+            mAnimationController.unregisterAllRemoteAnimations();
+            mAnimationController = null;
+        }
         super.unregisterOrganizer();
     }
 
-    void startOverrideSplitAnimation() {
+    /** Overrides the animation if the transition is on the given Task. */
+    void startOverrideSplitAnimation(int taskId) {
         if (mAnimationController == null) {
             mAnimationController = new TaskFragmentAnimationController(this);
         }
-        mAnimationController.registerRemoteAnimations();
+        mAnimationController.registerRemoteAnimations(taskId);
     }
 
-    void stopOverrideSplitAnimation() {
+    /** No longer overrides the animation if the transition is on the given Task. */
+    void stopOverrideSplitAnimation(int taskId) {
         if (mAnimationController != null) {
-            mAnimationController.unregisterRemoteAnimations();
+            mAnimationController.unregisterRemoteAnimations(taskId);
         }
     }
 

libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java

@@ -37,6 +37,7 @@ import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
+import android.util.ArraySet;
 import android.util.SparseArray;
 import android.window.TaskFragmentInfo;
 import android.window.WindowContainerTransaction;
@@ -75,10 +76,6 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
     private Consumer<List<SplitInfo>> mEmbeddingCallback;
     private final List<SplitInfo> mLastReportedSplitStates = new ArrayList<>();
 
-    // We currently only support split activity embedding within the one root Task.
-    // TODO(b/207720388): move to TaskContainer
-    private final Rect mParentBounds = new Rect();
-
     public SplitController() {
         mPresenter = new SplitPresenter(new MainThreadExecutor(), this);
         ActivityThread activityThread = ActivityThread.currentActivityThread();
@@ -95,7 +92,9 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
     public void setEmbeddingRules(@NonNull Set<EmbeddingRule> rules) {
         mSplitRules.clear();
         mSplitRules.addAll(rules);
-        updateAnimationOverride();
+        for (int i = mTaskContainers.size() - 1; i >= 0; i--) {
+            updateAnimationOverride(mTaskContainers.keyAt(i));
+        }
     }
 
     @NonNull
@@ -163,38 +162,49 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
 
     @Override
     public void onTaskFragmentVanished(@NonNull TaskFragmentInfo taskFragmentInfo) {
-        TaskFragmentContainer container = getContainer(taskFragmentInfo.getFragmentToken());
-        if (container == null) {
-            return;
+        final TaskFragmentContainer container = getContainer(taskFragmentInfo.getFragmentToken());
+        if (container != null) {
+            // Cleanup if the TaskFragment vanished is not requested by the organizer.
+            mPresenter.cleanupContainer(container, true /* shouldFinishDependent */);
+            updateCallbackIfNecessary();
         }
-
-        mPresenter.cleanupContainer(container, true /* shouldFinishDependent */);
-        updateCallbackIfNecessary();
+        cleanupTaskFragment(taskFragmentInfo.getFragmentToken());
     }
 
     @Override
     public void onTaskFragmentParentInfoChanged(@NonNull IBinder fragmentToken,
             @NonNull Configuration parentConfig) {
-        onParentBoundsMayChange(parentConfig.windowConfiguration.getBounds());
         TaskFragmentContainer container = getContainer(fragmentToken);
         if (container != null) {
+            onTaskBoundsMayChange(container.getTaskId(),
+                    parentConfig.windowConfiguration.getBounds());
             mPresenter.updateContainer(container);
             updateCallbackIfNecessary();
         }
     }
 
-    private void onParentBoundsMayChange(Activity activity) {
-        if (activity.isFinishing()) {
+    /** Called on receiving {@link #onTaskFragmentVanished(TaskFragmentInfo)} for cleanup. */
+    private void cleanupTaskFragment(@NonNull IBinder taskFragmentToken) {
+        for (int i = mTaskContainers.size() - 1; i >= 0; i--) {
+            final TaskContainer taskContainer = mTaskContainers.valueAt(i);
+            if (!taskContainer.mFinishedContainer.remove(taskFragmentToken)) {
+                continue;
+            }
+            if (taskContainer.isEmpty()) {
+                // Cleanup the TaskContainer if it becomes empty.
+                mPresenter.stopOverrideSplitAnimation(taskContainer.mTaskId);
+                mTaskContainers.remove(taskContainer.mTaskId);
+            }
             return;
         }
-
-        onParentBoundsMayChange(mPresenter.getParentContainerBounds(activity));
     }
 
-    private void onParentBoundsMayChange(Rect parentBounds) {
-        if (!parentBounds.isEmpty() && !mParentBounds.equals(parentBounds)) {
-            mParentBounds.set(parentBounds);
-            updateAnimationOverride();
+    private void onTaskBoundsMayChange(int taskId, @NonNull Rect taskBounds) {
+        final TaskContainer taskContainer = mTaskContainers.get(taskId);
+        if (taskContainer != null && !taskBounds.isEmpty()
+                && !taskContainer.mTaskBounds.equals(taskBounds)) {
+            taskContainer.mTaskBounds.set(taskBounds);
+            updateAnimationOverride(taskId);
         }
     }
 
@@ -202,9 +212,10 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
      * Updates if we should override transition animation. We only want to override if the Task
      * bounds is large enough for at least one split rule.
      */
-    private void updateAnimationOverride() {
-        if (mParentBounds.isEmpty()) {
-            // We don't know about the parent bounds yet.
+    private void updateAnimationOverride(int taskId) {
+        final TaskContainer taskContainer = mTaskContainers.get(taskId);
+        if (taskContainer == null || !taskContainer.isTaskBoundsInitialized()) {
+            // We don't know about the Task bounds yet.
             return;
         }
 
@@ -214,7 +225,7 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
             if (!(rule instanceof SplitRule)) {
                 continue;
             }
-            if (mPresenter.shouldShowSideBySide(mParentBounds, (SplitRule) rule)) {
+            if (mPresenter.shouldShowSideBySide(taskContainer.mTaskBounds, (SplitRule) rule)) {
                 supportSplit = true;
                 break;
             }
@@ -222,9 +233,9 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
 
         // We only want to override if it supports split.
         if (supportSplit) {
-            mPresenter.startOverrideSplitAnimation();
+            mPresenter.startOverrideSplitAnimation(taskId);
         } else {
-            mPresenter.stopOverrideSplitAnimation();
+            mPresenter.stopOverrideSplitAnimation(taskId);
         }
     }
 
@@ -243,11 +254,6 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
         final TaskFragmentContainer currentContainer = getContainerWithActivity(
                 launchedActivity.getActivityToken());
 
-        if (currentContainer == null) {
-            // Initial check before any TaskFragment is created.
-            onParentBoundsMayChange(launchedActivity);
-        }
-
         // Check if the activity is configured to always be expanded.
         if (shouldExpand(launchedActivity, null, splitRules)) {
             if (shouldContainerBeExpanded(currentContainer)) {
@@ -326,8 +332,6 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
             // onTaskFragmentParentInfoChanged
             return;
         }
-        // The bounds of the container may have been changed.
-        onParentBoundsMayChange(activity);
 
         // Check if activity requires a placeholder
         launchPlaceholderIfNecessary(activity);
@@ -357,9 +361,14 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
     TaskFragmentContainer newContainer(@Nullable Activity activity, int taskId) {
         final TaskFragmentContainer container = new TaskFragmentContainer(activity, taskId);
         if (!mTaskContainers.contains(taskId)) {
-            mTaskContainers.put(taskId, new TaskContainer());
+            mTaskContainers.put(taskId, new TaskContainer(taskId));
+        }
+        final TaskContainer taskContainer = mTaskContainers.get(taskId);
+        taskContainer.mContainers.add(container);
+        if (activity != null && !taskContainer.isTaskBoundsInitialized()) {
+            // Initial check before any TaskFragment has appeared.
+            onTaskBoundsMayChange(taskId, SplitPresenter.getTaskBoundsFromActivity(activity));
         }
-        mTaskContainers.get(taskId).mContainers.add(container);
         return container;
     }
 
@@ -391,11 +400,11 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
             return;
         }
         taskContainer.mContainers.remove(container);
-        if (taskContainer.mContainers.isEmpty()) {
-            mTaskContainers.remove(taskId);
-            // No more TaskFragment in this Task, so no need to check split container.
-            return;
-        }
+        // Marked as a pending removal which will be removed after it is actually removed on the
+        // server side (#onTaskFragmentVanished).
+        // In this way, we can keep track of the Task bounds until we no longer have any
+        // TaskFragment there.
+        taskContainer.mFinishedContainer.add(container.getTaskFragmentToken());
 
         final List<SplitContainer> containersToRemove = new ArrayList<>();
         for (SplitContainer splitContainer : taskContainer.mSplitContainers) {
@@ -543,7 +552,7 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
     }
 
     boolean launchPlaceholderIfNecessary(@NonNull Activity activity) {
-        final  TaskFragmentContainer container = getContainerWithActivity(
+        final TaskFragmentContainer container = getContainerWithActivity(
                 activity.getActivityToken());
         // Don't launch placeholder if the container is occluded.
         if (container != null && container != getTopActiveContainer(container.getTaskId())) {
@@ -1035,7 +1044,30 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
     /** Represents TaskFragments and split pairs below a Task. */
     @VisibleForTesting
     static class TaskContainer {
+        /** The unique task id. */
+        final int mTaskId;
+        /** Active TaskFragments in this Task. */
         final List<TaskFragmentContainer> mContainers = new ArrayList<>();
+        /** Active split pairs in this Task. */
         final List<SplitContainer> mSplitContainers = new ArrayList<>();
+        /**
+         * TaskFragments that the organizer has requested to be closed. They should be removed when
+         * the organizer receives {@link #onTaskFragmentVanished(TaskFragmentInfo)} event for them.
+         */
+        final Set<IBinder> mFinishedContainer = new ArraySet<>();
+        /** Available window bounds of this Task. */
+        final Rect mTaskBounds = new Rect();
+
+        TaskContainer(int taskId) {
+            mTaskId = taskId;
+        }
+
+        boolean isEmpty() {
+            return mContainers.isEmpty() && mFinishedContainer.isEmpty();
+        }
+
+        boolean isTaskBoundsInitialized() {
+            return !mTaskBounds.isEmpty();
+        }
     }
 }

libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java

@@ -419,7 +419,11 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
             }
         }
 
-        // TODO(b/190433398): Check if the client-side available info about parent bounds is enough.
+        return getTaskBoundsFromActivity(activity);
+    }
+
+    @NonNull
+    static Rect getTaskBoundsFromActivity(@NonNull Activity activity) {
         if (!activity.isInMultiWindowMode()) {
             // In fullscreen mode the max bounds should correspond to the task bounds.
             return activity.getResources().getConfiguration().windowConfiguration.getMaxBounds();

libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationController.java

@@ -24,11 +24,14 @@ import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_CLOSE;
 import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_OPEN;
 import static android.view.WindowManager.TRANSIT_OLD_TASK_OPEN;
 
+import android.util.ArraySet;
 import android.util.Log;
 import android.view.RemoteAnimationAdapter;
 import android.view.RemoteAnimationDefinition;
 import android.window.TaskFragmentOrganizer;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 /** Controls the TaskFragment remote animations. */
 class TaskFragmentAnimationController {
 
@@ -37,8 +40,10 @@ class TaskFragmentAnimationController {
 
     private final TaskFragmentOrganizer mOrganizer;
     private final TaskFragmentAnimationRunner mRemoteRunner = new TaskFragmentAnimationRunner();
-    private final RemoteAnimationDefinition mDefinition;
-    private boolean mIsRegister;
+    @VisibleForTesting
+    final RemoteAnimationDefinition mDefinition;
+    /** Task Ids that we have registered for remote animation. */
+    private final ArraySet<Integer> mRegisterTasks = new ArraySet<>();
 
     TaskFragmentAnimationController(TaskFragmentOrganizer organizer) {
         mOrganizer = organizer;
@@ -54,25 +59,32 @@ class TaskFragmentAnimationController {
         mDefinition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_CHANGE, animationAdapter);
     }
 
-    void registerRemoteAnimations() {
+    void registerRemoteAnimations(int taskId) {
         if (DEBUG) {
             Log.v(TAG, "registerRemoteAnimations");
         }
-        if (mIsRegister) {
+        if (mRegisterTasks.contains(taskId)) {
             return;
         }
-        mOrganizer.registerRemoteAnimations(mDefinition);
-        mIsRegister = true;
+        mOrganizer.registerRemoteAnimations(taskId, mDefinition);
+        mRegisterTasks.add(taskId);
     }
 
-    void unregisterRemoteAnimations() {
+    void unregisterRemoteAnimations(int taskId) {
         if (DEBUG) {
             Log.v(TAG, "unregisterRemoteAnimations");
         }
-        if (!mIsRegister) {
+        if (!mRegisterTasks.contains(taskId)) {
             return;
         }
-        mOrganizer.unregisterRemoteAnimations();
-        mIsRegister = false;
+        mOrganizer.unregisterRemoteAnimations(taskId);
+        mRegisterTasks.remove(taskId);
+    }
+
+    void unregisterAllRemoteAnimations() {
+        final ArraySet<Integer> tasks = new ArraySet<>(mRegisterTasks);
+        for (int taskId : tasks) {
+            unregisterRemoteAnimations(taskId);
+        }
     }
 }

libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2022 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 androidx.window.extensions.embedding;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.never;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Test class for {@link JetpackTaskFragmentOrganizer}.
+ *
+ * Build/Install/Run:
+ *  atest WMJetpackUnitTests:JetpackTaskFragmentOrganizerTest
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class JetpackTaskFragmentOrganizerTest {
+    private static final int TASK_ID = 10;
+
+    @Mock
+    private JetpackTaskFragmentOrganizer.TaskFragmentCallback mCallback;
+    private JetpackTaskFragmentOrganizer mOrganizer;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        mOrganizer = new JetpackTaskFragmentOrganizer(Runnable::run, mCallback);
+        mOrganizer.registerOrganizer();
+        spyOn(mOrganizer);
+    }
+
+    @Test
+    public void testUnregisterOrganizer() {
+        mOrganizer.startOverrideSplitAnimation(TASK_ID);
+        mOrganizer.startOverrideSplitAnimation(TASK_ID + 1);
+        mOrganizer.unregisterOrganizer();
+
+        verify(mOrganizer).unregisterRemoteAnimations(TASK_ID);
+        verify(mOrganizer).unregisterRemoteAnimations(TASK_ID + 1);
+    }
+
+    @Test
+    public void testStartOverrideSplitAnimation() {
+        assertNull(mOrganizer.mAnimationController);
+
+        mOrganizer.startOverrideSplitAnimation(TASK_ID);
+
+        assertNotNull(mOrganizer.mAnimationController);
+        verify(mOrganizer).registerRemoteAnimations(TASK_ID,
+                mOrganizer.mAnimationController.mDefinition);
+    }
+
+    @Test
+    public void testStopOverrideSplitAnimation() {
+        mOrganizer.stopOverrideSplitAnimation(TASK_ID);
+
+        verify(mOrganizer, never()).unregisterRemoteAnimations(anyInt());
+
+        mOrganizer.startOverrideSplitAnimation(TASK_ID);
+        mOrganizer.stopOverrideSplitAnimation(TASK_ID);
+
+        verify(mOrganizer).unregisterRemoteAnimations(TASK_ID);
+    }
+}

libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java

@@ -31,6 +31,12 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+/**
+ * Test class for {@link SplitController}.
+ *
+ * Build/Install/Run:
+ *  atest WMJetpackUnitTests:SplitController
+ */
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class SplitControllerTest {
@@ -46,7 +52,7 @@ public class SplitControllerTest {
 
     @Test
     public void testGetTopActiveContainer() {
-        TaskContainer taskContainer = new TaskContainer();
+        TaskContainer taskContainer = new TaskContainer(TASK_ID);
         // tf3 is finished so is not active.
         TaskFragmentContainer tf3 = mock(TaskFragmentContainer.class);
         doReturn(true).when(tf3).isFinished();

libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentAnimationControllerTest.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2022 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 androidx.window.extensions.embedding;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.never;
+
+import android.window.TaskFragmentOrganizer;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Test class for {@link TaskFragmentAnimationController}.
+ *
+ * Build/Install/Run:
+ *  atest WMJetpackUnitTests:TaskFragmentAnimationControllerTest
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class TaskFragmentAnimationControllerTest {
+    private static final int TASK_ID = 10;
+
+    @Mock
+    private TaskFragmentOrganizer mOrganizer;
+    private TaskFragmentAnimationController mAnimationController;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        mAnimationController = new TaskFragmentAnimationController(mOrganizer);
+    }
+
+    @Test
+    public void testRegisterRemoteAnimations() {
+        mAnimationController.registerRemoteAnimations(TASK_ID);
+
+        verify(mOrganizer).registerRemoteAnimations(TASK_ID, mAnimationController.mDefinition);
+
+        mAnimationController.registerRemoteAnimations(TASK_ID);
+
+        // No extra call if it has been registered.
+        verify(mOrganizer).registerRemoteAnimations(TASK_ID, mAnimationController.mDefinition);
+    }
+
+    @Test
+    public void testUnregisterRemoteAnimations() {
+        mAnimationController.unregisterRemoteAnimations(TASK_ID);
+
+        // No call if it is not registered.
+        verify(mOrganizer, never()).unregisterRemoteAnimations(anyInt());
+
+        mAnimationController.registerRemoteAnimations(TASK_ID);
+        mAnimationController.unregisterRemoteAnimations(TASK_ID);
+
+        verify(mOrganizer).unregisterRemoteAnimations(TASK_ID);
+
+        mAnimationController.unregisterRemoteAnimations(TASK_ID);
+
+        // No extra call if it has been unregistered.
+        verify(mOrganizer).unregisterRemoteAnimations(TASK_ID);
+    }
+
+    @Test
+    public void testUnregisterAllRemoteAnimations() {
+        mAnimationController.registerRemoteAnimations(TASK_ID);
+        mAnimationController.registerRemoteAnimations(TASK_ID + 1);
+        mAnimationController.unregisterAllRemoteAnimations();
+
+        verify(mOrganizer).unregisterRemoteAnimations(TASK_ID);
+        verify(mOrganizer).unregisterRemoteAnimations(TASK_ID + 1);
+    }
+}

services/core/java/com/android/server/wm/AppTransitionController.java

@@ -653,7 +653,7 @@ public class AppTransitionController {
         final ITaskFragmentOrganizer organizer = findTaskFragmentOrganizer(task);
         final RemoteAnimationDefinition definition = organizer != null
                 ? mDisplayContent.mAtmService.mTaskFragmentOrganizerController
-                    .getRemoteAnimationDefinition(organizer)
+                    .getRemoteAnimationDefinition(organizer, task.mTaskId)
                 : null;
         final RemoteAnimationAdapter adapter = definition != null
                 ? definition.getAdapter(transit, activityTypes)

services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java

@@ -32,6 +32,7 @@ import android.os.IBinder;
 import android.os.RemoteException;
 import android.util.ArrayMap;
 import android.util.Slog;
+import android.util.SparseArray;
 import android.view.RemoteAnimationDefinition;
 import android.window.ITaskFragmentOrganizer;
 import android.window.ITaskFragmentOrganizerController;
@@ -83,11 +84,12 @@ public class TaskFragmentOrganizerController extends ITaskFragmentOrganizerContr
                 new WeakHashMap<>();
 
         /**
-         * @see android.window.TaskFragmentOrganizer#registerRemoteAnimations(
-         * RemoteAnimationDefinition)
+         * Map from Task Id to {@link RemoteAnimationDefinition}.
+         * @see android.window.TaskFragmentOrganizer#registerRemoteAnimations(int,
+         * RemoteAnimationDefinition) )
          */
-        @Nullable
-        private RemoteAnimationDefinition mRemoteAnimationDefinition;
+        private final SparseArray<RemoteAnimationDefinition> mRemoteAnimationDefinitions =
+                new SparseArray<>();
 
         TaskFragmentOrganizerState(ITaskFragmentOrganizer organizer) {
             mOrganizer = organizer;
@@ -251,7 +253,7 @@ public class TaskFragmentOrganizerController extends ITaskFragmentOrganizerContr
     }
 
     @Override
-    public void registerRemoteAnimations(ITaskFragmentOrganizer organizer,
+    public void registerRemoteAnimations(ITaskFragmentOrganizer organizer, int taskId,
             RemoteAnimationDefinition definition) {
         final int pid = Binder.getCallingPid();
         final int uid = Binder.getCallingUid();
@@ -264,19 +266,20 @@ public class TaskFragmentOrganizerController extends ITaskFragmentOrganizerContr
             if (organizerState == null) {
                 throw new IllegalStateException("The organizer hasn't been registered.");
             }
-            if (organizerState.mRemoteAnimationDefinition != null) {
+            if (organizerState.mRemoteAnimationDefinitions.contains(taskId)) {
                 throw new IllegalStateException(
                         "The organizer has already registered remote animations="
-                                + organizerState.mRemoteAnimationDefinition);
+                                + organizerState.mRemoteAnimationDefinitions.get(taskId)
+                                + " for TaskId=" + taskId);
             }
 
             definition.setCallingPidUid(pid, uid);
-            organizerState.mRemoteAnimationDefinition = definition;
+            organizerState.mRemoteAnimationDefinitions.put(taskId, definition);
         }
     }
 
     @Override
-    public void unregisterRemoteAnimations(ITaskFragmentOrganizer organizer) {
+    public void unregisterRemoteAnimations(ITaskFragmentOrganizer organizer, int taskId) {
         final int pid = Binder.getCallingPid();
         final long uid = Binder.getCallingUid();
         synchronized (mGlobalLock) {
@@ -290,7 +293,7 @@ public class TaskFragmentOrganizerController extends ITaskFragmentOrganizerContr
                 return;
             }
 
-            organizerState.mRemoteAnimationDefinition = null;
+            organizerState.mRemoteAnimationDefinitions.remove(taskId);
         }
     }
 
@@ -300,11 +303,13 @@ public class TaskFragmentOrganizerController extends ITaskFragmentOrganizerContr
      */
     @Nullable
     public RemoteAnimationDefinition getRemoteAnimationDefinition(
-            ITaskFragmentOrganizer organizer) {
+            ITaskFragmentOrganizer organizer, int taskId) {
         synchronized (mGlobalLock) {
             final TaskFragmentOrganizerState organizerState =
                     mTaskFragmentOrganizerState.get(organizer.asBinder());
-            return organizerState != null ? organizerState.mRemoteAnimationDefinition : null;
+            return organizerState != null
+                    ? organizerState.mRemoteAnimationDefinitions.get(taskId)
+                    : null;
         }
     }
 

services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java

@@ -861,13 +861,13 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testOverrideTaskFragmentAdapter_overrideWithEmbeddedActivity() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
         // Create a TaskFragment with embedded activity.
-        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(
-                createTask(mDisplayContent), organizer);
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
         final ActivityRecord activity = taskFragment.getTopMostActivity();
         prepareActivityForAppTransition(activity);
         spyOn(mDisplayContent.mAppTransition);
@@ -882,11 +882,11 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testOverrideTaskFragmentAdapter_overrideWithNonEmbeddedActivity() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
-        final Task task = createTask(mDisplayContent);
         // Closing non-embedded activity.
         final ActivityRecord closingActivity = createActivityRecord(task);
         prepareActivityForAppTransition(closingActivity);
@@ -907,11 +907,11 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testOverrideTaskFragmentAdapter_overrideEmbeddedActivityWithDiffUid() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
-        final Task task = createTask(mDisplayContent);
         // Closing TaskFragment with embedded activity.
         final TaskFragment taskFragment1 = createTaskFragmentWithEmbeddedActivity(task, organizer);
         final ActivityRecord closingActivity = taskFragment1.getTopMostActivity();
@@ -934,16 +934,16 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testOverrideTaskFragmentAdapter_noOverrideWithTwoApps() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
         // Closing activity in Task1.
         final ActivityRecord closingActivity = createActivityRecord(mDisplayContent);
         prepareActivityForAppTransition(closingActivity);
         // Opening TaskFragment with embedded activity in Task2.
-        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(
-                createTask(mDisplayContent), organizer);
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
         final ActivityRecord openingActivity = taskFragment.getTopMostActivity();
         prepareActivityForAppTransition(openingActivity);
         spyOn(mDisplayContent.mAppTransition);
@@ -958,11 +958,11 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testOverrideTaskFragmentAdapter_noOverrideNonEmbeddedActivityWithDiffUid() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
-        final Task task = createTask(mDisplayContent);
         // Closing TaskFragment with embedded activity.
         final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
         final ActivityRecord closingActivity = taskFragment.getTopMostActivity();
@@ -986,13 +986,13 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testOverrideTaskFragmentAdapter_noOverrideWithWallpaper() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
         // Create a TaskFragment with embedded activity.
-        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(
-                createTask(mDisplayContent), organizer);
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
         final ActivityRecord activity = taskFragment.getTopMostActivity();
         prepareActivityForAppTransition(activity);
         // Set wallpaper as visible.
@@ -1012,13 +1012,13 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
     @Test
     public void testOverrideTaskFragmentAdapter_inputProtectedForUntrustedAnimation() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
         // Create a TaskFragment with embedded activities, one is trusted embedded, and the other
         // one is untrusted embedded.
-        final Task task = createTask(mDisplayContent);
         final TaskFragment taskFragment = new TaskFragmentBuilder(mAtm)
                 .setParentTask(task)
                 .createActivityCount(2)
@@ -1071,12 +1071,12 @@ public class AppTransitionControllerTest extends WindowTestsBase {
      */
     @Test
     public void testOverrideTaskFragmentAdapter_inputProtectedForTrustedAnimation() {
+        final Task task = createTask(mDisplayContent);
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
-        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+        setupTaskFragmentRemoteAnimation(organizer, task.mTaskId, remoteAnimationRunner);
 
         // Create a TaskFragment with only trusted embedded activity
-        final Task task = createTask(mDisplayContent);
         final TaskFragment taskFragment = new TaskFragmentBuilder(mAtm)
                 .setParentTask(task)
                 .createActivityCount(1)
@@ -1170,7 +1170,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
     }
 
     /** Registers remote animation for the organizer. */
-    private void setupTaskFragmentRemoteAnimation(TaskFragmentOrganizer organizer,
+    private void setupTaskFragmentRemoteAnimation(TaskFragmentOrganizer organizer, int taskId,
             TestRemoteAnimationRunner remoteAnimationRunner) {
         final RemoteAnimationAdapter adapter = new RemoteAnimationAdapter(
                 remoteAnimationRunner, 10, 1);
@@ -1181,7 +1181,8 @@ public class AppTransitionControllerTest extends WindowTestsBase {
         definition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_OPEN, adapter);
         definition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_CLOSE, adapter);
         mAtm.mTaskFragmentOrganizerController.registerOrganizer(iOrganizer);
-        mAtm.mTaskFragmentOrganizerController.registerRemoteAnimations(iOrganizer, definition);
+        mAtm.mTaskFragmentOrganizerController.registerRemoteAnimations(iOrganizer, taskId,
+                definition);
     }
 
     private void prepareAndTriggerAppTransition(@Nullable ActivityRecord openingActivity,

services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java

@@ -67,6 +67,7 @@ import org.junit.runner.RunWith;
 @Presubmit
 @RunWith(WindowTestRunner.class)
 public class TaskFragmentOrganizerControllerTest extends WindowTestsBase {
+    private static final int TASK_ID = 10;
 
     private TaskFragmentOrganizerController mController;
     private TaskFragmentOrganizer mOrganizer;
@@ -219,13 +220,13 @@ public class TaskFragmentOrganizerControllerTest extends WindowTestsBase {
     @Test
     public void testRegisterRemoteAnimations() {
         mController.registerOrganizer(mIOrganizer);
-        mController.registerRemoteAnimations(mIOrganizer, mDefinition);
+        mController.registerRemoteAnimations(mIOrganizer, TASK_ID, mDefinition);
 
-        assertEquals(mDefinition, mController.getRemoteAnimationDefinition(mIOrganizer));
+        assertEquals(mDefinition, mController.getRemoteAnimationDefinition(mIOrganizer, TASK_ID));
 
-        mController.unregisterRemoteAnimations(mIOrganizer);
+        mController.unregisterRemoteAnimations(mIOrganizer, TASK_ID);
 
-        assertNull(mController.getRemoteAnimationDefinition(mIOrganizer));
+        assertNull(mController.getRemoteAnimationDefinition(mIOrganizer, TASK_ID));
     }
 
     @Test