project.ProjectProxy.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of evaluation-proxy-core Show documentation
Core package for Amplitude's evaluation proxy.
There is a newer version: 0.9.1
Show newest version
package com.amplitude.project

import com.amplitude.Configuration
import com.amplitude.EvaluationProxyResponse
import com.amplitude.assignment.Assignment
import com.amplitude.assignment.AssignmentTracker
import com.amplitude.cohort.CohortApiV1
import com.amplitude.cohort.CohortLoader
import com.amplitude.cohort.CohortStorage
import com.amplitude.cohort.GetCohortResponse
import com.amplitude.cohort.USER_GROUP_TYPE
import com.amplitude.deployment.DeploymentApiV2
import com.amplitude.deployment.DeploymentLoader
import com.amplitude.deployment.DeploymentStorage
import com.amplitude.experiment.evaluation.EvaluationEngineImpl
import com.amplitude.experiment.evaluation.EvaluationVariant
import com.amplitude.experiment.evaluation.topologicalSort
import com.amplitude.util.getGroupedCohortIds
import com.amplitude.util.json
import com.amplitude.util.logger
import com.amplitude.util.toEvaluationContext
import io.ktor.http.HttpStatusCode
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.launch
import kotlinx.serialization.encodeToString

internal class ProjectProxy(
    private val project: Project,
    configuration: Configuration,
    private val assignmentTracker: AssignmentTracker,
    private val deploymentStorage: DeploymentStorage,
    private val cohortStorage: CohortStorage,
) {
    companion object {
        val log by logger()
    }

    private val engine = EvaluationEngineImpl()

    private val projectApi = ProjectApiV1(configuration.managementServerUrl, project.managementKey)
    private val deploymentApi = DeploymentApiV2(configuration.serverUrl)
    private val cohortApi = CohortApiV1(configuration.cohortServerUrl, project.apiKey, project.secretKey)
    private val cohortLoader = CohortLoader(configuration.maxCohortSize, cohortApi, cohortStorage)
    private val deploymentLoader = DeploymentLoader(deploymentApi, deploymentStorage, cohortLoader)
    private val projectRunner =
        ProjectRunner(
            project,
            configuration,
            projectApi,
            deploymentLoader,
            deploymentStorage,
            cohortLoader,
            cohortStorage,
        )

    suspend fun start() {
        log.info("Starting project. projectId=${project.id}")
        projectRunner.start()
    }

    suspend fun shutdown() {
        log.info("Shutting down project. project.id=${project.id}")
        projectRunner.stop()
    }

    suspend fun getFlagConfigs(deploymentKey: String?): EvaluationProxyResponse {
        if (deploymentKey.isNullOrEmpty()) {
            return EvaluationProxyResponse.error(HttpStatusCode.Unauthorized, "Invalid deployment")
        }
        val result = deploymentStorage.getAllFlags(deploymentKey).values.toList()
        return EvaluationProxyResponse.error(HttpStatusCode.OK, json.encodeToString(result))
    }

    suspend fun getCohort(
        cohortId: String?,
        lastModified: Long?,
        maxCohortSize: Int?,
    ): EvaluationProxyResponse {
        if (cohortId.isNullOrEmpty()) {
            return EvaluationProxyResponse.error(HttpStatusCode.NotFound, "Cohort not found")
        }
        val cohortDescription =
            cohortStorage.getCohortDescription(cohortId)
                ?: return EvaluationProxyResponse.error(HttpStatusCode.NotFound, "Cohort not found")
        if (cohortDescription.size > (maxCohortSize ?: Int.MAX_VALUE)) {
            return EvaluationProxyResponse.error(
                HttpStatusCode.PayloadTooLarge,
                "Cohort $cohortId sized ${cohortDescription.size} is greater than max cohort size $maxCohortSize",
            )
        }
        if (cohortDescription.lastModified == lastModified) {
            return EvaluationProxyResponse.error(HttpStatusCode.NoContent, "Cohort not modified")
        }
        val cohort =
            cohortStorage.getCohort(cohortId)
                ?: return EvaluationProxyResponse.error(HttpStatusCode.NotFound, "Cohort members not found")
        return EvaluationProxyResponse.json(HttpStatusCode.OK, GetCohortResponse.fromCohort(cohort))
    }

    suspend fun getCohortMemberships(
        deploymentKey: String?,
        groupType: String?,
        groupName: String?,
    ): EvaluationProxyResponse {
        if (deploymentKey.isNullOrEmpty()) {
            return EvaluationProxyResponse.error(HttpStatusCode.Unauthorized, "Invalid deployment")
        }
        if (groupType.isNullOrEmpty()) {
            return EvaluationProxyResponse.error(HttpStatusCode.BadRequest, "Invalid group type")
        }
        if (groupName.isNullOrEmpty()) {
            return EvaluationProxyResponse.error(HttpStatusCode.BadRequest, "Invalid group name")
        }
        val cohortIds = deploymentStorage.getAllFlags(deploymentKey).values.getGroupedCohortIds()[groupType]
        if (cohortIds.isNullOrEmpty()) {
            return EvaluationProxyResponse.json(HttpStatusCode.OK, emptySet())
        }
        val result = cohortStorage.getCohortMemberships(groupType, groupName, cohortIds)
        return EvaluationProxyResponse.json(HttpStatusCode.OK, result)
    }

    suspend fun evaluate(
        deploymentKey: String?,
        user: Map?,
        flagKeys: Set? = null,
    ): EvaluationProxyResponse {
        if (deploymentKey.isNullOrEmpty()) {
            return EvaluationProxyResponse.error(HttpStatusCode.Unauthorized, "Invalid deployment")
        }
        val result = evaluateInternal(deploymentKey, user, flagKeys)
        return EvaluationProxyResponse(HttpStatusCode.OK, json.encodeToString(result))
    }

    suspend fun evaluateV1(
        deploymentKey: String?,
        user: Map?,
        flagKeys: Set? = null,
    ): EvaluationProxyResponse {
        if (deploymentKey.isNullOrEmpty()) {
            return EvaluationProxyResponse(HttpStatusCode.Unauthorized, "Invalid deployment")
        }
        val result =
            evaluateInternal(deploymentKey, user, flagKeys).filter { entry ->
                val default = entry.value.metadata?.get("default") as? Boolean ?: false
                val deployed = entry.value.metadata?.get("deployed") as? Boolean ?: true
                (!default && deployed)
            }
        return EvaluationProxyResponse(HttpStatusCode.OK, json.encodeToString(result))
    }

    private suspend fun evaluateInternal(
        deploymentKey: String,
        user: Map?,
        flagKeys: Set? = null,
    ): Map {
        // Get flag configs for the deployment from storage and topo sort.
        val storageFlags = deploymentStorage.getAllFlags(deploymentKey)
        if (storageFlags.isEmpty()) {
            return mapOf()
        }
        val flags = topologicalSort(storageFlags, flagKeys ?: setOf())
        if (flags.isEmpty()) {
            return mapOf()
        }
        val groupedCohortIds = flags.getGroupedCohortIds()
        // Enrich user with cohort IDs and build the evaluation context
        val userId = user?.get("user_id") as? String
        val enrichedUser = user?.toMutableMap() ?: mutableMapOf()
        val userCohortIds = groupedCohortIds[USER_GROUP_TYPE]
        if (userId != null && userCohortIds != null) {
            enrichedUser["cohort_ids"] = cohortStorage.getCohortMemberships(USER_GROUP_TYPE, userId, userCohortIds)
        }
        val groups = enrichedUser["groups"] as? Map<*, *>
        if (!groups.isNullOrEmpty()) {
            val groupCohortIds = mutableMapOf>>()
            for (entry in groups.entries) {
                val groupType = entry.key as? String
                val groupName = (entry.value as? Collection<*>)?.firstOrNull() as? String
                val groupTypeCohortIds = groupedCohortIds[groupType]
                if (groupType != null && groupName != null && groupTypeCohortIds != null) {
                    val cohortIds = cohortStorage.getCohortMemberships(groupType, groupName, groupTypeCohortIds)
                    if (groupCohortIds.isNotEmpty()) {
                        groupCohortIds.putIfAbsent(groupType, mutableMapOf(groupName to cohortIds))
                    }
                }
            }
            if (groupCohortIds.isNotEmpty()) {
                enrichedUser["group_cohort_ids"] = groupCohortIds
            }
        }
        val evaluationContext = enrichedUser.toEvaluationContext()
        // Evaluate results
        log.debug("evaluate - context={}", evaluationContext)
        val result = engine.evaluate(evaluationContext, flags)
        if (enrichedUser.isNotEmpty()) {
            coroutineScope {
                launch {
                    assignmentTracker.track(Assignment(evaluationContext, result))
                }
            }
        }
        return result
    }

    // Internal

    internal suspend fun getDeployments(): Set {
        return deploymentStorage.getDeployments().keys
    }
}