feat: add stop() for LocalEvaluationClient (#38)

amplitude · Dec 12, 2024 · 7b97728 · 7b97728
1 parent 519ae84
commit 7b97728
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 2 deletions.
diff --git a/src/main/kotlin/LocalEvaluationClient.kt b/src/main/kotlin/LocalEvaluationClient.kt
@@ -84,6 +84,10 @@ class LocalEvaluationClient internal constructor(
         }
     }
 
+    fun stop() {
+        deploymentRunner.stop()
+    }
+
     private fun createAssignmentService(deploymentKey: String): AssignmentService? {
         if (config.assignmentConfiguration == null) return null
         return AmplitudeAssignmentService(
@@ -98,6 +102,7 @@ class LocalEvaluationClient internal constructor(
             metrics = metrics,
         )
     }
+
     @JvmOverloads
     @Deprecated(
         "Use the evaluateV2 method. EvaluateV2 returns variant objects with default values (e.g. null/off) if the user is evaluated, but not assigned a variant.",

diff --git a/src/main/kotlin/deployment/DeploymentRunner.kt b/src/main/kotlin/deployment/DeploymentRunner.kt
@@ -97,8 +97,10 @@ internal class DeploymentRunner(
     }
 
     fun stop() {
-        poller.shutdown()
-        flagConfigUpdater.shutdown()
+        lock.once {
+            poller.shutdown()
+            flagConfigUpdater.shutdown()
+        }
     }
 
     private fun getCohortPollingInterval(): Long {

diff --git a/src/test/kotlin/LocalEvaluationClientTest.kt b/src/test/kotlin/LocalEvaluationClientTest.kt
@@ -29,6 +29,7 @@ class LocalEvaluationClientTest {
         val variants = client.evaluate(ExperimentUser(userId = "test_user"))
         val variant = variants["sdk-local-evaluation-ci-test"]
         Assert.assertEquals(Variant(key = "on", value = "on", payload = "payload"), variant?.copy(metadata = null))
+        client.stop()
     }
 
     @Test
@@ -53,6 +54,7 @@ class LocalEvaluationClientTest {
         )
         val variant = variants["this-flag-doesnt-exit"]
         Assert.assertNull(variant)
+        client.stop()
     }
 
     @Test
@@ -65,6 +67,7 @@ class LocalEvaluationClientTest {
         val millis = duration / 1000.0 / 1000.0
         println("1 flag: $millis")
         Assert.assertTrue(millis < 20)
+        client.stop()
     }
 
     @Test
@@ -80,6 +83,7 @@ class LocalEvaluationClientTest {
         val millis = total / 1000.0 / 1000.0
         println("10 flags: $millis")
         Assert.assertTrue(millis < 40)
+        client.stop()
     }
 
     @Test
@@ -95,6 +99,7 @@ class LocalEvaluationClientTest {
         val millis = total / 1000.0 / 1000.0
         println("100 flags: $millis")
         Assert.assertTrue(millis < 80)
+        client.stop()
     }
 
     @Test
@@ -110,6 +115,7 @@ class LocalEvaluationClientTest {
         val millis = total / 1000.0 / 1000.0
         println("1000 flags: $millis")
         Assert.assertTrue(millis < 160)
+        client.stop()
     }
 
     @Test
@@ -119,6 +125,7 @@ class LocalEvaluationClientTest {
         val variants = client.evaluate(ExperimentUser(userId = "user_id", deviceId = "device_id"))
         val variant = variants["sdk-ci-local-dependencies-test"]
         Assert.assertEquals(variant?.copy(metadata = null), Variant(key = "control", value = "control", payload = null))
+        client.stop()
     }
 
     @Test
@@ -131,6 +138,7 @@ class LocalEvaluationClientTest {
         )
         val variant = variants["sdk-ci-local-dependencies-test"]
         Assert.assertEquals(variant?.copy(metadata = null), Variant(key = "control", value = "control", payload = null))
+        client.stop()
     }
 
     @Test
@@ -143,6 +151,7 @@ class LocalEvaluationClientTest {
         )
         val variant = variants["sdk-ci-local-dependencies-test"]
         Assert.assertEquals(variant, null)
+        client.stop()
     }
 
     @Test
@@ -152,6 +161,7 @@ class LocalEvaluationClientTest {
         val variants = client.evaluate(ExperimentUser(userId = "user_id", deviceId = "device_id"))
         val variant = variants["sdk-ci-local-dependencies-test-holdout"]
         Assert.assertEquals(variant, null)
+        client.stop()
     }
 
     @Test
@@ -178,6 +188,7 @@ class LocalEvaluationClientTest {
         val groupVariant = client.evaluateV2(user, setOf("sdk-local-evaluation-group-cohort-ci-test"))["sdk-local-evaluation-group-cohort-ci-test"]
         assertEquals("off", groupVariant?.key)
         assertNull(groupVariant?.value)
+        client.stop()
     }
 
     @Test
@@ -200,6 +211,7 @@ class LocalEvaluationClientTest {
         val userVariant = client.evaluateV2(user, setOf("sdk-local-evaluation-user-cohort-ci-test"))["sdk-local-evaluation-user-cohort-ci-test"]
         assertEquals("on", userVariant?.key)
         assertEquals("on", userVariant?.value)
+        client.stop()
     }
 
     @Test
@@ -222,6 +234,7 @@ class LocalEvaluationClientTest {
         val userVariant = client.evaluateV2(user, setOf("sdk-local-evaluation-user-cohort-ci-test"))["sdk-local-evaluation-user-cohort-ci-test"]
         assertEquals("on", userVariant?.key)
         assertEquals("on", userVariant?.value)
+        client.stop()
     }
 
     @Test
@@ -246,6 +259,7 @@ class LocalEvaluationClientTest {
         val groupVariant = client.evaluateV2(user, setOf("sdk-local-evaluation-group-cohort-ci-test"))["sdk-local-evaluation-group-cohort-ci-test"]
         assertEquals("on", groupVariant?.key)
         assertEquals("on", groupVariant?.value)
+        client.stop()
     }
 
     @Test
@@ -270,6 +284,7 @@ class LocalEvaluationClientTest {
         val groupVariant = client.evaluateV2(user, setOf("sdk-local-evaluation-group-cohort-ci-test"))["sdk-local-evaluation-group-cohort-ci-test"]
         assertEquals("on", groupVariant?.key)
         assertEquals("on", groupVariant?.value)
+        client.stop()
     }
 
     @Test
@@ -283,6 +298,7 @@ class LocalEvaluationClientTest {
         val variants = client.evaluate(ExperimentUser(userId = "test_user"))
         val variant = variants["sdk-local-evaluation-ci-test"]
         Assert.assertEquals(Variant(key = "on", value = "on", payload = "payload"), variant?.copy(metadata = null))
+        client.stop()
     }
 
     @Test
@@ -310,5 +326,6 @@ class LocalEvaluationClientTest {
         val userVariant = client.evaluateV2(user, setOf("sdk-local-evaluation-user-cohort-ci-test"))["sdk-local-evaluation-user-cohort-ci-test"]
         assertEquals("on", userVariant?.key)
         assertEquals("on", userVariant?.value)
+        client.stop()
     }
 }