fix: add metric for cohort download too large (#33)

amplitude · Sep 23, 2024 · 5a855c9 · 5a855c9
1 parent e9ee829
commit 5a855c9
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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")