Introduce NotifCollectionCache

/*

 * 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.statusbar.notification.collection

import androidx.test.ext.junit.runners.AndroidJUnit4

import androidx.test.filters.SmallTest

import com.android.systemui.SysuiTestCase

import com.android.systemui.util.time.FakeSystemClock

import com.google.common.truth.Truth.assertThat

import kotlin.test.assertNotNull

import kotlin.test.assertNull

import org.junit.After

import org.junit.Test

import org.junit.runner.RunWith

@SmallTest

@RunWith(AndroidJUnit4::class)

class NotifCollectionCacheTest : SysuiTestCase() {

    companion object {

        const val A = "a"

        const val B = "b"

        const val C = "c"

    }

    val systemClock = FakeSystemClock()

    val underTest =

        NotifCollectionCache<String>(purgeTimeoutMillis = 200L, systemClock = systemClock)

    @After

    fun cleanUp() {

        underTest.clear()

    }

    @Test

    fun fetch_isOnlyCalledOncePerEntry() {

        val fetchList = mutableListOf<String>()

        val fetch = { key: String ->

            fetchList.add(key)

            key

        }

        // Construct the cache and make sure fetch is called

        assertThat(underTest.getOrFetch(A, fetch)).isEqualTo(A)

        assertThat(underTest.getOrFetch(B, fetch)).isEqualTo(B)

        assertThat(underTest.getOrFetch(C, fetch)).isEqualTo(C)

        assertThat(fetchList).containsExactly(A, B, C).inOrder()

        // Verify that further calls don't trigger fetch again

        underTest.getOrFetch(A, fetch)

        underTest.getOrFetch(A, fetch)

        underTest.getOrFetch(B, fetch)

        underTest.getOrFetch(C, fetch)

        assertThat(fetchList).containsExactly(A, B, C).inOrder()

        // Verify that fetch gets called again if the entries are cleared

        underTest.clear()

        underTest.getOrFetch(A, fetch)

        assertThat(fetchList).containsExactly(A, B, C, A).inOrder()

    }

    @Test

    fun purge_beforeTimeout_doesNothing() {

        // Populate cache

        val fetch = { key: String -> key }

        underTest.getOrFetch(A, fetch)

        underTest.getOrFetch(B, fetch)

        underTest.getOrFetch(C, fetch)

        // B starts off with ♥ ︎♥︎

        assertThat(underTest.getLives(B)).isEqualTo(2)

        // First purge run removes a ︎♥︎

        underTest.purge(listOf(A, C))

        assertNotNull(underTest.cache[B])

        assertThat(underTest.getLives(B)).isEqualTo(1)

        // Second purge run done too early does nothing to B

        systemClock.advanceTime(100L)

        underTest.purge(listOf(A, C))

        assertNotNull(underTest.cache[B])

        assertThat(underTest.getLives(B)).isEqualTo(1)

        // Purge done after timeout (200ms) clears B

        systemClock.advanceTime(100L)

        underTest.purge(listOf(A, C))

        assertNull(underTest.cache[B])

    }

    @Test

    fun get_resetsLives() {

        // Populate cache

        val fetch = { key: String -> key }

        underTest.getOrFetch(A, fetch)

        underTest.getOrFetch(B, fetch)

        underTest.getOrFetch(C, fetch)

        // Bring B down to one ︎♥︎

        underTest.purge(listOf(A, C))

        assertThat(underTest.getLives(B)).isEqualTo(1)

        // Get should restore B to ♥ ︎♥︎

        underTest.getOrFetch(B, fetch)

        assertThat(underTest.getLives(B)).isEqualTo(2)

        // Subsequent purge should remove a life regardless of timing

        underTest.purge(listOf(A, C))

        assertThat(underTest.getLives(B)).isEqualTo(1)

    }

    @Test

    fun purge_resetsLives() {

        // Populate cache

        val fetch = { key: String -> key }

        underTest.getOrFetch(A, fetch)

        underTest.getOrFetch(B, fetch)

        underTest.getOrFetch(C, fetch)

        // Bring B down to one ︎♥︎

        underTest.purge(listOf(A, C))

        assertThat(underTest.getLives(B)).isEqualTo(1)

        // When B is back to wantedKeys, it is restored to to ♥ ︎♥ ︎︎

        underTest.purge(listOf(B))

        assertThat(underTest.getLives(B)).isEqualTo(2)

        assertThat(underTest.getLives(A)).isEqualTo(1)

        assertThat(underTest.getLives(C)).isEqualTo(1)

        // Subsequent purge should remove a life regardless of timing

        underTest.purge(listOf(A, C))

        assertThat(underTest.getLives(B)).isEqualTo(1)

    }

    @Test

    fun purge_worksWithMoreLives() {

        val multiLivesCache =

            NotifCollectionCache<String>(

                retainCount = 3,

                purgeTimeoutMillis = 100L,

                systemClock = systemClock,

            )

        // Populate cache

        val fetch = { key: String -> key }

        multiLivesCache.getOrFetch(A, fetch)

        multiLivesCache.getOrFetch(B, fetch)

        multiLivesCache.getOrFetch(C, fetch)

        // B starts off with ♥ ︎♥︎ ♥ ︎♥︎

        assertThat(multiLivesCache.getLives(B)).isEqualTo(4)

        // First purge run removes a ︎♥︎

        multiLivesCache.purge(listOf(A, C))

        assertNotNull(multiLivesCache.cache[B])

        assertThat(multiLivesCache.getLives(B)).isEqualTo(3)

        // Second purge run done too early does nothing to B

        multiLivesCache.purge(listOf(A, C))

        assertNotNull(multiLivesCache.cache[B])

        assertThat(multiLivesCache.getLives(B)).isEqualTo(3)

        // Staggered purge runs remove further ︎♥︎

        systemClock.advanceTime(100L)

        multiLivesCache.purge(listOf(A, C))

        assertNotNull(multiLivesCache.cache[B])

        assertThat(multiLivesCache.getLives(B)).isEqualTo(2)

        systemClock.advanceTime(100L)

        multiLivesCache.purge(listOf(A, C))

        assertNotNull(multiLivesCache.cache[B])

        assertThat(multiLivesCache.getLives(B)).isEqualTo(1)

        systemClock.advanceTime(100L)

        multiLivesCache.purge(listOf(A, C))

        assertNull(multiLivesCache.cache[B])

    }

    @Test

    fun purge_worksWithNoLives() {

        val noLivesCache =

            NotifCollectionCache<String>(

                retainCount = 0,

                purgeTimeoutMillis = 0L,

                systemClock = systemClock,

            )

        val fetch = { key: String -> key }

        noLivesCache.getOrFetch(A, fetch)

        noLivesCache.getOrFetch(B, fetch)

        noLivesCache.getOrFetch(C, fetch)

        // Purge immediately removes entry

        noLivesCache.purge(listOf(A, C))

        assertNotNull(noLivesCache.cache[A])

        assertNull(noLivesCache.cache[B])

        assertNotNull(noLivesCache.cache[C])

    }

    private fun <V> NotifCollectionCache<V>.getLives(key: String) = this.cache[key]?.lives

}

/*

 * 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.statusbar.notification.collection

import com.android.internal.annotations.VisibleForTesting

import com.android.systemui.util.time.SystemClock

import com.android.systemui.util.time.SystemClockImpl

import java.util.concurrent.ConcurrentHashMap

/**

 * A cache in which entries can "survive" getting purged [retainCount] times, given consecutive

 * [purge] calls made at least [purgeTimeoutMillis] apart. See also [purge].

 *

 * This cache is safe for multithreaded usage, and is recommended for objects that take a while to

 * resolve (such as drawables, or things that require binder calls). As such, [getOrFetch] is

 * recommended to be run on a background thread, while [purge] can be done from any thread.

 */

class NotifCollectionCache<V>(

    private val retainCount: Int = 1,

    private val purgeTimeoutMillis: Long = 1000L,

    private val systemClock: SystemClock = SystemClockImpl(),

) {

    @get:VisibleForTesting val cache = ConcurrentHashMap<String, CacheEntry>()

    init {

        if (retainCount < 0) {

            throw IllegalArgumentException("retainCount cannot be negative")

        }

    }

    inner class CacheEntry(val key: String, val value: V) {

        /**

         * The "lives" represent how many times the entry will remain in the cache when purging it

         * is attempted.

         */

        @get:VisibleForTesting var lives: Int = retainCount + 1

        /**

         * The last time this entry lost a "life". Starts at a negative value chosen so that the

         * first purge is always considered "valid".

         */

        private var lastValidPurge: Long = -purgeTimeoutMillis

        fun resetLives() {

            // Lives/timeouts don't matter if retainCount is 0

            if (retainCount == 0) {

                return

            }

            synchronized(key) {

                lives = retainCount + 1

                lastValidPurge = -purgeTimeoutMillis

            }

            // Add it to the cache again just in case it was deleted before we could reset the lives

            cache[key] = this

        }

        fun tryPurge(): Boolean {

            // Lives/timeouts don't matter if retainCount is 0

            if (retainCount == 0) {

                return true

            }

            // Using uptimeMillis since it's guaranteed to be monotonic, as we don't want a

            // timezone/clock change to break us

            val now = systemClock.uptimeMillis()

            // Cannot purge the same entry from two threads simultaneously

            synchronized(key) {

                if (now - lastValidPurge < purgeTimeoutMillis) {

                    return false

                }

                lastValidPurge = now

                return --lives <= 0

            }

        }

    }

    /**

     * Get value from cache, or fetch it and add it to cache if not found. This can be called from

     * any thread, but is usually expected to be called from the background.

     *

     * @param key key for the object to be obtained

     * @param fetch method to fetch the object and add it to the cache if not present; note that

     *   there is no guarantee that two [fetch] cannot run in parallel for the same [key] (if

     *   [getOrFetch] is called simultaneously from different threads), so be mindful of potential

     *   side effects

     */

    fun getOrFetch(key: String, fetch: (String) -> V): V {

        val entry = cache[key]

        if (entry != null) {

            // Refresh lives on access

            entry.resetLives()

            return entry.value

        }

        val value = fetch(key)

        cache[key] = CacheEntry(key, value)

        return value

    }

    /**

     * Clear entries that are NOT in [wantedKeys] if appropriate. This can be called from any

     * thread.

     *

     * If retainCount > 0, a given entry will need to not be present in [wantedKeys] for

     * ([retainCount] + 1) consecutive [purge] calls made within at least [purgeTimeoutMillis] of

     * each other in order to be cleared. This count will be reset for any given entry 1) if

     * [getOrFetch] is called for the entry or 2) if the entry is present in [wantedKeys] in a

     * subsequent [purge] call. We prioritize keeping the entry if possible, so if [purge] is called

     * simultaneously with [getOrFetch] on different threads for example, we will try to keep it in

     * the cache, although it is not guaranteed. If avoiding cache misses is a concern, consider

     * increasing the [retainCount] or [purgeTimeoutMillis].

     *

     * For example, say [retainCount] = 1 and [purgeTimeoutMillis] = 1000 and we start with entries

     * (a, b, c) in the cache:

     * ```kotlin

     * purge((a, c)); // marks b for deletion

     * Thread.sleep(500)

     * purge((a, c)); // does nothing as it was called earlier than the min 1s

     * Thread.sleep(500)

     * purge((b, c)); // b is no longer marked for deletion, but now a is

     * Thread.sleep(1000);

     * purge((c));    // deletes a from the cache and marks b for deletion, etc.

     * ```

     */

    fun purge(wantedKeys: List<String>) {

        for ((key, entry) in cache) {

            if (key in wantedKeys) {

                entry.resetLives()

            } else if (entry.tryPurge()) {

                cache.remove(key)

            }

        }

    }

    /** Clear all entries from the cache. */

    fun clear() {

        cache.clear()

    }

}