Add mechanism to boost the shell background thread as needed

- This will be used when things offloaded to the BG thread need to
  be boosted in thread priority (ie. when loading icons along
  critical paths)

Bug: 370895085
Flag: EXEMPT bugfix
Test: atest WMShellUnitTests:HandlerExecutorTest
Change-Id: I531922a2f7b993df1ce4723b4b313bb070f728e7
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/BoostExecutor.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/common/BoostExecutor.kt
new file mode 100644
index 0000000..498d0e4
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/BoostExecutor.kt
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.common
+
+import android.os.Looper
+import java.util.concurrent.Executor
+
+/** Executor implementation which can be boosted temporarily to a different thread priority.  */
+interface BoostExecutor : Executor {
+    /**
+     * Requests that the executor is boosted until {@link #resetBoost()} is called.
+     */
+    fun setBoost() {}
+
+    /**
+     * Requests that the executor is not boosted (only resets if there are no other boost requests
+     * in progress).
+     */
+    fun resetBoost() {}
+
+    /**
+     * Returns whether the executor is boosted.
+     */
+    fun isBoosted() : Boolean {
+        return false
+    }
+
+    /**
+     * Returns the looper for this executor.
+     */
+    fun getLooper() : Looper? {
+        return Looper.myLooper()
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/HandlerExecutor.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/HandlerExecutor.java
index 736d954..803f16c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/HandlerExecutor.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/HandlerExecutor.java
@@ -16,15 +16,50 @@
 
 package com.android.wm.shell.common;
 
+import static android.os.Process.THREAD_PRIORITY_DEFAULT;
+import static android.os.Process.setThreadPriority;
+
 import android.annotation.NonNull;
 import android.os.Handler;
+import android.os.HandlerThread;
+import android.os.Looper;
+
+import androidx.annotation.VisibleForTesting;
+
+import java.util.function.BiConsumer;
 
 /** Executor implementation which is backed by a Handler. */
 public class HandlerExecutor implements ShellExecutor {
+    @NonNull
     private final Handler mHandler;
+    // See android.os.Process#THREAD_PRIORITY_*
+    private final int mDefaultThreadPriority;
+    private final int mBoostedThreadPriority;
+    // Number of current requests to boost thread priority
+    private int mBoostCount;
+    private final Object mBoostLock = new Object();
+    // Default function for setting thread priority (tid, priority)
+    private BiConsumer<Integer, Integer> mSetThreadPriorityFn =
+            HandlerExecutor::setThreadPriorityInternal;
 
     public HandlerExecutor(@NonNull Handler handler) {
+        this(handler, THREAD_PRIORITY_DEFAULT, THREAD_PRIORITY_DEFAULT);
+    }
+
+    /**
+     * Used only if this executor can be boosted, if so, it can be boosted to the given
+     * {@param boostPriority}.
+     */
+    public HandlerExecutor(@NonNull Handler handler, int defaultThreadPriority,
+            int boostedThreadPriority) {
         mHandler = handler;
+        mDefaultThreadPriority = defaultThreadPriority;
+        mBoostedThreadPriority = boostedThreadPriority;
+    }
+
+    @VisibleForTesting
+    void replaceSetThreadPriorityFn(BiConsumer<Integer, Integer> setThreadPriorityFn) {
+        mSetThreadPriorityFn = setThreadPriorityFn;
     }
 
     @Override
@@ -56,9 +91,54 @@
     }
 
     @Override
+    public void setBoost() {
+        synchronized (mBoostLock) {
+            if (mDefaultThreadPriority == mBoostedThreadPriority) {
+                // Nothing to boost
+                return;
+            }
+            if (mBoostCount == 0) {
+                mSetThreadPriorityFn.accept(
+                        ((HandlerThread) mHandler.getLooper().getThread()).getThreadId(),
+                        mBoostedThreadPriority);
+            }
+            mBoostCount++;
+        }
+    }
+
+    @Override
+    public void resetBoost() {
+        synchronized (mBoostLock) {
+            mBoostCount--;
+            if (mBoostCount == 0) {
+                mSetThreadPriorityFn.accept(
+                        ((HandlerThread) mHandler.getLooper().getThread()).getThreadId(),
+                        mDefaultThreadPriority);
+            }
+        }
+    }
+
+    @Override
+    public boolean isBoosted() {
+        synchronized (mBoostLock) {
+            return mBoostCount > 0;
+        }
+    }
+
+    @Override
+    @NonNull
+    public Looper getLooper() {
+        return mHandler.getLooper();
+    }
+
+    @Override
     public void assertCurrentThread() {
         if (!mHandler.getLooper().isCurrentThread()) {
             throw new IllegalStateException("must be called on " + mHandler);
         }
     }
+
+    private static void setThreadPriorityInternal(Integer tid, Integer priority) {
+        setThreadPriority(tid, priority);
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/ShellExecutor.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ShellExecutor.java
index 2c2961f..9e5071e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/ShellExecutor.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ShellExecutor.java
@@ -18,15 +18,15 @@
 
 import java.lang.reflect.Array;
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
 /**
  * Super basic Executor interface that adds support for delayed execution and removing callbacks.
- * Intended to wrap Handler while better-supporting testing.
+ * Intended to wrap Handler while better-supporting testing.  Not every ShellExecutor implementation
+ * may support boosting.
  */
-public interface ShellExecutor extends Executor {
+public interface ShellExecutor extends BoostExecutor {
 
     /**
      * Executes the given runnable. If the caller is running on the same looper as this executor,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellConcurrencyModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellConcurrencyModule.java
index c5644a8..d7ddbde 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellConcurrencyModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellConcurrencyModule.java
@@ -18,6 +18,7 @@
 
 import static android.os.Process.THREAD_PRIORITY_BACKGROUND;
 import static android.os.Process.THREAD_PRIORITY_DISPLAY;
+import static android.os.Process.THREAD_PRIORITY_FOREGROUND;
 import static android.os.Process.THREAD_PRIORITY_TOP_APP_BOOST;
 
 import android.content.Context;
@@ -205,13 +206,14 @@
     }
 
     /**
-     * Provides a Shell background thread Executor for low priority background tasks.
+     * Provides a Shell background thread Executor for low priority background tasks.  The thread
+     * may also be boosted to THREAD_PRIORITY_FOREGROUND if necessary.
      */
     @WMSingleton
     @Provides
     @ShellBackgroundThread
     public static ShellExecutor provideSharedBackgroundExecutor(
             @ShellBackgroundThread Handler handler) {
-        return new HandlerExecutor(handler);
+        return new HandlerExecutor(handler, THREAD_PRIORITY_BACKGROUND, THREAD_PRIORITY_FOREGROUND);
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/docs/threading.md b/libs/WindowManager/Shell/src/com/android/wm/shell/docs/threading.md
index 9d01535..837a6dd3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/docs/threading.md
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/docs/threading.md
@@ -36,7 +36,8 @@
   thread)
   - This is always another thread even if config_enableShellMainThread is not set true
   - **Note**:
-    - This thread runs with `THREAD_PRIORITY_BACKGROUND` priority
+    - This thread runs with `THREAD_PRIORITY_BACKGROUND` priority but can be requested to be boosted
+      to `THREAD_PRIORITY_FOREGROUND`
 - `ShellAnimationThread` (currently only used for Transitions and Splitscreen, but potentially all
   animations could be offloaded here)
 - `ShellSplashScreenThread` (only for use with splashscreens)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/HandlerExecutorTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/HandlerExecutorTest.kt
new file mode 100644
index 0000000..799b48c
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/HandlerExecutorTest.kt
@@ -0,0 +1,173 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.common
+
+import android.os.Handler
+import android.os.HandlerThread
+import android.os.Looper
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito
+import com.android.wm.shell.ShellTestCase
+import com.google.common.truth.Truth.assertThat
+import java.util.function.BiConsumer
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.mock
+import org.mockito.MockitoSession
+import org.mockito.kotlin.whenever
+
+/**
+ * Tests for HandlerExecutor.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:HandlerExecutorTest
+ */
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class HandlerExecutorTest : ShellTestCase() {
+
+    class TestSetThreadPriorityFn : BiConsumer<Int, Int> {
+        var lastSetPriority = UNSET_THREAD_PRIORITY
+            private set
+        var callCount = 0
+            private set
+
+        override fun accept(tid: Int, priority: Int) {
+            lastSetPriority = priority
+            callCount++
+        }
+
+        fun reset() {
+            lastSetPriority = UNSET_THREAD_PRIORITY
+            callCount = 0
+        }
+    }
+
+    val testSetPriorityFn = TestSetThreadPriorityFn()
+
+    @Test
+    fun defaultExecutorDisallowBoost() {
+        val executor = createTestHandlerExecutor()
+
+        executor.setBoost()
+
+        assertThat(executor.isBoosted()).isFalse()
+    }
+
+    @Test
+    fun boostExecutor_resetWhenNotSet_expectNoOp() {
+        val executor = createTestHandlerExecutor(DEFAULT_THREAD_PRIORITY, BOOSTED_THREAD_PRIORITY)
+        val mockSession: MockitoSession = ExtendedMockito.mockitoSession()
+            .mockStatic(android.os.Process::class.java)
+            .startMocking()
+
+        try {
+            // Try to reset and ensure we never try to set the thread priority
+            executor.resetBoost()
+
+            assertThat(testSetPriorityFn.callCount).isEqualTo(0)
+            assertThat(executor.isBoosted()).isFalse()
+        } finally {
+            mockSession.finishMocking()
+        }
+    }
+
+    @Test
+    fun boostExecutor_setResetBoost_expectThreadPriorityUpdated() {
+        val executor = createTestHandlerExecutor(DEFAULT_THREAD_PRIORITY, BOOSTED_THREAD_PRIORITY)
+        val mockSession: MockitoSession = ExtendedMockito.mockitoSession()
+            .mockStatic(android.os.Process::class.java)
+            .startMocking()
+
+        try {
+            // Boost and ensure the boosted thread priority is requested
+            executor.setBoost()
+
+            assertThat(testSetPriorityFn.lastSetPriority).isEqualTo(BOOSTED_THREAD_PRIORITY)
+            assertThat(testSetPriorityFn.callCount).isEqualTo(1)
+            assertThat(executor.isBoosted()).isTrue()
+
+            // Reset and ensure the default thread priority is requested
+            executor.resetBoost()
+
+            assertThat(testSetPriorityFn.lastSetPriority).isEqualTo(DEFAULT_THREAD_PRIORITY)
+            assertThat(testSetPriorityFn.callCount).isEqualTo(2)
+            assertThat(executor.isBoosted()).isFalse()
+        } finally {
+            mockSession.finishMocking()
+        }
+    }
+
+    @Test
+    fun boostExecutor_overlappingBoost_expectResetOnlyWhenNotOverlapping() {
+        val executor = createTestHandlerExecutor(DEFAULT_THREAD_PRIORITY, BOOSTED_THREAD_PRIORITY)
+        val mockSession: MockitoSession = ExtendedMockito.mockitoSession()
+            .mockStatic(android.os.Process::class.java)
+            .startMocking()
+
+        try {
+            // Set and ensure we only update the thread priority once
+            executor.setBoost()
+            executor.setBoost()
+
+            assertThat(testSetPriorityFn.lastSetPriority).isEqualTo(BOOSTED_THREAD_PRIORITY)
+            assertThat(testSetPriorityFn.callCount).isEqualTo(1)
+            assertThat(executor.isBoosted()).isTrue()
+
+            // Reset and ensure we are still boosted and the thread priority doesn't change
+            executor.resetBoost()
+
+            assertThat(testSetPriorityFn.lastSetPriority).isEqualTo(BOOSTED_THREAD_PRIORITY)
+            assertThat(testSetPriorityFn.callCount).isEqualTo(1)
+            assertThat(executor.isBoosted()).isTrue()
+
+            // Reset again and ensure we update the thread priority accordingly
+            executor.resetBoost()
+
+            assertThat(testSetPriorityFn.lastSetPriority).isEqualTo(DEFAULT_THREAD_PRIORITY)
+            assertThat(testSetPriorityFn.callCount).isEqualTo(2)
+            assertThat(executor.isBoosted()).isFalse()
+        } finally {
+            mockSession.finishMocking()
+        }
+    }
+
+    /**
+     * Creates a test handler executor backed by a mocked handler thread.
+     */
+    private fun createTestHandlerExecutor(
+        defaultThreadPriority: Int = DEFAULT_THREAD_PRIORITY,
+        boostedThreadPriority: Int = DEFAULT_THREAD_PRIORITY
+    ) : HandlerExecutor {
+        val handler = mock(Handler::class.java)
+        val looper = mock(Looper::class.java)
+        val thread = mock(HandlerThread::class.java)
+        whenever(handler.looper).thenReturn(looper)
+        whenever(looper.thread).thenReturn(thread)
+        whenever(thread.threadId).thenReturn(1234)
+        val executor = HandlerExecutor(handler, defaultThreadPriority, boostedThreadPriority)
+        executor.replaceSetThreadPriorityFn(testSetPriorityFn)
+        return executor
+    }
+
+    companion object {
+        private const val UNSET_THREAD_PRIORITY = 0
+        private const val DEFAULT_THREAD_PRIORITY = 1
+        private const val BOOSTED_THREAD_PRIORITY = 1000
+    }
+}
\ No newline at end of file