Support Sample Rate in message param (#3)

e.g. `gorets:1|c|@0.1` See https://github.com/statsd/statsd/blob/master/docs/metric_types.md#sampling
kpavlov · May 29, 2023 · 9b6cb7a · 9b6cb7a
1 parent 3e32976
commit 9b6cb7a
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 7 deletions.
diff --git a/src/main/kotlin/me/kpavlov/mocks/statsd/client/StatsDClient.kt b/src/main/kotlin/me/kpavlov/mocks/statsd/client/StatsDClient.kt
@@ -18,12 +18,22 @@ public class StatsDClient(
 ) {
     private val socket: DatagramSocket = DatagramSocket()
 
-    public fun incrementCounter(metric: String) {
-        send("$metric:1|c")
+    @JvmOverloads
+    public fun incrementCounter(metric: String, sampleRate: Double? = null) {
+        if (sampleRate != null) {
+            send("$metric:1|c|@$sampleRate")
+        } else {
+            send("$metric:1|c")
+        }
     }
 
-    public fun time(metric: String, value: Long) {
-        send("$metric:$value|ms")
+    @JvmOverloads
+    public fun time(metric: String, value: Long, sampleRate: Double? = null) {
+        if (sampleRate != null) {
+            send("$metric:$value|ms|@$sampleRate")
+        } else {
+            send("$metric:$value|ms")
+        }
     }
 
     public fun gauge(metric: String, value: Double) {

diff --git a/src/main/kotlin/me/kpavlov/mocks/statsd/server/StatsDServer.kt b/src/main/kotlin/me/kpavlov/mocks/statsd/server/StatsDServer.kt
@@ -62,10 +62,23 @@ public open class StatsDServer(port: Int = DEFAULT_PORT) {
         message.split("\n").forEach(this::handleMetric)
     }
 
+    /**
+     * Split metric:
+     * ```
+     * <metric name>:<value>|c[|@<sample rate>]
+     * ```
+     */
     private fun handleMetric(message: String) {
         val metricData = message.split(":")
         val metricName = metricData[0]
-        val metricValue = metricData[1].split("|")[0].toDouble()
+        val valueParts = metricData[1].split("|")
+        val metricValue = valueParts[0].toDouble()
+        val metricType = valueParts[1]
+        var sampleRate = if (valueParts.size == 3) {
+            valueParts[2].removePrefix("@").toDouble()
+        } else {
+            null
+        }
         metrics.merge(metricName, metricValue, Double::plus)
         logger.debug("Updated value: {} = {}", metricName, metrics[metricName])
     }

diff --git a/src/test/kotlin/me/kpavlov/mocks/statsd/server/BaseStatsDServerTest.kt b/src/test/kotlin/me/kpavlov/mocks/statsd/server/BaseStatsDServerTest.kt
@@ -15,7 +15,7 @@ internal abstract class BaseStatsDServerTest {
     protected lateinit var client: StatsDClient
 
     @Test
-    fun `Server should capture Time`() {
+    fun `Server should capture Timer`() {
         val name = "timeMetric"
         val value = 31L
         client.time(name, value)
@@ -28,6 +28,20 @@ internal abstract class BaseStatsDServerTest {
         statsd.verifyNoMoreCalls(expectedMessage)
     }
 
+    @Test
+    fun `Server should capture Timer with sample rate`() {
+        val name = "sampleTimeMetric"
+        val value = 31L
+        client.time(name, value, 0.1)
+        await untilAsserted {
+            assertThat(statsd.metric(name)).isEqualTo(value.toDouble())
+        }
+        val expectedMessage = "$name:$value|ms|@0.1"
+        assertThat(statsd.calls()).containsOnlyOnce(expectedMessage)
+        statsd.verifyCall(expectedMessage)
+        statsd.verifyNoMoreCalls(expectedMessage)
+    }
+
     @Test
     fun `Server should capture Counter`() {
         val name = "counterMetric"

diff --git a/src/test/kotlin/me/kpavlov/mocks/statsd/server/StatsDServerTest.kt b/src/test/kotlin/me/kpavlov/mocks/statsd/server/StatsDServerTest.kt
@@ -7,7 +7,6 @@ import org.junit.jupiter.api.TestInstance
 
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 internal class StatsDServerTest : BaseStatsDServerTest() {
-
     @BeforeAll
     fun beforeAll() {
         statsd = MockStatsDServer(RANDOM_PORT)