Add support for custom lifecycle manager for PerDisplayRepository

/*

 * Copyright (C) 2025 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.app.displaylib

import kotlinx.coroutines.flow.MutableStateFlow

import kotlinx.coroutines.flow.StateFlow

/**

 * Reports the display ids that should have a per-display instance, if any.

 *

 * This can be overridden to support different policies (e.g. display being connected, display

 * having decorations, etc..). A [PerDisplayRepository] instance is expected to be cleaned up when a

 * displayId is removed from this set.

 */

interface DisplayInstanceLifecycleManager {

    /** Set of display ids that are allowed to have an instance. */

    val displayIds: StateFlow<Set<Int>>

}

/** Meant to be used in tests. */

class FakeDisplayInstanceLifecycleManager : DisplayInstanceLifecycleManager {

    override val displayIds = MutableStateFlow<Set<Int>>(emptySet())

}

import android.util.Log

import android.view.Display

import com.android.app.tracing.coroutines.flow.stateInTraced

import com.android.app.tracing.coroutines.launchTraced as launch

import com.android.app.tracing.traceSection

import dagger.assisted.Assisted

import java.util.concurrent.ConcurrentHashMap

import javax.inject.Qualifier

import kotlinx.coroutines.CoroutineScope

import kotlinx.coroutines.flow.SharingStarted

import kotlinx.coroutines.flow.StateFlow

import kotlinx.coroutines.flow.collectLatest

import kotlinx.coroutines.flow.combine

/**

 * Used to create instances of type `T` for a specific display.

 *

 * This class manages a cache of per-display instances of type `T`, creating them using a provided

 * [PerDisplayInstanceProvider] and optionally tearing them down using a

 * [PerDisplayInstanceProviderWithTeardown] when displays are disconnected.

 * [PerDisplayInstanceProviderWithTeardown] when based on [lifecycleManager].

 *

 * It listens to the [DisplayRepository] to detect when displays are added or removed, and

 * automatically manages the lifecycle of the per-display instances.

 * An instance will be destroyed when either

 * - The display is not connected anymore

 * - or based on [lifecycleManager]. If no lifecycle manager is provided, instances are destroyed

 *   when the display is disconnected.

 *

 * [DisplayInstanceLifecycleManager] can decide to delete instances for a display even before it is

 * disconnected. An example of usecase for it, is to delete instances when screen decorations are

 * removed.

 *

 * Note that this is a [PerDisplayStoreImpl] 2.0 that doesn't require [CoreStartable] bindings,

 * providing all args in the constructor.

constructor(

    @Assisted override val debugName: String,

    @Assisted private val instanceProvider: PerDisplayInstanceProvider<T>,

    @Assisted lifecycleManager: DisplayInstanceLifecycleManager? = null,

    @DisplayLibBackground bgApplicationScope: CoroutineScope,

    private val displayRepository: DisplayRepository,

    private val initCallback: PerDisplayRepository.InitCallback,

    private val perDisplayInstances = ConcurrentHashMap<Int, T?>()

    private val allowedDisplays: StateFlow<Set<Int>> =

        if (lifecycleManager == null) {

                displayRepository.displayIds

            } else {

                // If there is a lifecycle manager, we still consider the smallest subset between

                // the ones connected and the ones from the lifecycle. This is to safeguard against

                // leaks, in case of lifecycle manager misbehaving (as it's provided by clients, and

                // we can't guarantee it's correct).

                combine(lifecycleManager.displayIds, displayRepository.displayIds) {

                    lifecycleAllowedDisplayIds,

                    connectedDisplays ->

                    lifecycleAllowedDisplayIds.intersect(connectedDisplays)

                }

            }

            .stateInTraced(

                "allowed displays for $debugName",

                bgApplicationScope,

                SharingStarted.WhileSubscribed(),

                setOf(Display.DEFAULT_DISPLAY),

            )

    init {

        bgApplicationScope.launch("$debugName#start") { start() }

    }

    private suspend fun start() {

        initCallback.onInit(debugName, this)

        displayRepository.displayIds.collectLatest { displayIds ->

        allowedDisplays.collectLatest { displayIds ->

            val toRemove = perDisplayInstances.keys - displayIds

            toRemove.forEach { displayId ->

                Log.d(TAG, "<$debugName> destroying instance for displayId=$displayId.")

            return null

        }

        if (displayId !in allowedDisplays.value) {

            Log.e(

                TAG,

                "<$debugName: Display with id $displayId exists but it's not " +

                    "allowed by lifecycle manager.",

            )

            return null

        }

        // If it doesn't exist, create it and put it in the map.

        return perDisplayInstances.computeIfAbsent(displayId) { key ->

            Log.d(TAG, "<$debugName> creating instance for displayId=$key, as it wasn't available.")

        fun create(

            debugName: String,

            instanceProvider: PerDisplayInstanceProvider<T>,

            overrideLifecycleManager: DisplayInstanceLifecycleManager? = null,

        ): PerDisplayInstanceRepositoryImpl<T>

    }