All Downloads are FREE. Search and download functionalities are using the official Maven repository.

cohort.CohortStorage.kt Maven / Gradle / Ivy

@file:OptIn(ExperimentalApi::class)

package com.amplitude.experiment.cohort

import com.amplitude.experiment.ExperimentalApi
import com.amplitude.experiment.ProxyConfiguration
import com.amplitude.experiment.util.Cache
import com.amplitude.experiment.util.Logger
import java.util.concurrent.locks.ReentrantReadWriteLock
import kotlin.concurrent.read
import kotlin.concurrent.write

internal interface CohortStorage {
    fun getCohortsForUser(userId: String, cohortIds: Set): Set
    fun getCohortDescription(cohortId: String): CohortDescription?
    fun getCohortDescriptions(): Map
    fun putCohort(cohortDescription: CohortDescription, userIds: Set)
    fun deleteCohort(cohortId: String)
}

internal class ProxyCohortStorage(
    proxyConfig: ProxyConfiguration,
    private val cohortMembershipApi: CohortMembershipApi
) : CohortStorage {

    private val cohortCache = Cache>(
        proxyConfig.cohortCacheCapacity,
        proxyConfig.cohortCacheTtlMillis
    )
    private val inMemoryStorage = InMemoryCohortStorage()

    override fun getCohortsForUser(userId: String, cohortIds: Set): Set {
        val localCohortIds = inMemoryStorage.getCohortDescriptions().keys
        return if (localCohortIds.containsAll(cohortIds)) {
            inMemoryStorage.getCohortsForUser(userId, cohortIds)
        } else {
            cohortCache[userId]
                ?: try {
                    cohortMembershipApi.getCohortsForUser(userId).also { proxyCohortMemberships ->
                        cohortCache[userId] = proxyCohortMemberships
                    }
                } catch (e: Exception) {
                    Logger.e("Failed to get cohort membership from proxy.", e)
                    // Fall back on in memory storage in the case of proxy failure.
                    inMemoryStorage.getCohortsForUser(userId, cohortIds)
                }
        }
    }

    override fun getCohortDescription(cohortId: String): CohortDescription? {
        return inMemoryStorage.getCohortDescription(cohortId)
    }

    override fun getCohortDescriptions(): Map {
        return inMemoryStorage.getCohortDescriptions()
    }

    override fun putCohort(cohortDescription: CohortDescription, userIds: Set) {
        inMemoryStorage.putCohort(cohortDescription, userIds)
    }

    override fun deleteCohort(cohortId: String) {
        inMemoryStorage.deleteCohort(cohortId)
    }
}

internal class InMemoryCohortStorage : CohortStorage {
    private val lock = ReentrantReadWriteLock()
    private val cohortStore = mutableMapOf>()
    private val descriptionStore = mutableMapOf()

    override fun getCohortsForUser(userId: String, cohortIds: Set): Set {
        val result = mutableSetOf()
        lock.read {
            for (entry in cohortStore.entries) {
                if (cohortIds.contains(entry.key) && entry.value.contains(userId)) {
                    result.add(entry.key)
                }
            }
        }
        return result
    }

    override fun getCohortDescription(cohortId: String): CohortDescription? {
        return lock.read {
            descriptionStore[cohortId]
        }
    }

    override fun getCohortDescriptions(): Map {
        return lock.read {
            descriptionStore.toMap()
        }
    }

    override fun putCohort(cohortDescription: CohortDescription, userIds: Set) {
        lock.write {
            cohortStore[cohortDescription.id] = userIds.toMutableSet()
            descriptionStore[cohortDescription.id] = cohortDescription
        }
    }

    override fun deleteCohort(cohortId: String) {
        lock.write {
            cohortStore.remove(cohortId)
            descriptionStore.remove(cohortId)
        }
    }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy