fix: CohortTooLargeException not thrown when in proxy mode (#36)

amplitude · Nov 1, 2024 · cc10582 · cc10582
1 parent b5e0684
commit cc10582
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/src/main/kotlin/cohort/CohortApi.kt b/src/main/kotlin/cohort/CohortApi.kt
@@ -14,10 +14,13 @@ import okhttp3.OkHttpClient
 import java.util.Base64
 import java.util.concurrent.ExecutionException
 
-internal class CohortTooLargeException(cohortId: String, maxCohortSize: Int) : RuntimeException(
+open class CohortTooLargeException(cohortId: String, maxCohortSize: Int) : RuntimeException(
     "Cohort $cohortId exceeds the maximum cohort size defined in the SDK configuration $maxCohortSize"
 )
 
+internal class ProxyCohortTooLargeException(cohortId: String, maxCohortSize: Int) :
+    CohortTooLargeException(cohortId, maxCohortSize)
+
 internal class CohortNotModifiedException(cohortId: String) : RuntimeException(
     "Cohort $cohortId has not been modified."
 )
@@ -68,11 +71,15 @@ internal class DynamicCohortApi(
             } catch (e: CohortNotModifiedException) {
                 throw e
             } catch (e: CohortTooLargeException) {
-                throw e
+                throw ProxyCohortTooLargeException(cohortId, maxCohortSize)
             } catch (e: Exception) {
                 Logger.w("Downloading cohort $cohortId from proxy failed. Falling back to Amplitude.", e)
                 metrics.onCohortDownloadOriginFallback(e)
-                getCohort(serverUrl, cohortId, cohort)
+                try {
+                    getCohort(serverUrl, cohortId, cohort)
+                } catch (e: CohortTooLargeException) {
+                    throw ProxyCohortTooLargeException(cohortId, maxCohortSize)
+                }
             }
         } else {
             getCohort(serverUrl, cohortId, cohort)

diff --git a/src/main/kotlin/cohort/CohortLoader.kt b/src/main/kotlin/cohort/CohortLoader.kt
@@ -43,6 +43,8 @@ internal class CohortLoader(
                         cohortStorage.putCohort(cohort)
                     } catch (e: CohortNotModifiedException) {
                         // Do nothing
+                    } catch (e: ProxyCohortTooLargeException) {
+                        // Do nothing
                     } catch (e: CohortTooLargeException) {
                         metrics.onCohortDownloadTooLarge(e)
                         throw e