fix: remove map based flag config storage (#7)

amplitude · Mar 14, 2023 · 48c0628 · 48c0628
1 parent 63d7e46
commit 48c0628
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 62 deletions.
diff --git a/src/main/kotlin/LocalEvaluationClient.kt b/src/main/kotlin/LocalEvaluationClient.kt
@@ -7,7 +7,6 @@ import com.amplitude.experiment.flag.FlagConfigApiImpl
 import com.amplitude.experiment.flag.FlagConfigService
 import com.amplitude.experiment.flag.FlagConfigServiceConfig
 import com.amplitude.experiment.flag.FlagConfigServiceImpl
-import com.amplitude.experiment.flag.InMemoryFlagConfigStorage
 import com.amplitude.experiment.util.Once
 import com.amplitude.experiment.util.toSerialExperimentUser
 import com.amplitude.experiment.util.toVariant
@@ -23,11 +22,9 @@ class LocalEvaluationClient internal constructor(
     private val httpClient = OkHttpClient()
     private val serverUrl: HttpUrl = config.serverUrl.toHttpUrl()
     private val evaluation: EvaluationEngine = EvaluationEngineImpl()
-    private val flagConfigStorage = InMemoryFlagConfigStorage()
     private val flagConfigService: FlagConfigService = FlagConfigServiceImpl(
         FlagConfigServiceConfig(config.flagConfigPollerIntervalMillis),
         FlagConfigApiImpl(apiKey, serverUrl, httpClient),
-        flagConfigStorage
     )
 
     fun start() {
@@ -38,10 +35,10 @@ class LocalEvaluationClient internal constructor(
 
     @JvmOverloads
     fun evaluate(user: ExperimentUser, flagKeys: List<String> = listOf()): Map<String, Variant> {
-        val flagConfigs = flagConfigService.getFlagConfigs(flagKeys)
+        val flagConfigs = flagConfigService.getFlagConfigs()
         val flagResults = evaluation.evaluate(flagConfigs, user.toSerialExperimentUser().convert())
         return flagResults.mapNotNull { entry ->
-            if (!entry.value.isDefaultVariant) {
+            if (!entry.value.isDefaultVariant && (flagKeys.isEmpty() || flagKeys.contains(entry.key))) {
                 entry.key to SerialVariant(entry.value.variant).toVariant()
             } else {
                 null

diff --git a/src/main/kotlin/flag/FlagConfigApi.kt b/src/main/kotlin/flag/FlagConfigApi.kt
@@ -11,7 +11,7 @@ import java.util.concurrent.CompletableFuture
 
 internal object GetFlagConfigsRequest
 
-internal typealias GetFlagConfigsResponse = Map<String, FlagConfig>
+internal typealias GetFlagConfigsResponse = List<FlagConfig>
 
 internal interface FlagConfigApi {
     fun getFlagConfigs(request: GetFlagConfigsRequest): CompletableFuture<GetFlagConfigsResponse>
@@ -35,9 +35,7 @@ internal class FlagConfigApiImpl(
                 .addHeader("X-Amp-Exp-Library", "experiment-jvm-server/$LIBRARY_VERSION")
                 .build()
         ).thenApply { result ->
-            result.associate {
-                it.flagKey to it.convert()
-            }
+            result.map { it.convert() }
         }
     }
 }
diff --git a/src/main/kotlin/flag/FlagConfigService.kt b/src/main/kotlin/flag/FlagConfigService.kt
@@ -13,30 +13,31 @@ internal data class FlagConfigServiceConfig(
     val flagConfigPollerIntervalMillis: Long,
 )
 
-internal typealias FlagConfigInterceptor = (Map<String, FlagConfig>) -> Unit
+internal typealias FlagConfigInterceptor = (List<FlagConfig>) -> Unit
 
 internal interface FlagConfigService {
     fun start()
     fun stop()
-    fun getFlagConfigs(keys: List<String> = listOf()): List<FlagConfig>
+    fun getFlagConfigs(): List<FlagConfig>
     fun addFlagConfigInterceptor(listener: FlagConfigInterceptor)
 }
 
 internal class FlagConfigServiceImpl(
     private val config: FlagConfigServiceConfig,
     private val flagConfigApi: FlagConfigApi,
-    private val flagConfigStorage: FlagConfigStorage,
 ) : FlagConfigService {
 
     private val lock = Once()
     private val poller = Executors.newSingleThreadScheduledExecutor()
 
     private val interceptorsLock = ReentrantReadWriteLock()
     private val interceptors: MutableSet<FlagConfigInterceptor> = mutableSetOf()
+    private val flagConfigsLock = ReentrantReadWriteLock()
+    private val flagConfigs: MutableList<FlagConfig> = mutableListOf()
 
     private fun refresh() {
         Logger.d("Refreshing flag configs.")
-        val flagConfigs = getFlagConfigs()
+        val flagConfigs = fetchFlagConfigs()
         storeFlagConfigs(flagConfigs)
         Logger.d("Refreshed ${flagConfigs.size} flag configs.")
     }
@@ -57,11 +58,9 @@ internal class FlagConfigServiceImpl(
         poller.shutdown()
     }
 
-    override fun getFlagConfigs(keys: List<String>): List<FlagConfig> {
-        return if (keys.isEmpty()) {
-            flagConfigStorage.getAll().values.toList()
-        } else {
-            keys.mapNotNull { flagConfigStorage.get(it) }
+    override fun getFlagConfigs(): List<FlagConfig> {
+        return flagConfigsLock.read {
+            flagConfigs
         }
     }
 
@@ -71,7 +70,7 @@ internal class FlagConfigServiceImpl(
         }
     }
 
-    private fun getFlagConfigs(): Map<String, FlagConfig> {
+    private fun fetchFlagConfigs(): List<FlagConfig> {
         return flagConfigApi.getFlagConfigs(GetFlagConfigsRequest).get().apply {
             interceptorsLock.read {
                 interceptors.forEach { interceptor ->
@@ -81,7 +80,10 @@ internal class FlagConfigServiceImpl(
         }
     }
 
-    private fun storeFlagConfigs(flagConfigs: Map<String, FlagConfig>) {
-        flagConfigStorage.overwrite(flagConfigs)
+    private fun storeFlagConfigs(flagConfigs: List<FlagConfig>) {
+        flagConfigsLock.write {
+            this.flagConfigs.clear()
+            this.flagConfigs.addAll(flagConfigs)
+        }
     }
 }
diff --git a/src/main/kotlin/flag/FlagConfigStorage.kt b/src/main/kotlin/flag/FlagConfigStorage.kt
diff --git a/src/test/kotlin/LocalEvaluationClientTest.kt b/src/test/kotlin/LocalEvaluationClientTest.kt
@@ -50,7 +50,7 @@ class LocalEvaluationClientTest {
         }
         val millis = duration / 1000.0 / 1000.0
         println("1 flag: $millis")
-        Assert.assertTrue(millis < 10)
+        Assert.assertTrue(millis < 20)
     }
 
     @Test
@@ -65,7 +65,7 @@ class LocalEvaluationClientTest {
         }
         val millis = total / 1000.0 / 1000.0
         println("10 flags: $millis")
-        Assert.assertTrue(millis < 20)
+        Assert.assertTrue(millis < 40)
     }
 
     @Test
@@ -80,7 +80,7 @@ class LocalEvaluationClientTest {
         }
         val millis = total / 1000.0 / 1000.0
         println("100 flags: $millis")
-        Assert.assertTrue(millis < 20)
+        Assert.assertTrue(millis < 80)
     }
 
     @Test
@@ -95,6 +95,6 @@ class LocalEvaluationClientTest {
         }
         val millis = total / 1000.0 / 1000.0
         println("1000 flags: $millis")
-        Assert.assertTrue(millis < 100)
+        Assert.assertTrue(millis < 160)
     }
 }