cohort.CohortStorage.kt Maven / Gradle / Ivy

Go to download
package com.amplitude.experiment.cohort

import java.util.concurrent.locks.ReentrantReadWriteLock
import kotlin.concurrent.read
import kotlin.concurrent.write

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

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

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

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

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

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

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