Make WakeLock handling thread safe

We have an implementation of a wakelock which reference counts requestors - however, the implementation is not thread safe.

This commit:

1. Makes the implementation thread safe
2. Retroactively flag guards the use of WakeLocks on background thread from previous commit.
3. Slightly improves logging so we don't churn strings if logging is disabled.

Bug:326544949
Bug:316128516

Test: newly added and existing unit tests
Flag: ACONFIG com.android.systemui.flags.DELAYED_WAKELOCK_RELEASE_ON_BACKGROUND_THREAD DISABLED
Change-Id: I9ce86fdd69fdae9f88b5e06f5b85836fb4f9aa5e
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index f5c4843..ba77380 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -542,6 +542,13 @@
     namespace: "systemui"
     description: "Binds Keyguard Media Controller Visibility to MediaContainerView"
     bug: "298213983"
+}
+
+flag {
+    name: "delayed_wakelock_release_on_background_thread"
+    namespace: "systemui"
+    description: "Released delayed wakelocks on background threads to avoid janking screen transitions."
+    bug: "316128516"
     metadata {
         purpose: PURPOSE_BUGFIX
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/util/wakelock/ClientTrackingWakeLockTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/wakelock/ClientTrackingWakeLockTest.kt
new file mode 100644
index 0000000..fdfcdc4
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/wakelock/ClientTrackingWakeLockTest.kt
@@ -0,0 +1,130 @@
+/*
+ * 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.systemui.util.wakelock
+
+import android.os.Build
+import android.os.PowerManager
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import org.junit.After
+import org.junit.Assert
+import org.junit.Assume
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ClientTrackingWakeLockTest : SysuiTestCase() {
+
+    private val WHY = "test"
+    private val WHY_2 = "test2"
+
+    lateinit var mWakeLock: ClientTrackingWakeLock
+    lateinit var mInner: PowerManager.WakeLock
+
+    @Before
+    fun setUp() {
+        mInner =
+            WakeLock.createWakeLockInner(mContext, "WakeLockTest", PowerManager.PARTIAL_WAKE_LOCK)
+        mWakeLock = ClientTrackingWakeLock(mInner, null, 20000)
+    }
+
+    @After
+    fun tearDown() {
+        mInner.setReferenceCounted(false)
+        mInner.release()
+    }
+
+    @Test
+    fun createPartialInner_notHeldYet() {
+        Assert.assertFalse(mInner.isHeld)
+    }
+
+    @Test
+    fun wakeLock_acquire() {
+        mWakeLock.acquire(WHY)
+        Assert.assertTrue(mInner.isHeld)
+    }
+
+    @Test
+    fun wakeLock_release() {
+        mWakeLock.acquire(WHY)
+        mWakeLock.release(WHY)
+        Assert.assertFalse(mInner.isHeld)
+    }
+
+    @Test
+    fun wakeLock_acquiredReleasedMultipleSources_stillHeld() {
+        mWakeLock.acquire(WHY)
+        mWakeLock.acquire(WHY_2)
+        mWakeLock.release(WHY)
+
+        Assert.assertTrue(mInner.isHeld)
+        mWakeLock.release(WHY_2)
+        Assert.assertFalse(mInner.isHeld)
+    }
+
+    @Test
+    fun wakeLock_releasedTooManyTimes_stillReleased_noThrow() {
+        Assume.assumeFalse(Build.IS_ENG)
+        mWakeLock.acquire(WHY)
+        mWakeLock.acquire(WHY_2)
+        mWakeLock.release(WHY)
+        mWakeLock.release(WHY_2)
+        mWakeLock.release(WHY)
+        Assert.assertFalse(mInner.isHeld)
+    }
+
+    @Test
+    fun wakeLock_wrap() {
+        val ran = BooleanArray(1)
+        val wrapped = mWakeLock.wrap { ran[0] = true }
+        Assert.assertTrue(mInner.isHeld)
+        Assert.assertFalse(ran[0])
+        wrapped.run()
+        Assert.assertTrue(ran[0])
+        Assert.assertFalse(mInner.isHeld)
+    }
+
+    @Test
+    fun prodBuild_wakeLock_releaseWithoutAcquire_noThrow() {
+        Assume.assumeFalse(Build.IS_ENG)
+        // shouldn't throw an exception on production builds
+        mWakeLock.release(WHY)
+    }
+
+    @Test
+    fun acquireSeveralLocks_stringReportsCorrectCount() {
+        mWakeLock.acquire(WHY)
+        mWakeLock.acquire(WHY_2)
+        mWakeLock.acquire(WHY)
+        mWakeLock.acquire(WHY)
+        mWakeLock.acquire(WHY_2)
+        Assert.assertEquals(5, mWakeLock.activeClients())
+
+        mWakeLock.release(WHY_2)
+        mWakeLock.release(WHY_2)
+        Assert.assertEquals(3, mWakeLock.activeClients())
+
+        mWakeLock.release(WHY)
+        mWakeLock.release(WHY)
+        mWakeLock.release(WHY)
+        Assert.assertEquals(0, mWakeLock.activeClients())
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/util/wakelock/ClientTrackingWakeLock.kt b/packages/SystemUI/src/com/android/systemui/util/wakelock/ClientTrackingWakeLock.kt
new file mode 100644
index 0000000..db300eb
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/util/wakelock/ClientTrackingWakeLock.kt
@@ -0,0 +1,63 @@
+/*
+ * 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.systemui.util.wakelock
+
+import android.os.PowerManager
+import android.util.Log
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicInteger
+
+/**
+ * [PowerManager.WakeLock] wrapper that tracks acquire/release reasons and logs them if owning
+ * logger is enabled.
+ */
+class ClientTrackingWakeLock(
+    private val pmWakeLock: PowerManager.WakeLock,
+    private val logger: WakeLockLogger?,
+    private val maxTimeout: Long
+) : WakeLock {
+
+    private val activeClients = ConcurrentHashMap<String, AtomicInteger>()
+
+    override fun acquire(why: String) {
+        val count = activeClients.computeIfAbsent(why) { _ -> AtomicInteger(0) }.incrementAndGet()
+        logger?.logAcquire(pmWakeLock, why, count)
+        pmWakeLock.acquire(maxTimeout)
+    }
+
+    override fun release(why: String) {
+        val count = activeClients[why]?.decrementAndGet() ?: -1
+        if (count < 0) {
+            Log.wtf(WakeLock.TAG, "Releasing WakeLock with invalid reason: $why")
+            // Restore count just in case.
+            activeClients[why]?.incrementAndGet()
+            return
+        }
+
+        logger?.logRelease(pmWakeLock, why, count)
+        pmWakeLock.release()
+    }
+
+    override fun wrap(r: Runnable): Runnable = WakeLock.wrapImpl(this, r)
+
+    fun activeClients(): Int =
+        activeClients.reduceValuesToInt(Long.MAX_VALUE, AtomicInteger::get, 0, Integer::sum)
+
+    override fun toString(): String {
+        return "active clients=${activeClients()}"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/util/wakelock/DelayedWakeLock.java b/packages/SystemUI/src/com/android/systemui/util/wakelock/DelayedWakeLock.java
index 039109e..d2ed71c 100644
--- a/packages/SystemUI/src/com/android/systemui/util/wakelock/DelayedWakeLock.java
+++ b/packages/SystemUI/src/com/android/systemui/util/wakelock/DelayedWakeLock.java
@@ -19,8 +19,11 @@
 import android.content.Context;
 import android.os.Handler;
 
+import com.android.systemui.Flags;
 import com.android.systemui.dagger.qualifiers.Background;
+import com.android.systemui.dagger.qualifiers.Main;
 
+import dagger.Lazy;
 import dagger.assisted.Assisted;
 import dagger.assisted.AssistedFactory;
 import dagger.assisted.AssistedInject;
@@ -37,10 +40,13 @@
     private final WakeLock mInner;
 
     @AssistedInject
-    public DelayedWakeLock(@Background Handler handler, Context context, WakeLockLogger logger,
+    public DelayedWakeLock(@Background Lazy<Handler> bgHandler,
+                           @Main Lazy<Handler> mainHandler,
+                           Context context, WakeLockLogger logger,
             @Assisted String tag) {
         mInner = WakeLock.createPartial(context, logger, tag);
-        mHandler = handler;
+        mHandler = Flags.delayedWakelockReleaseOnBackgroundThread() ? bgHandler.get()
+                : mainHandler.get();
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java b/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java
index 6128fee..707751a 100644
--- a/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java
+++ b/packages/SystemUI/src/com/android/systemui/util/wakelock/WakeLock.java
@@ -22,6 +22,8 @@
 
 import androidx.annotation.VisibleForTesting;
 
+import com.android.systemui.Flags;
+
 import java.util.HashMap;
 
 import javax.inject.Inject;
@@ -112,6 +114,11 @@
     @VisibleForTesting
     static WakeLock wrap(
             final PowerManager.WakeLock inner, WakeLockLogger logger, long maxTimeout) {
+        if (Flags.delayedWakelockReleaseOnBackgroundThread()) {
+            return new ClientTrackingWakeLock(inner, logger, maxTimeout);
+        }
+
+        // Non-thread safe implementation, remove when flag above is removed.
         return new WakeLock() {
             private final HashMap<String, Integer> mActiveClients = new HashMap<>();
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java b/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java
index 23a9207..ed07ad2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/wakelock/WakeLockTest.java
@@ -21,21 +21,40 @@
 
 import android.os.Build;
 import android.os.PowerManager;
+import android.platform.test.flag.junit.FlagsParameterization;
+import android.platform.test.flag.junit.SetFlagsRule;
 
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
+import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 
 import org.junit.After;
+import org.junit.Assume;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import java.util.List;
 
 @SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(Parameterized.class)
 public class WakeLockTest extends SysuiTestCase {
 
+    @Parameterized.Parameters(name = "{0}")
+    public static List<FlagsParameterization> getFlags() {
+        return FlagsParameterization.allCombinationsOf(
+                Flags.FLAG_DELAYED_WAKELOCK_RELEASE_ON_BACKGROUND_THREAD);
+    }
+
+    @Rule public final SetFlagsRule mSetFlagsRule;
+
+    public WakeLockTest(FlagsParameterization flags) {
+        mSetFlagsRule = new SetFlagsRule(SetFlagsRule.DefaultInitValueType.NULL_DEFAULT, flags);
+    }
+
     private static final String WHY = "test";
     WakeLock mWakeLock;
     PowerManager.WakeLock mInner;
@@ -91,10 +110,7 @@
 
     @Test
     public void prodBuild_wakeLock_releaseWithoutAcquire_noThrow() {
-        if (Build.IS_ENG) {
-            return;
-        }
-
+        Assume.assumeFalse(Build.IS_ENG);
         // shouldn't throw an exception on production builds
         mWakeLock.release(WHY);
     }