cohort.CohortLoader.kt Maven / Gradle / Ivy

Go to download
package com.amplitude.cohort

import com.amplitude.CohortDescriptionFetch
import com.amplitude.CohortDescriptionFetchFailure
import com.amplitude.CohortDownload
import com.amplitude.CohortDownloadFailure
import com.amplitude.Metrics
import com.amplitude.util.logger
import kotlinx.coroutines.Job
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.joinAll
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

internal class CohortLoader(
    @Volatile var maxCohortSize: Int,
    private val cohortApi: CohortApi,
    private val cohortStorage: CohortStorage
) {

    companion object {
        val log by logger()
    }
    private val jobsMutex = Mutex()
    private val jobs = mutableMapOf()

    suspend fun loadCohorts(cohortIds: Set) = coroutineScope {
        val jobs = mutableListOf()
        for (cohortId in cohortIds) {
            jobs += launch { loadCohort(cohortId) }
        }
        jobs.joinAll()
    }

    private suspend fun loadCohort(cohortId: String) = coroutineScope {
        log.trace("loadCohort: start - cohortId={}", cohortId)
        val networkCohort = Metrics.with(
            { CohortDescriptionFetch },
            { e -> CohortDescriptionFetchFailure(e) }
        ) {
            cohortApi.getCohortDescription(cohortId)
        }
        val storageCohort = cohortStorage.getCohortDescription(cohortId)
        val shouldDownloadCohort = networkCohort.size <= maxCohortSize &&
            networkCohort.lastComputed > (storageCohort?.lastComputed ?: -1)
        if (shouldDownloadCohort) {
            jobsMutex.withLock {
                jobs.getOrPut(cohortId) {
                    launch {
                        log.info("Downloading cohort. $networkCohort")
                        val cohortMembers = Metrics.with({ CohortDownload }, { e -> CohortDownloadFailure(e) }) {
                            cohortApi.getCohortMembers(networkCohort)
                        }
                        cohortStorage.putCohort(networkCohort, cohortMembers)
                        jobsMutex.withLock { jobs.remove(cohortId) }
                        log.info("Cohort download complete. $networkCohort")
                    }
                }
            }.join()
        }
        log.trace("loadCohort: end - cohortId={}", cohortId)
    }
}