fix: update config

src/main/kotlin/LocalEvaluationClient.kt

@@ -43,7 +43,7 @@ class LocalEvaluationClient internal constructor(
     private val metrics: LocalEvaluationMetrics = LocalEvaluationMetricsWrapper(config.metrics)
     private val flagConfigApi = DynamicFlagConfigApi(apiKey, serverUrl, getProxyUrl(config), httpClient)
     private val flagConfigStorage = InMemoryFlagConfigStorage()
-    private val cohortStorage = if (config.cohortSyncConfiguration != null) {
+    private val cohortStorage = if (config.cohortSyncConfig != null) {
         InMemoryCohortStorage()
     } else {
         null
@@ -145,11 +145,11 @@ class LocalEvaluationClient internal constructor(
 }
 
 private fun getCohortDownloadApi(config: LocalEvaluationConfig, httpClient: OkHttpClient): CohortApi? {
-    return if (config.cohortSyncConfiguration != null) {
+    return if (config.cohortSyncConfig != null) {
         DynamicCohortApi(
-            apiKey = config.cohortSyncConfiguration.apiKey,
-            secretKey = config.cohortSyncConfiguration.secretKey,
-            maxCohortSize = config.cohortSyncConfiguration.maxCohortSize,
+            apiKey = config.cohortSyncConfig.apiKey,
+            secretKey = config.cohortSyncConfig.secretKey,
+            maxCohortSize = config.cohortSyncConfig.maxCohortSize,
             serverUrl = getCohortServerUrl(config),
             proxyUrl = getProxyUrl(config),
             httpClient = httpClient,
@@ -171,12 +171,13 @@ private fun getServerUrl(config: LocalEvaluationConfig): HttpUrl {
 }
 
 private fun getProxyUrl(config: LocalEvaluationConfig): HttpUrl? {
-    return config.evaluationProxyConfiguration?.proxyUrl?.toHttpUrl()
+    return config.evaluationProxyConfig?.proxyUrl?.toHttpUrl()
 }
 
 private fun getCohortServerUrl(config: LocalEvaluationConfig): HttpUrl {
     return if (config.serverZone == LocalEvaluationConfig.Defaults.SERVER_ZONE) {
-        config.cohortServerUrl.toHttpUrl()
+        config.cohortSyncConfig?.cohortServerUrl?.toHttpUrl()
+            ?: US_COHORT_SERVER_URL.toHttpUrl()
     } else {
         when (config.serverZone) {
             ServerZone.US -> US_COHORT_SERVER_URL.toHttpUrl()

src/main/kotlin/LocalEvaluationConfig.kt

@@ -1,23 +1,21 @@
-@file:OptIn(ExperimentalApi::class)
-
 package com.amplitude.experiment
 
 import com.amplitude.Middleware
+import com.amplitude.experiment.LocalEvaluationConfig.Defaults
 
 /**
  * Configuration options. This is an immutable object that can be created using
  * a [LocalEvaluationConfig.Builder]. Example usage:
  *
  * `LocalEvaluationConfig.builder().serverUrl("https://api.lab.amplitude.com/").build()`
  */
+@OptIn(ExperimentalApi::class)
 class LocalEvaluationConfig internal constructor(
     @JvmField
     val debug: Boolean = Defaults.DEBUG,
     @JvmField
     val serverUrl: String = Defaults.SERVER_URL,
     @JvmField
-    val cohortServerUrl: String = Defaults.COHORT_SERVER_URL,
-    @JvmField
     val serverZone: ServerZone = Defaults.SERVER_ZONE,
     @JvmField
     val flagConfigPollerIntervalMillis: Long = Defaults.FLAG_CONFIG_POLLER_INTERVAL_MILLIS,
@@ -26,9 +24,9 @@ class LocalEvaluationConfig internal constructor(
     @JvmField
     val assignmentConfiguration: AssignmentConfiguration? = Defaults.ASSIGNMENT_CONFIGURATION,
     @JvmField
-    val cohortSyncConfiguration: CohortSyncConfiguration? = Defaults.COHORT_SYNC_CONFIGURATION,
+    val cohortSyncConfig: CohortSyncConfig? = Defaults.COHORT_SYNC_CONFIGURATION,
     @JvmField
-    val evaluationProxyConfiguration: EvaluationProxyConfiguration? = Defaults.EVALUATION_PROXY_CONFIGURATION,
+    val evaluationProxyConfig: EvaluationProxyConfig? = Defaults.EVALUATION_PROXY_CONFIGURATION,
     @JvmField
     val metrics: LocalEvaluationMetrics? = Defaults.LOCAL_EVALUATION_METRICS,
 ) {
@@ -81,12 +79,12 @@ class LocalEvaluationConfig internal constructor(
         /**
          * null
          */
-        val COHORT_SYNC_CONFIGURATION: CohortSyncConfiguration? = null
+        val COHORT_SYNC_CONFIGURATION: CohortSyncConfig? = null
 
         /**
          * null
          */
-        val EVALUATION_PROXY_CONFIGURATION: EvaluationProxyConfiguration? = null
+        val EVALUATION_PROXY_CONFIGURATION: EvaluationProxyConfig? = null
 
         /**
          * null
@@ -132,12 +130,13 @@ class LocalEvaluationConfig internal constructor(
             this.assignmentConfiguration = assignmentConfiguration
         }
 
-        fun enableCohortSync(cohortSyncConfiguration: CohortSyncConfiguration) = apply {
-            this.cohortSyncConfiguration = cohortSyncConfiguration
+        fun cohortSyncConfig(cohortSyncConfig: CohortSyncConfig) = apply {
+            this.cohortSyncConfiguration = cohortSyncConfig
         }
 
-        fun enableEvaluationProxy(evaluationProxyConfiguration: EvaluationProxyConfiguration) = apply {
-            this.evaluationProxyConfiguration = evaluationProxyConfiguration
+        @ExperimentalApi
+        fun evaluationProxyConfig(evaluationProxyConfig: EvaluationProxyConfig) = apply {
+            this.evaluationProxyConfiguration = evaluationProxyConfig
         }
 
         @ExperimentalApi
@@ -152,8 +151,8 @@ class LocalEvaluationConfig internal constructor(
                 flagConfigPollerIntervalMillis = flagConfigPollerIntervalMillis,
                 flagConfigPollerRequestTimeoutMillis = flagConfigPollerRequestTimeoutMillis,
                 assignmentConfiguration = assignmentConfiguration,
-                cohortSyncConfiguration = cohortSyncConfiguration,
-                evaluationProxyConfiguration = evaluationProxyConfiguration,
+                cohortSyncConfig = cohortSyncConfiguration,
+                evaluationProxyConfig = evaluationProxyConfiguration,
                 metrics = metrics,
             )
         }
@@ -164,8 +163,8 @@ class LocalEvaluationConfig internal constructor(
             "flagConfigPollerIntervalMillis=$flagConfigPollerIntervalMillis, " +
             "flagConfigPollerRequestTimeoutMillis=$flagConfigPollerRequestTimeoutMillis, " +
             "assignmentConfiguration=$assignmentConfiguration, " +
-            "cohortSyncConfiguration=$cohortSyncConfiguration, " +
-            "evaluationProxyConfiguration=$evaluationProxyConfiguration, " +
+            "cohortSyncConfiguration=$cohortSyncConfig, " +
+            "evaluationProxyConfiguration=$evaluationProxyConfig, " +
             "metrics=$metrics)"
     }
 }
@@ -180,15 +179,17 @@ data class AssignmentConfiguration @JvmOverloads constructor(
     val middleware: List<Middleware> = listOf(),
 )
 
-
-data class CohortSyncConfiguration @JvmOverloads constructor(
+@ExperimentalApi
+data class CohortSyncConfig @JvmOverloads constructor(
     val apiKey: String,
     val secretKey: String,
     val maxCohortSize: Int = Int.MAX_VALUE,
+    val cohortServerUrl: String = Defaults.COHORT_SERVER_URL,
+    val cohortPollingIntervalMillis: Long = 60000L
 )
 
 @ExperimentalApi
-data class EvaluationProxyConfiguration @JvmOverloads constructor(
+data class EvaluationProxyConfig @JvmOverloads constructor(
     val proxyUrl: String,
     val cohortCacheCapacity: Int = 1000000,
     val cohortCacheTtlMillis: Long = 60000L,

src/main/kotlin/assignment/AssignmentService.kt

@@ -3,7 +3,6 @@ package com.amplitude.experiment.assignment
 import com.amplitude.Amplitude
 import com.amplitude.AmplitudeCallbacks
 import com.amplitude.Event
-import com.amplitude.experiment.ExperimentalApi
 import com.amplitude.experiment.LocalEvaluationMetrics
 import com.amplitude.experiment.util.LocalEvaluationMetricsWrapper
 import com.amplitude.experiment.util.wrapMetrics

src/main/kotlin/cohort/CohortApi.kt

@@ -55,8 +55,8 @@ internal class DynamicCohortApi(
 
     private val token = Base64.getEncoder().encodeToString("$apiKey:$secretKey".toByteArray())
     private val backoffConfig = BackoffConfig(
-        attempts = 3,
-        min = 500,
+        attempts = 5,
+        min = 100,
         max = 2000,
         scalar = 2.0,
     )
@@ -115,9 +115,8 @@ internal class DynamicCohortApi(
         })
         try {
             return future.get().toCohort()
-        } catch(e: ExecutionException) {
+        } catch (e: ExecutionException) {
             throw e.cause ?: e
         }
     }
-
 }

src/main/kotlin/cohort/CohortStorage.kt

@@ -2,7 +2,7 @@
 
 package com.amplitude.experiment.cohort
 
-import com.amplitude.experiment.EvaluationProxyConfiguration
+import com.amplitude.experiment.EvaluationProxyConfig
 import com.amplitude.experiment.ExperimentalApi
 import com.amplitude.experiment.LocalEvaluationMetrics
 import com.amplitude.experiment.util.Cache
@@ -26,7 +26,7 @@ internal interface CohortStorage {
 }
 
 internal class ProxyCohortStorage(
-    private val proxyConfig: EvaluationProxyConfiguration,
+    private val proxyConfig: EvaluationProxyConfig,
     private val membershipApi: CohortMembershipApi,
     private val metrics: LocalEvaluationMetrics = LocalEvaluationMetricsWrapper()
 ) : CohortStorage {
@@ -126,7 +126,6 @@ internal class InMemoryCohortStorage : CohortStorage {
         return result
     }
 
-
     override fun putCohort(cohort: Cohort) {
         lock.write {
             cohortStore[cohort.id] = cohort

src/main/kotlin/deployment/DeploymentRunner.kt

@@ -1,5 +1,8 @@
+@file:OptIn(ExperimentalApi::class)
+
 package com.amplitude.experiment.deployment
 
+import com.amplitude.experiment.ExperimentalApi
 import com.amplitude.experiment.LocalEvaluationConfig
 import com.amplitude.experiment.LocalEvaluationMetrics
 import com.amplitude.experiment.cohort.CohortApi
@@ -33,6 +36,7 @@ internal class DeploymentRunner(
     } else {
         null
     }
+    private val cohortPollingInterval: Long = getCohortPollingInterval()
 
     fun start() = lock.once {
         refresh()
@@ -59,10 +63,9 @@ internal class DeploymentRunner(
                     } catch (t: Throwable) {
                         Logger.e("Refresh cohorts failed.", t)
                     }
-                },
-                60,
-                60,
-                TimeUnit.SECONDS
+                }, cohortPollingInterval,
+                cohortPollingInterval,
+                TimeUnit.MILLISECONDS
             )
         }
     }
@@ -122,9 +125,16 @@ internal class DeploymentRunner(
             val storageCohortIds = cohortStorage.getCohorts().keys
             val deletedCohortIds = storageCohortIds - flagCohortIds
             for (deletedCohortId in deletedCohortIds) {
-                    cohortStorage.deleteCohort(deletedCohortId)
+                cohortStorage.deleteCohort(deletedCohortId)
             }
         }
         Logger.d("Refreshed ${flagConfigs.size} flag configs.")
     }
+
+    private fun getCohortPollingInterval(): Long {
+        if (config.cohortSyncConfig == null || config.cohortSyncConfig.cohortPollingIntervalMillis < 60000) {
+            return 60000
+        }
+        return config.cohortSyncConfig.cohortPollingIntervalMillis
+    }
 }

src/main/kotlin/flag/FlagConfigApi.kt

@@ -55,7 +55,7 @@ internal class DynamicFlagConfigApi(
         }
         try {
             return future.get()
-        } catch(e: ExecutionException) {
+        } catch (e: ExecutionException) {
             throw e.cause ?: e
         }
     }

src/main/kotlin/util/Backoff.kt

@@ -64,8 +64,8 @@ private class Backoff<T>(
                     if (shouldRetry && nextAttempt < config.attempts) {
                         val nextDelay = min(delay * config.scalar, config.max.toDouble()).toLong()
                         val jitter = Random.nextLong(
-                            (nextDelay - (nextDelay*0.1).toLong()) * -1,
-                            nextDelay + (nextDelay+0.1).toLong()
+                            (nextDelay - (nextDelay * 0.1).toLong()) * -1,
+                            nextDelay + (nextDelay + 0.1).toLong()
                         )
                         backoff(nextAttempt, nextDelay + jitter, function, retry)
                     } else {

src/test/kotlin/LocalEvaluationClientTest.kt

@@ -2,8 +2,6 @@ package com.amplitude.experiment
 
 import com.amplitude.experiment.cohort.Cohort
 import com.amplitude.experiment.cohort.CohortApi
-import com.amplitude.experiment.util.Logger
-import com.amplitude.experiment.util.SystemLogger
 import io.mockk.every
 import io.mockk.mockk
 import org.junit.Assert
@@ -151,7 +149,7 @@ class LocalEvaluationClientTest {
     @Test
     fun `evaluate with user and group, not targeted`() {
         val cohortConfig = LocalEvaluationConfig(
-            cohortSyncConfiguration = CohortSyncConfiguration("api", "secret")
+            cohortSyncConfig = CohortSyncConfig("api", "secret")
         )
         val cohortApi = mockk<CohortApi>().apply {
             every { getCohort(eq("52gz3yi7"), allAny()) } returns Cohort("52gz3yi7", "User", 2, 1722363790000, setOf("1", "2"))
@@ -177,7 +175,7 @@ class LocalEvaluationClientTest {
     @Test
     fun `evaluate with user, cohort segment targeted`() {
         val cohortConfig = LocalEvaluationConfig(
-            cohortSyncConfiguration = CohortSyncConfiguration("api", "secret")
+            cohortSyncConfig = CohortSyncConfig("api", "secret")
         )
         val cohortApi = mockk<CohortApi>().apply {
             every { getCohort(eq("52gz3yi7"), allAny()) } returns Cohort("52gz3yi7", "User", 2, 1722363790000, setOf("1", "2"))
@@ -198,7 +196,7 @@ class LocalEvaluationClientTest {
     @Test
     fun `evaluate with user, cohort tester targeted`() {
         val cohortConfig = LocalEvaluationConfig(
-            cohortSyncConfiguration = CohortSyncConfiguration("api", "secret")
+            cohortSyncConfig = CohortSyncConfig("api", "secret")
         )
         val cohortApi = mockk<CohortApi>().apply {
             every { getCohort(eq("52gz3yi7"), allAny()) } returns Cohort("52gz3yi7", "User", 2, 1722363790000, setOf("1", "2"))
@@ -220,7 +218,7 @@ class LocalEvaluationClientTest {
     @Test
     fun `evaluate with group, cohort segment targeted`() {
         val cohortConfig = LocalEvaluationConfig(
-            cohortSyncConfiguration = CohortSyncConfiguration("api", "secret")
+            cohortSyncConfig = CohortSyncConfig("api", "secret")
         )
         val cohortApi = mockk<CohortApi>().apply {
             every { getCohort(eq("52gz3yi7"), allAny()) } returns Cohort("52gz3yi7", "User", 2, 1722363790000, setOf("1", "2"))
@@ -243,7 +241,7 @@ class LocalEvaluationClientTest {
     @Test
     fun `evaluate with group, cohort tester targeted`() {
         val cohortConfig = LocalEvaluationConfig(
-            cohortSyncConfiguration = CohortSyncConfiguration("api", "secret")
+            cohortSyncConfig = CohortSyncConfig("api", "secret")
         )
         val cohortApi = mockk<CohortApi>().apply {
             every { getCohort(eq("52gz3yi7"), allAny()) } returns Cohort("52gz3yi7", "User", 2, 1722363790000, setOf("1", "2"))

src/test/kotlin/cohort/CohortDownloadApiTest.kt

@@ -20,7 +20,7 @@ class CohortDownloadApiTest {
     private val httpClient = OkHttpClient()
     private val server = MockWebServer()
     private val url = server.url("/")
-    private val api = DynamicCohortApi(apiKey,secretKey, maxCohortSize, url, null, httpClient)
+    private val api = DynamicCohortApi(apiKey, secretKey, maxCohortSize, url, null, httpClient)
     @Test
     fun `cohort download, success`() {
         val response = cohortResponse("a", setOf("1"))
@@ -87,6 +87,8 @@ class CohortDownloadApiTest {
         server.enqueue(MockResponse().setResponseCode(501))
         server.enqueue(MockResponse().setResponseCode(502))
         server.enqueue(MockResponse().setResponseCode(503))
+        server.enqueue(MockResponse().setResponseCode(503))
+        server.enqueue(MockResponse().setResponseCode(503))
         // Should not be sent in response
         server.enqueue(MockResponse().setResponseCode(204))
         try {

src/test/kotlin/cohort/CohortStorageTest.kt

@@ -2,7 +2,7 @@
 
 package com.amplitude.experiment.cohort
 
-import com.amplitude.experiment.EvaluationProxyConfiguration
+import com.amplitude.experiment.EvaluationProxyConfig
 import com.amplitude.experiment.ExperimentalApi
 import io.mockk.every
 import io.mockk.mockk
@@ -67,7 +67,7 @@ class CohortStorageTest {
         every { cohortMembershipApi.getCohortMemberships(eq("User"), eq("u1")) } returns setOf("a")
         every { cohortMembershipApi.getCohortMemberships(eq("group"), eq("g1")) } returns setOf("b")
         val storage = ProxyCohortStorage(
-            EvaluationProxyConfiguration(""),
+            EvaluationProxyConfig(""),
             cohortMembershipApi
         )
         val cohortA = Cohort("a", "User", 1, 100, setOf("u1"))