
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