From 5a855c904b011b931467a29f48914bd7867f83db Mon Sep 17 00:00:00 2001
From: Brian Giori <brian.giori@amplitude.com>
Date: Mon, 23 Sep 2024 16:05:57 -0700
Subject: [PATCH] fix: add metric for cohort download too large (#33)

---
 src/main/kotlin/LocalEvaluationConfig.kt       | 1 +
 src/main/kotlin/cohort/CohortLoader.kt         | 3 ++-
 src/main/kotlin/deployment/DeploymentRunner.kt | 9 +++++++--
 src/main/kotlin/util/Metrics.kt                | 5 +++++
 src/test/kotlin/cohort/CohortLoaderTest.kt     | 6 +++++-
 5 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/main/kotlin/LocalEvaluationConfig.kt b/src/main/kotlin/LocalEvaluationConfig.kt
index 67a5a72..0475204 100644
--- a/src/main/kotlin/LocalEvaluationConfig.kt
+++ b/src/main/kotlin/LocalEvaluationConfig.kt
@@ -208,6 +208,7 @@ interface LocalEvaluationMetrics {
     fun onFlagConfigFetchFailure(exception: Exception)
     fun onFlagConfigFetchOriginFallback(exception: Exception)
     fun onCohortDownload()
+    fun onCohortDownloadTooLarge(exception: Exception)
     fun onCohortDownloadFailure(exception: Exception)
     fun onCohortDownloadOriginFallback(exception: Exception)
     fun onProxyCohortMembership()
diff --git a/src/main/kotlin/cohort/CohortLoader.kt b/src/main/kotlin/cohort/CohortLoader.kt
index f43d2a4..97c5721 100644
--- a/src/main/kotlin/cohort/CohortLoader.kt
+++ b/src/main/kotlin/cohort/CohortLoader.kt
@@ -44,7 +44,8 @@ internal class CohortLoader(
                     } catch (e: CohortNotModifiedException) {
                         // Do nothing
                     } catch (e: CohortTooLargeException) {
-                        Logger.e("Cohort too large", e)
+                        metrics.onCohortDownloadTooLarge(e)
+                        throw e
                     }
                 }
             }, executor).whenComplete { _, _ -> jobs.remove(cohortId) }
diff --git a/src/main/kotlin/deployment/DeploymentRunner.kt b/src/main/kotlin/deployment/DeploymentRunner.kt
index 54d81d6..fcec6fe 100644
--- a/src/main/kotlin/deployment/DeploymentRunner.kt
+++ b/src/main/kotlin/deployment/DeploymentRunner.kt
@@ -60,12 +60,17 @@ internal class DeploymentRunner(
                     try {
                         val cohortIds = flagConfigStorage.getFlagConfigs().values.getAllCohortIds()
                         for (cohortId in cohortIds) {
-                            cohortLoader.loadCohort(cohortId)
+                            cohortLoader.loadCohort(cohortId).handle { _, exception ->
+                                if (exception != null) {
+                                    Logger.e("Failed to load cohort $cohortId", exception)
+                                }
+                            }
                         }
                     } catch (t: Throwable) {
                         Logger.e("Refresh cohorts failed.", t)
                     }
-                }, cohortPollingInterval,
+                },
+                cohortPollingInterval,
                 cohortPollingInterval,
                 TimeUnit.MILLISECONDS
             )
diff --git a/src/main/kotlin/util/Metrics.kt b/src/main/kotlin/util/Metrics.kt
index 2748652..bc06dbb 100644
--- a/src/main/kotlin/util/Metrics.kt
+++ b/src/main/kotlin/util/Metrics.kt
@@ -76,6 +76,11 @@ internal class LocalEvaluationMetricsWrapper(
         executor?.execute { metrics.onCohortDownload() }
     }
 
+    override fun onCohortDownloadTooLarge(exception: Exception) {
+        val metrics = metrics ?: return
+        executor?.execute { metrics.onCohortDownloadTooLarge(exception) }
+    }
+
     override fun onCohortDownloadFailure(exception: Exception) {
         val metrics = metrics ?: return
         executor?.execute { metrics.onCohortDownloadFailure(exception) }
diff --git a/src/test/kotlin/cohort/CohortLoaderTest.kt b/src/test/kotlin/cohort/CohortLoaderTest.kt
index 4664a70..1878bb9 100644
--- a/src/test/kotlin/cohort/CohortLoaderTest.kt
+++ b/src/test/kotlin/cohort/CohortLoaderTest.kt
@@ -35,7 +35,11 @@ class CohortLoaderTest {
         `when`(api.getCohort("b", null)).thenReturn(cohortB)
         val storage = InMemoryCohortStorage()
         val loader = CohortLoader(api, storage)
-        loader.loadCohort("a").get()
+        try {
+            loader.loadCohort("a").get()
+        } catch (t: Throwable) {
+            // Expected
+        }
         loader.loadCohort("b").get()
         val storageDescriptionA = storage.getCohort("a")
         val storageDescriptionB = storage.getCohort("b")