Implement caching Toggle.Store and Toggle@enabled() flow API

Author: aitorvs

build.gradle

@@ -53,7 +53,7 @@ subprojects {
         apply plugin: 'org.jetbrains.dokka'
     }
 
-    String[] allowAndroidTestsIn = ["app", "sync-lib", "httpsupgrade-impl"]
+    String[] allowAndroidTestsIn = ["app", "sync-lib", "httpsupgrade-impl", "feature-toggles-impl"]
     if (!allowAndroidTestsIn.contains(project.name)) {
         project.projectDir.eachFile(groovy.io.FileType.DIRECTORIES) { File parent ->
             if (parent.name == "src") {

feature-toggles/feature-toggles-api/build.gradle

@@ -36,5 +36,9 @@ dependencies {
 
     implementation Google.dagger
     implementation "org.apache.commons:commons-math3:_"
+    implementation("com.google.guava:guava:_") {
+        exclude group: 'com.google.guava', module: 'listenablefuture'
+    }
+    implementation KotlinX.coroutines.core
 }
 

feature-toggles/feature-toggles-api/src/main/java/com/duckduckgo/feature/toggles/api/FeatureToggles.kt

@@ -20,7 +20,12 @@ import com.duckduckgo.feature.toggles.api.Toggle.FeatureName
 import com.duckduckgo.feature.toggles.api.Toggle.State
 import com.duckduckgo.feature.toggles.api.Toggle.State.Cohort
 import com.duckduckgo.feature.toggles.api.Toggle.State.CohortName
+import com.duckduckgo.feature.toggles.api.internal.CachedToggleStore
+import com.duckduckgo.feature.toggles.api.internal.CachedToggleStore.Listener
 import com.duckduckgo.feature.toggles.internal.api.FeatureTogglesCallback
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.callbackFlow
 import org.apache.commons.math3.distribution.EnumeratedIntegerDistribution
 import java.lang.reflect.Method
 import java.lang.reflect.Proxy
@@ -119,7 +124,7 @@ class FeatureToggles private constructor(
             }.getOrNull() != null
 
             return ToggleImpl(
-                store = store,
+                store = if (store is CachedToggleStore) store else CachedToggleStore(store),
                 key = getToggleNameForMethod(method),
                 defaultValue = resolvedDefaultValue,
                 isInternalAlwaysEnabled = isInternalAlwaysEnabledAnnotated,
@@ -172,6 +177,39 @@ interface Toggle {
      */
     suspend fun enroll(): Boolean
 
+    /**
+     * Returns a cold [Flow] of [Boolean] values representing whether this toggle is enabled.
+     *
+     * ### Behavior
+     * - When a collector starts, the current toggle value is emitted immediately.
+     * - Subsequent emissions occur whenever the underlying [store] writes a new [State].
+     * - The flow is cold: a listener is only registered while it is being collected.
+     * - When collection is cancelled or completed, the registered listener is automatically unregistered.
+     *
+     * ### Thread-safety
+     * Emissions are delivered on the coroutine context where the flow is collected.
+     * Multiple collectors will each register their own listener instance.
+     *
+     * ### Example
+     * ```
+     * viewModelScope.launch {
+     *     toggle.enabled()
+     *         .distinctUntilChanged()
+     *         .collect { enabled ->
+     *             if (enabled) {
+     *                 showOnboarding()
+     *             } else {
+     *                 showLoading()
+     *             }
+     *         }
+     * }
+     * ```
+     *
+     * @return a cold [Flow] that emits the current enabled state and any subsequent changes
+     *         until the collector is cancelled.
+     */
+    fun enabled(): Flow<Boolean>
+
     /**
      * This method
      *    - Returns whether the feature flag state is enabled or disabled.
@@ -386,6 +424,28 @@ internal class ToggleImpl constructor(
         return enrollInternal()
     }
 
+    override fun enabled(): Flow<Boolean> = callbackFlow {
+        // emit current value when someone starts collecting
+        trySend(isEnabled())
+
+        val unsubscribe = when (val s = store) {
+            is CachedToggleStore -> {
+                s.setListener(
+                    object : Listener {
+                        override fun onToggleStored(newValue: State) {
+                            // emit value just stored
+                            trySend(isEnabled())
+                        }
+                    },
+                )
+            }
+            else -> { -> Unit }
+        }
+
+        // when flow collection is cancelled/closed, run the unsubscribe to avoid leaking the listener
+        awaitClose { unsubscribe() }
+    }
+
     private fun enrollInternal(force: Boolean = false): Boolean {
         // if the Toggle is not enabled, then we don't enroll
         if (isEnabled() == false) {

feature-toggles/feature-toggles-api/src/main/java/com/duckduckgo/feature/toggles/api/internal/CachedToggleStore.kt

@@ -0,0 +1,110 @@
+/*
+ * Copyright (c) 2025 DuckDuckGo
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.duckduckgo.feature.toggles.api.internal
+
+import com.duckduckgo.feature.toggles.api.Toggle
+import com.duckduckgo.feature.toggles.api.Toggle.State
+import com.google.common.cache.CacheBuilder
+import com.google.common.cache.CacheLoader
+import com.google.common.cache.LoadingCache
+import org.jetbrains.annotations.TestOnly
+import java.util.Optional
+import kotlin.jvm.optionals.getOrNull
+
+class CachedToggleStore constructor(
+    private val store: Toggle.Store,
+) : Toggle.Store {
+    @Volatile
+    private var listener: Listener? = null
+
+    private val cache: LoadingCache<String, Optional<State>> =
+        CacheBuilder
+            .newBuilder()
+            .maximumSize(50)
+            .build(
+                object : CacheLoader<String, Optional<State>>() {
+                    override fun load(key: String): Optional<State> = Optional.ofNullable(store.get(key))
+                },
+            )
+
+    override fun set(
+        key: String,
+        state: State,
+    ) {
+        cache.asMap().compute(key) { k, _ ->
+            store.set(k, state)
+            Optional.of(state)
+        }
+        // Notify AFTER compute() to avoid deadlocks or re-entrancy into the cache/store.
+        // If the store.set() above throws, this never runs (which is what we want).
+        // Swallow listener exceptions so they don't break writes.
+        listener?.runCatching { onToggleStored(state) }
+    }
+
+    /**
+     * Registers a [Listener] to observe changes in toggle states stored by this [CachedToggleStore].
+     *
+     * Only a single listener is supported at a time. When a new listener is set, it replaces any
+     * previously registered listener. To avoid memory leaks, callers should always invoke the returned
+     * unsubscribe function when the listener is no longer needed (for example, when the collector
+     * of a [kotlinx.coroutines.flow.callbackFlow] is closed).
+     *
+     * Example usage:
+     * ```
+     * val unsubscribe = cachedToggleStore.setListener(object : CachedToggleStore.Listener {
+     *     override fun onToggleStored(newValue: Toggle.State, oldValue: Toggle.State?) {
+     *         // React to state change
+     *     }
+     * })
+     *
+     * // Later, when no longer interested:
+     * unsubscribe()
+     * ```
+     *
+     * @param listener the [Listener] instance that will receive callbacks for each `set` operation.
+     * @return a function that removes the listener when invoked. The returned function is safe to call
+     *         multiple times and will only clear the listener if it is the same instance that was
+     *         originally registered.
+     */
+    fun setListener(listener: Listener?): () -> Unit {
+        this.listener = listener
+
+        return { if (this.listener === listener) this.listener = null }
+    }
+
+    /**
+     * DO NOT USE outside tests
+     */
+    @TestOnly
+    fun invalidateAll() {
+        cache.invalidateAll()
+    }
+
+    override fun get(key: String): State? {
+        val value = cache.get(key).getOrNull()
+        if (value == null) {
+            // avoid negative caching
+            cache.invalidate(key)
+        }
+
+        return value
+    }
+
+    interface Listener {
+        fun onToggleStored(newValue: State)
+    }
+}

feature-toggles/feature-toggles-impl/build.gradle

@@ -52,10 +52,24 @@ dependencies {
     testImplementation project(':data-store-test')
     testImplementation "org.mockito.kotlin:mockito-kotlin:_"
     testImplementation "com.squareup.moshi:moshi-kotlin:_"
+    testImplementation CashApp.turbine
     testImplementation Testing.robolectric
     testImplementation AndroidX.test.ext.junit
     testImplementation Square.retrofit2.converter.moshi
     testImplementation Testing.junit4
+    androidTestImplementation AndroidX.test.runner
+    androidTestImplementation AndroidX.test.ext.junit
+    androidTestImplementation Square.retrofit2.converter.moshi
+
+    coreLibraryDesugaring Android.tools.desugarJdkLibs
+}
+
+configurations.all {
+    exclude(group: "com.google.guava", module: "listenablefuture")
+}
+
+tasks.register('androidTestsBuild') {
+    dependsOn 'assembleDebugAndroidTest'
 }
 
 anvil {
@@ -68,4 +82,7 @@ android {
   lintOptions {
       baseline file("lint-baseline.xml")
   }
+  compileOptions {
+    coreLibraryDesugaringEnabled = true
+  }
 }

feature-toggles/feature-toggles-impl/src/test/java/com/duckduckgo/feature/toggles/api/FeatureTogglesTest.kt

@@ -17,14 +17,17 @@
 package com.duckduckgo.feature.toggles.api
 
 import android.annotation.SuppressLint
+import app.cash.turbine.test
 import com.duckduckgo.appbuildconfig.api.BuildFlavor
 import com.duckduckgo.feature.toggles.api.Cohorts.CONTROL
 import com.duckduckgo.feature.toggles.api.Cohorts.TREATMENT
 import com.duckduckgo.feature.toggles.api.Toggle.DefaultFeatureValue
 import com.duckduckgo.feature.toggles.api.Toggle.FeatureName
 import com.duckduckgo.feature.toggles.api.Toggle.State
 import com.duckduckgo.feature.toggles.api.Toggle.State.CohortName
+import com.duckduckgo.feature.toggles.api.internal.CachedToggleStore
 import com.duckduckgo.feature.toggles.internal.api.FeatureTogglesCallback
+import kotlinx.coroutines.flow.drop
 import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
@@ -41,13 +44,13 @@ class FeatureTogglesTest {
 
     private lateinit var feature: TestFeature
     private lateinit var provider: FakeProvider
-    private lateinit var toggleStore: FakeToggleStore
+    private lateinit var toggleStore: Toggle.Store
     private lateinit var callback: FakeFeatureTogglesCallback
 
     @Before
     fun setup() {
         provider = FakeProvider()
-        toggleStore = FakeToggleStore()
+        toggleStore = CachedToggleStore(FakeToggleStore())
         callback = FakeFeatureTogglesCallback()
         feature = FeatureToggles.Builder()
             .store(toggleStore)
@@ -121,6 +124,40 @@ class FeatureTogglesTest {
         feature.noDefaultValue().isEnabled()
     }
 
+    @Test
+    fun whenEnabledByDefaultThenEmitEnabled() = runTest {
+        feature.enabledByDefault().enabled().test {
+            assertTrue(awaitItem())
+            expectNoEvents()
+        }
+    }
+
+    @Test
+    fun whenEnabledByDefaultAndSetEnabledThenEmitTwoEnables() = runTest {
+        feature.enabledByDefault().enabled().test {
+            assertTrue(awaitItem())
+            feature.enabledByDefault().setRawStoredState(Toggle.State(enable = false))
+            assertFalse(awaitItem())
+            expectNoEvents()
+        }
+    }
+
+    @Test
+    fun enableValuesSetBeforeRegistrationGetLost() = runTest {
+        feature.enabledByDefault().setRawStoredState(Toggle.State(enable = false))
+        feature.enabledByDefault().enabled().test {
+            assertFalse(awaitItem())
+            expectNoEvents()
+        }
+    }
+
+    @Test
+    fun whenDroppingEmissionThenNoValueEmitted() = runTest {
+        feature.enabledByDefault().enabled().drop(1).test {
+            expectNoEvents()
+        }
+    }
+
     @Test(expected = IllegalArgumentException::class)
     fun whenWrongReturnValueThenThrow() {
         feature.wrongReturnValue()