feat: support subset flags fetch (#19)

sdk/src/main/java/com/amplitude/experiment/DefaultExperimentClient.kt

@@ -29,6 +29,7 @@ import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.TimeoutException
 import kotlin.jvm.Throws
+import org.json.JSONArray
 
 internal class DefaultExperimentClient internal constructor(
     private val apiKey: String,
@@ -65,10 +66,14 @@ internal class DefaultExperimentClient internal constructor(
         }
 
     override fun fetch(user: ExperimentUser?): Future<ExperimentClient> {
+        return fetch(user, null)
+    }
+
+    override fun fetch(user: ExperimentUser?, options: FetchOptions?): Future<ExperimentClient> {
         this.user = user ?: this.user
         return executorService.submit(Callable {
             val fetchUser = getUserMergedWithProviderOrWait(10000)
-            fetchInternal(fetchUser, config.fetchTimeoutMillis, config.retryFetchOnFailure)
+            fetchInternal(fetchUser, config.fetchTimeoutMillis, config.retryFetchOnFailure, options)
             this
         })
     }
@@ -174,16 +179,16 @@ internal class DefaultExperimentClient internal constructor(
     }
 
     @Throws
-    private fun fetchInternal(user: ExperimentUser, timeoutMillis: Long, retry: Boolean) {
+    private fun fetchInternal(user: ExperimentUser, timeoutMillis: Long, retry: Boolean, options: FetchOptions?) {
         if (retry) {
             stopRetries()
         }
         try {
-            val variants = doFetch(user, timeoutMillis).get()
-            storeVariants(variants)
+            val variants = doFetch(user, timeoutMillis, options).get()
+            storeVariants(variants, options)
         } catch (e: Exception) {
             if (retry) {
-                startRetries(user)
+                startRetries(user, options)
             }
             throw e
         }
@@ -192,6 +197,7 @@ internal class DefaultExperimentClient internal constructor(
     private fun doFetch(
         user: ExperimentUser,
         timeoutMillis: Long,
+        options: FetchOptions?
     ): Future<Map<String, Variant>> {
         if (user.userId == null && user.deviceId == null) {
             Logger.w("user id and device id are null; amplitude may not resolve identity")
@@ -205,12 +211,20 @@ internal class DefaultExperimentClient internal constructor(
         val url = serverUrl.newBuilder()
             .addPathSegments("sdk/vardata")
             .build()
-        val request = Request.Builder()
+        val builder = Request.Builder()
             .get()
             .url(url)
             .addHeader("Authorization", "Api-Key $apiKey")
             .addHeader("X-Amp-Exp-User", userBase64)
-            .build()
+        if (!options?.flagKeys.isNullOrEmpty()) {
+            val flagKeysBase64 = JSONArray(options?.flagKeys)
+                .toString()
+                .toByteArray(Charsets.UTF_8)
+                .toByteString()
+                .base64()
+            builder.addHeader("X-Amp-Exp-Flag-Keys", flagKeysBase64)
+        }
+        val request = builder.build()
         val call = httpClient.newCall(request)
         call.timeout().timeout(timeoutMillis, TimeUnit.MILLISECONDS)
         val future = AsyncFuture<Map<String, Variant>>(call)
@@ -233,10 +247,10 @@ internal class DefaultExperimentClient internal constructor(
         return future
     }
 
-    private fun startRetries(user: ExperimentUser) = synchronized(backoffLock) {
+    private fun startRetries(user: ExperimentUser, options: FetchOptions?) = synchronized(backoffLock) {
         backoff?.cancel()
         backoff = executorService.backoff(backoffConfig) {
-            fetchInternal(user, fetchBackoffTimeoutMillis, false)
+            fetchInternal(user, fetchBackoffTimeoutMillis, false, options)
         }
     }
 
@@ -261,10 +275,17 @@ internal class DefaultExperimentClient internal constructor(
         return variants
     }
 
-    private fun storeVariants(variants: Map<String, Variant>) = synchronized(storageLock) {
-        storage.clear()
+    private fun storeVariants(variants: Map<String, Variant>, options: FetchOptions?) = synchronized(storageLock) {
+        val failedFlagKeys = options?.flagKeys ?.toMutableList() ?: mutableListOf()
+        if (options?.flagKeys == null) {
+            storage.clear()
+        }
         for (entry in variants.entries) {
             storage.put(entry.key, entry.value)
+            failedFlagKeys.remove(entry.key)
+        }
+        for (key in failedFlagKeys) {
+            storage.remove(key)
         }
         Logger.d("Stored variants: $variants")
     }

sdk/src/main/java/com/amplitude/experiment/ExperimentClient.kt

@@ -6,7 +6,6 @@ import java.util.concurrent.Future
  * An experiment client manages a set of experiments and flags for a given user.
  */
 interface ExperimentClient {
-
     /**
      * Assign the given user to the SDK and asynchronously fetch all variants
      * from the server. Subsequent calls may omit the user from the argument to
@@ -25,6 +24,25 @@ interface ExperimentClient {
      */
     fun fetch(user: ExperimentUser? = null): Future<ExperimentClient>
 
+    /**
+     * Assign the given user to the SDK and asynchronously fetch all variants
+     * from the server. Subsequent calls may omit the user from the argument to
+     * use the user from the previous call, or set previously using [setUser].
+     *
+     * If an [ExperimentUserProvider] has been set, the argument user will
+     * be merged with the provider user, preferring user fields from the
+     * argument user and falling back on the provider for fields which are null
+     * or undefined.
+     *
+     * @param user The user to fetch variants for. If null use the user stored
+     *             in the client.
+     * @param options Optional fetch options, could config to fetch subset flags.
+     * @returns Future that resolves when the request for variants completes.
+     * @see ExperimentUser
+     * @see ExperimentUserProvider
+     */
+    fun fetch(user: ExperimentUser? = null, options: FetchOptions? = null): Future<ExperimentClient>
+
     /**
      * Returns the stored variant for the provided key.
      *

sdk/src/main/java/com/amplitude/experiment/FetchOptions.kt

@@ -0,0 +1,5 @@
+package com.amplitude.experiment
+
+data class FetchOptions @JvmOverloads constructor(
+    @JvmField val flagKeys: List<String>? = null
+)

sdk/src/main/java/com/amplitude/experiment/storage/InMemoryStorage.kt

@@ -15,6 +15,10 @@ internal class InMemoryStorage : Storage {
         return data[key]
     }
 
+    override fun remove(key: String) {
+        data.remove(key)
+    }
+
     override fun getAll(): Map<String, Variant> {
         return data.toMap()
     }

sdk/src/main/java/com/amplitude/experiment/storage/SharedPrefsStorage.kt

@@ -34,6 +34,10 @@ internal class SharedPrefsStorage(
         return sharedPrefs.getString(key, null).toVariant()
     }
 
+    override fun remove(key: String) {
+        sharedPrefs.edit().remove(key).apply()
+    }
+
     override fun getAll(): Map<String, Variant> {
         val result = mutableMapOf<String, Variant>()
         for ((key, value) in sharedPrefs.all) {

sdk/src/main/java/com/amplitude/experiment/storage/Storage.kt

@@ -5,6 +5,7 @@ import com.amplitude.experiment.Variant
 internal interface Storage {
     fun put(key: String, variant: Variant)
     fun get(key: String): Variant?
+    fun remove(key: String)
     fun getAll(): Map<String, Variant>
     fun clear()
 }

sdk/src/test/java/com/amplitude/experiment/ExperimentClientTest.kt

@@ -187,6 +187,30 @@ class ExperimentClientTest {
         Assert.assertEquals(newUser, client.getUser())
     }
 
+    @Test
+    fun `test fetch user with flag`() {
+        client.fetch(testUser, FetchOptions(flagKeys = listOf(KEY))).get()
+        val variant = client.variant(KEY)
+        Assert.assertNotNull(variant)
+        Assert.assertEquals(serverVariant, variant)
+    }
+
+    @Test
+    fun `test fetch user with invalid flags`() {
+        val invalidKey = "invalid"
+        client.fetch(testUser, FetchOptions(flagKeys = listOf(KEY, INITIAL_KEY, invalidKey))).get()
+        val variant = client.variant(KEY)
+        Assert.assertNotNull(variant)
+        Assert.assertEquals(serverVariant, variant)
+
+        val firstFallback = Variant("first")
+        val initialVariant = client.variant(INITIAL_KEY, firstFallback)
+        Assert.assertEquals(firstFallback, initialVariant)
+
+        val invalidVariant = client.variant(invalidKey)
+        Assert.assertEquals(fallbackVariant, invalidVariant)
+    }
+
     @Test
     fun `test exposure event through analytics provider when variant called`() {
         var didExposureGetTracked = false