Merge pull request #191 from statsig-io/diag_core_api

xinlili-statsig · web-flow · commit f730341524bf · 2023-12-01T10:22:42.000-08:00
Diagnostics core api
diff --git a/src/main/kotlin/com/statsig/sdk/Diagnostics.kt b/src/main/kotlin/com/statsig/sdk/Diagnostics.kt
@@ -6,8 +6,14 @@ const val NANO_IN_MS = 1_000_000.0
 const val MAX_SAMPLING_RATE = 10_000
 internal class Diagnostics(private var isDisabled: Boolean, private var logger: StatsigLogger) {
     var diagnosticsContext: ContextType = ContextType.INITIALIZE
-    private val samplingRates: MutableMap<String, Int> = mutableMapOf("dcs" to 0, "log" to 0, "initialize" to MAX_SAMPLING_RATE, "idlist" to 0)
-    private var markers: DiagnosticsMarkers = mutableMapOf()
+    private val samplingRates: MutableMap<String, Int> = mutableMapOf(
+        "dcs" to 0,
+        "log" to 0,
+        "initialize" to MAX_SAMPLING_RATE,
+        "idlist" to 0,
+        "api_call" to 0,
+    )
+    internal var markers: DiagnosticsMarkers = mutableMapOf()
 
     fun setSamplingRate(rates: Map<String, Int>) {
         rates.forEach { entry ->
@@ -23,10 +29,11 @@ internal class Diagnostics(private var isDisabled: Boolean, private var logger:
         }
     }
 
-    fun markStart(key: KeyType, step: StepType? = null, additionalMarker: Marker? = null) {
+    fun markStart(key: KeyType, step: StepType? = null, context: ContextType? = null, additionalMarker: Marker? = null) {
         if (isDisabled) {
             return
         }
+        val contextType = context ?: diagnosticsContext
         val marker = Marker(key = key, action = ActionType.START, timestamp = System.nanoTime() / NANO_IN_MS, step = step)
         when (key) {
             KeyType.GET_ID_LIST -> {
@@ -38,13 +45,20 @@ internal class Diagnostics(private var isDisabled: Boolean, private var logger:
                 }
             }
         }
-        this.addMarker(marker)
+        when (contextType) {
+            ContextType.API_CALL -> {
+                marker.configName = additionalMarker?.configName
+                marker.markerID = additionalMarker?.markerID
+            }
+        }
+        this.addMarker(marker, contextType)
     }
 
-    fun markEnd(key: KeyType, success: Boolean, step: StepType? = null, additionalMarker: Marker? = null) {
+    fun markEnd(key: KeyType, success: Boolean, step: StepType? = null, context: ContextType? = null, additionalMarker: Marker? = null) {
         if (isDisabled) {
             return
         }
+        val contextType = context ?: diagnosticsContext
         val marker = Marker(key = key, action = ActionType.END, success = success, timestamp = System.nanoTime() / NANO_IN_MS, step = step)
         when (key) {
             KeyType.DOWNLOAD_CONFIG_SPECS -> {
@@ -72,24 +86,31 @@ internal class Diagnostics(private var isDisabled: Boolean, private var logger:
                 marker.reason = additionalMarker?.reason
             }
         }
-        this.addMarker(marker)
+        when (contextType) {
+            ContextType.API_CALL -> {
+                marker.configName = additionalMarker?.configName
+                marker.markerID = additionalMarker?.markerID
+            }
+        }
+        this.addMarker(marker, contextType)
     }
 
-    private fun shouldLogDiagnostics(context: ContextType): Boolean {
+    internal fun shouldLogDiagnostics(context: ContextType): Boolean {
         val samplingKey: String =
             when (context) {
                 ContextType.CONFIG_SYNC -> "dcs"
                 ContextType.INITIALIZE -> "initialize"
+                ContextType.API_CALL -> "api_call"
             }
         val rand = Math.random() * MAX_SAMPLING_RATE
         return samplingRates[samplingKey] ?: 0 > rand
     }
 
-    private fun addMarker(marker: Marker) {
-        if (this.markers[diagnosticsContext] == null) {
-            this.markers[diagnosticsContext] = mutableListOf()
+    private fun addMarker(marker: Marker, context: ContextType) {
+        if (this.markers[context] == null) {
+            this.markers[context] = mutableListOf()
         }
-        this.markers[diagnosticsContext]?.add(marker)
+        this.markers[context]?.add(marker)
         this.markers.values
     }
 
diff --git a/src/main/kotlin/com/statsig/sdk/DiagnosticsData.kt b/src/main/kotlin/com/statsig/sdk/DiagnosticsData.kt
@@ -3,6 +3,7 @@ package com.statsig.sdk
 import com.google.gson.annotations.SerializedName
 
 data class Marker(
+    @SerializedName("markerID") var markerID: String? = null,
     @SerializedName("key") val key: KeyType? = null,
     @SerializedName("action") val action: ActionType? = null,
     @SerializedName("timestamp") val timestamp: Double? = null,
@@ -13,6 +14,7 @@ data class Marker(
     @SerializedName("idListCount") var idListCount: Int? = null,
     @SerializedName("reason") var reason: String? = null,
     @SerializedName("sdkRegion") var sdkRegion: String? = null,
+    @SerializedName("configName") var configName: String? = null,
 )
 
 enum class ContextType {
@@ -21,6 +23,9 @@ enum class ContextType {
 
     @SerializedName("config_sync")
     CONFIG_SYNC,
+
+    @SerializedName("api_call")
+    API_CALL,
 }
 
 enum class KeyType {
@@ -38,6 +43,35 @@ enum class KeyType {
 
     @SerializedName("overall")
     OVERALL,
+
+    @SerializedName("check_gate")
+    CHECK_GATE,
+
+    @SerializedName("get_config")
+    GET_CONFIG,
+
+    @SerializedName("get_experiment")
+    GET_EXPERIMENT,
+
+    @SerializedName("get_layer")
+    GET_LAYER, ;
+
+    companion object {
+        fun convertFromString(value: String): KeyType? {
+            return when (value) {
+                in "checkGate", "checkGateWithExposureLoggingDisabled" ->
+                    CHECK_GATE
+                in "getExperiment", "getExperimentWithExposureLoggingDisabled" ->
+                    GET_EXPERIMENT
+                in "getConfig", "getConfigWithExposureLoggingDisabled" ->
+                    GET_CONFIG
+                in "getLayer", "getLayerWithExposureLoggingDisabled" ->
+                    GET_LAYER
+                else ->
+                    null
+            }
+        }
+    }
 }
 
 enum class StepType {
diff --git a/src/main/kotlin/com/statsig/sdk/ErrorBoundary.kt b/src/main/kotlin/com/statsig/sdk/ErrorBoundary.kt
@@ -13,8 +13,9 @@ internal class ErrorBoundary(private val apiKey: String, private val options: St
     internal var uri = URI("https://statsigapi.net/v1/sdk_exception")
     private val seen = HashSet<String>()
     private val maxInfoLength = 3000
-
     private val client = OkHttpClient()
+    internal var diagnostics: Diagnostics? = null
+
     private companion object {
         val MEDIA_TYPE = "application/json; charset=utf-8".toMediaType()
     }
@@ -34,10 +35,17 @@ internal class ErrorBoundary(private val apiKey: String, private val options: St
     }
 
     suspend fun <T> capture(tag: String, task: suspend () -> T, recover: suspend () -> T, configName: String? = null): T {
+        var markerID: String? = null
+        var keyType: KeyType? = null
         return try {
-            task()
+            keyType = KeyType.convertFromString(tag)
+            markerID = markStart(keyType, configName)
+            val result = task()
+            markEnd(keyType, true, configName, markerID)
+            return result
         } catch (ex: Throwable) {
             onException(tag, ex, configName)
+            markEnd(keyType, false, configName, markerID)
             recover()
         }
     }
@@ -99,4 +107,21 @@ internal class ErrorBoundary(private val apiKey: String, private val options: St
 
         logException(tag, ex, configName)
     }
+
+    private fun markStart(keyType: KeyType?, configName: String?): String? {
+        if (diagnostics == null || keyType == null) {
+            return null
+        }
+        val markerID = keyType.name + "_" + (diagnostics?.markers?.get(ContextType.API_CALL)?.count() ?: 0)
+
+        diagnostics?.markStart(keyType, context = ContextType.API_CALL, additionalMarker = Marker(markerID = markerID, configName = configName))
+        return markerID
+    }
+
+    private fun markEnd(keyType: KeyType?, success: Boolean, configName: String?, markerID: String?) {
+        if (diagnostics == null || keyType == null) {
+            return
+        }
+        diagnostics?.markEnd(keyType, success, context = ContextType.API_CALL, additionalMarker = Marker(markerID = markerID, configName = configName))
+    }
 }
diff --git a/src/main/kotlin/com/statsig/sdk/SpecStore.kt b/src/main/kotlin/com/statsig/sdk/SpecStore.kt
@@ -141,7 +141,7 @@ internal class SpecStore constructor(
         }
         val body = response.body ?: return
         val jsonResponse = gson.fromJson<Map<String, IDList>>(body.string()) ?: return
-        diagnostics.markStart(KeyType.GET_ID_LIST_SOURCES, StepType.PROCESS, Marker(idListCount = jsonResponse.size))
+        diagnostics.markStart(KeyType.GET_ID_LIST_SOURCES, StepType.PROCESS, additionalMarker = Marker(idListCount = jsonResponse.size))
         val tasks = mutableListOf<Job>()
 
         for ((name, serverList) in jsonResponse) {
diff --git a/src/main/kotlin/com/statsig/sdk/StatsigLogger.kt b/src/main/kotlin/com/statsig/sdk/StatsigLogger.kt
@@ -47,7 +47,7 @@ internal class StatsigLogger(
         }
     }
     private val gson = GsonBuilder().setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE).create()
-
+    internal var diagnostics: Diagnostics? = null
     fun log(event: StatsigEvent) {
         events.add(event)
 
@@ -131,6 +131,9 @@ internal class StatsigLogger(
     }
 
     fun logDiagnostics(context: ContextType, markers: Collection<Marker>) {
+        if (markers.isEmpty()) {
+            return
+        }
         val event = StatsigEvent(DIAGNOSTICS_EVENT)
         event.eventMetadata = mapOf(
             "context" to context.toString().lowercase(),
@@ -140,8 +143,25 @@ internal class StatsigLogger(
         log(event)
     }
 
+    fun addAPICallDiagnostics() {
+        val markers = diagnostics?.markers?.get(ContextType.API_CALL)
+        if (markers.isNullOrEmpty() ||
+            diagnostics?.shouldLogDiagnostics(ContextType.API_CALL) != true
+        ) {
+            return
+        }
+        val event = StatsigEvent(DIAGNOSTICS_EVENT)
+        event.eventMetadata = mapOf(
+            "context" to "api_call",
+            "markers" to gson.toJson(markers),
+            "setupOptions" to gson.toJson(statsigOptions.getLoggingCopy()),
+        )
+        events.add(event)
+    }
+
     suspend fun flush() {
         withContext(Dispatchers.IO) {
+            addAPICallDiagnostics()
             if (events.size() == 0) {
                 return@withContext
             }
diff --git a/src/main/kotlin/com/statsig/sdk/StatsigServer.kt b/src/main/kotlin/com/statsig/sdk/StatsigServer.kt
@@ -327,7 +327,10 @@ private class StatsigServerImpl() :
             return DynamicConfig.empty(experimentName)
         }
         return this.errorBoundary.capture("getExperiment", {
-            return@capture getConfig(user, experimentName)
+            val normalizedUser = normalizeUser(user)
+            val result = getConfigImpl(user, experimentName)
+            logConfigImpl(normalizedUser, experimentName, result)
+            return@capture getDynamicConfigFromEvalResult(result, user, experimentName)
         }, {
             return@capture DynamicConfig.empty(experimentName)
         }, configName = experimentName)
@@ -784,6 +787,8 @@ private class StatsigServerImpl() :
 
     private fun setupAndStartDiagnostics() {
         diagnostics = Diagnostics(options.disableDiagnostics, logger)
+        errorBoundary.diagnostics = diagnostics
+        logger.diagnostics = diagnostics
         diagnostics.markStart(KeyType.OVERALL)
     }
 
diff --git a/src/test/java/com/statsig/sdk/DiagnosticsTest.kt b/src/test/java/com/statsig/sdk/DiagnosticsTest.kt
@@ -59,9 +59,98 @@ class DiagnosticsTest {
         assertEquals(Gson().toJson(options.getLoggingCopy()), diagnosticsEvent!!.eventMetadata!!["statsigOptions"])
     }
 
+    @Test
+    fun testCoreAPI() = runBlocking {
+        val downloadConfigSpecsResponseWithSampling = StringBuilder(downloadConfigSpecsResponse).insert(downloadConfigSpecsResponse.length - 2, ",\n \"diagnostics\": {\"initialize\": \"0\", \"api_call\": \"10000\"}").toString()
+        setupWebServer(downloadConfigSpecsResponseWithSampling)
+        driver.initializeAsync("secret-testcase", options).get()
+        val user = StatsigUser("testUser")
+        driver.checkGate(user, "always_on_gate")
+        driver.getConfig(user, "test_config")
+        driver.getExperiment(user, "sample_experiment")
+        driver.getLayer(user, "a_layer")
+        driver.shutdown()
+        val events = TestUtil.captureEvents(eventLogInputCompletable)
+        val diagnosticsEvent = events.find { it.eventName == "statsig::diagnostics" }
+        val markers: List<Marker> = gson.fromJson(diagnosticsEvent!!.eventMetadata!!["markers"], object : TypeToken<List<Marker>>() {}.type)
+        assertEquals(8, markers.size)
+        verifyMarker(
+            markers[0],
+            Marker(
+                markerID = "CHECK_GATE_0",
+                action = ActionType.START,
+                configName = "always_on_gate",
+                key = KeyType.CHECK_GATE,
+            ),
+        )
+        verifyMarker(
+            markers[1],
+            Marker(
+                markerID = "CHECK_GATE_0",
+                action = ActionType.END,
+                configName = "always_on_gate",
+                key = KeyType.CHECK_GATE,
+            ),
+        )
+        verifyMarker(
+            markers[2],
+            Marker(
+                markerID = "GET_CONFIG_2",
+                action = ActionType.START,
+                configName = "test_config",
+                key = KeyType.GET_CONFIG,
+            ),
+        )
+        verifyMarker(
+            markers[3],
+            Marker(
+                markerID = "GET_CONFIG_2",
+                action = ActionType.END,
+                configName = "test_config",
+                key = KeyType.GET_CONFIG,
+            ),
+        )
+        verifyMarker(
+            markers[4],
+            Marker(
+                markerID = "GET_EXPERIMENT_4",
+                action = ActionType.START,
+                configName = "sample_experiment",
+                key = KeyType.GET_EXPERIMENT,
+            ),
+        )
+        verifyMarker(
+            markers[5],
+            Marker(
+                markerID = "GET_EXPERIMENT_4",
+                action = ActionType.END,
+                configName = "sample_experiment",
+                key = KeyType.GET_EXPERIMENT,
+            ),
+        )
+        verifyMarker(
+            markers[6],
+            Marker(
+                markerID = "GET_LAYER_6",
+                action = ActionType.START,
+                configName = "a_layer",
+                key = KeyType.GET_LAYER,
+            ),
+        )
+        verifyMarker(
+            markers[7],
+            Marker(
+                markerID = "GET_LAYER_6",
+                action = ActionType.END,
+                configName = "a_layer",
+                key = KeyType.GET_LAYER,
+            ),
+        )
+    }
+
     @Test
     fun testSamping() = runBlocking {
-        val downloadConfigSpecsResponseWithSampling = StringBuilder(downloadConfigSpecsResponse).insert(downloadConfigSpecsResponse.length - 2, ",\n \"diagnostics\": {\"initialize\": \"0\"}").toString()
+        val downloadConfigSpecsResponseWithSampling = StringBuilder(downloadConfigSpecsResponse).insert(downloadConfigSpecsResponse.length - 2, ",\n \"diagnostics\": {\"initialize\": \"0\", \"api_call\": \"0\"}").toString()
         setupWebServer(downloadConfigSpecsResponseWithSampling)
         driver.initializeAsync("secret-testcase", options).get()
         driver.shutdown()
@@ -99,5 +188,7 @@ class DiagnosticsTest {
         Assert.assertEquals(expected.key, actual.key)
         Assert.assertEquals(expected.action, actual.action)
         Assert.assertEquals(expected.step, actual.step)
+        Assert.assertEquals(expected.configName, actual.configName)
+        Assert.assertEquals(expected.markerID, actual.markerID)
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -141,7 +141,7 @@ internal class SpecStore constructor(`
`141`	`141`	`}`
`142`	`142`	`val body = response.body ?: return`
`143`	`143`	`val jsonResponse = gson.fromJson<Map<String, IDList>>(body.string()) ?: return`
`144`		`- diagnostics.markStart(KeyType.GET_ID_LIST_SOURCES, StepType.PROCESS, Marker(idListCount = jsonResponse.size))`
	`144`	`+ diagnostics.markStart(KeyType.GET_ID_LIST_SOURCES, StepType.PROCESS, additionalMarker = Marker(idListCount = jsonResponse.size))`
`145`	`145`	`val tasks = mutableListOf<Job>()`
`146`	`146`
`147`	`147`	`for ((name, serverList) in jsonResponse) {`