From dfb19bd8a27bb9a5b9890d9ceaa004c2f9daf97c Mon Sep 17 00:00:00 2001
From: Nikita Nazarov <nikita.nazarov@jetbrains.com>
Date: Mon, 13 Sep 2021 00:12:01 +0300
Subject: [PATCH] Add method that allows IDEA debugger to retrieve enhanced
 stack trace in a JSON format

---
 .../jvm/src/debug/internal/DebugProbesImpl.kt | 35 ++++++++++---
 ...nhanceStackTraceWithTreadDumpAsJsonTest.kt | 49 +++++++++++++++++++
 2 files changed, 78 insertions(+), 6 deletions(-)
 create mode 100644 kotlinx-coroutines-debug/test/EnhanceStackTraceWithTreadDumpAsJsonTest.kt

diff --git a/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt b/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt
index f0fc50e658..e6cddbca1c 100644
--- a/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt
+++ b/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt
@@ -183,7 +183,6 @@ internal object DebugProbesImpl {
      */
     @OptIn(ExperimentalStdlibApi::class)
     public fun dumpCoroutinesInfoAsJsonAndReferences(): Array<Any> {
-        fun Any.toStringWithQuotes() = "\"$this\""
         val coroutinesInfo = dumpCoroutinesInfo()
         val size = coroutinesInfo.size
         val lastObservedThreads = ArrayList<Thread?>(size)
@@ -196,11 +195,11 @@ internal object DebugProbesImpl {
             coroutinesInfoAsJson.add(
                 """
                 {
-                   "name": $name,
-                   "id": ${context[CoroutineId.Key]?.id},
-                   "dispatcher": $dispatcher,
-                   "sequenceNumber": ${info.sequenceNumber},
-                   "state": "${info.state}"
+                    "name": $name,
+                    "id": ${context[CoroutineId.Key]?.id},
+                    "dispatcher": $dispatcher,
+                    "sequenceNumber": ${info.sequenceNumber},
+                    "state": "${info.state}"
                 } 
                 """.trimIndent()
             )
@@ -216,6 +215,30 @@ internal object DebugProbesImpl {
         )
     }
 
+    /*
+     * Internal (JVM-public) method used by IDEA debugger as of 1.6.0-RC.
+     */
+    public fun enhanceStackTraceWithThreadDumpAsJson(info: DebugCoroutineInfo): String {
+        val stackTraceElements = enhanceStackTraceWithThreadDump(info, info.lastObservedStackTrace)
+        val stackTraceElementsInfoAsJson = mutableListOf<String>()
+        for (element in stackTraceElements) {
+            stackTraceElementsInfoAsJson.add(
+                """
+                {
+                    "declaringClass": "${element.className}",
+                    "methodName": "${element.methodName}",
+                    "fileName": ${element.fileName?.toStringWithQuotes()},
+                    "lineNumber": ${element.lineNumber}
+                }
+                """.trimIndent()
+            )
+        }
+
+        return "[${stackTraceElementsInfoAsJson.joinToString()}]"
+    }
+
+    private fun Any.toStringWithQuotes() = "\"$this\""
+
     /*
      * Internal (JVM-public) method used by IDEA debugger as of 1.4-M3.
      */
diff --git a/kotlinx-coroutines-debug/test/EnhanceStackTraceWithTreadDumpAsJsonTest.kt b/kotlinx-coroutines-debug/test/EnhanceStackTraceWithTreadDumpAsJsonTest.kt
new file mode 100644
index 0000000000..fcf9f1a9a9
--- /dev/null
+++ b/kotlinx-coroutines-debug/test/EnhanceStackTraceWithTreadDumpAsJsonTest.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+@file:Suppress("INVISIBLE_MEMBER", "INVISIBLE_REFERENCE")
+package kotlinx.coroutines.debug
+
+import com.google.gson.*
+import kotlinx.coroutines.*
+import kotlinx.coroutines.debug.internal.*
+import org.junit.Test
+import kotlin.test.*
+
+class EnhanceStackTraceWithTreadDumpAsJsonTest : DebugTestBase() {
+    private data class StackTraceElementInfoFromJson(
+        val declaringClass: String,
+        val methodName: String,
+        val fileName: String?,
+        val lineNumber: Int
+    )
+
+    @Test
+    fun testEnhancedStackTraceFormatWithDeferred() = runTest {
+        val deferred = async {
+            suspendingMethod()
+            assertTrue(true)
+        }
+        yield()
+
+        val coroutineInfo = DebugProbesImpl.dumpCoroutinesInfo()
+        assertEquals(coroutineInfo.size, 2)
+        val info = coroutineInfo[1]
+        val enhancedStackTraceAsJsonString = DebugProbesImpl.enhanceStackTraceWithThreadDumpAsJson(info)
+        val enhancedStackTraceFromJson = Gson().fromJson(enhancedStackTraceAsJsonString, Array<StackTraceElementInfoFromJson>::class.java)
+        val enhancedStackTrace = DebugProbesImpl.enhanceStackTraceWithThreadDump(info, info.lastObservedStackTrace)
+        assertEquals(enhancedStackTrace.size, enhancedStackTraceFromJson.size)
+        for ((frame, frameFromJson) in enhancedStackTrace.zip(enhancedStackTraceFromJson)) {
+            assertEquals(frame.className, frameFromJson.declaringClass)
+            assertEquals(frame.methodName, frameFromJson.methodName)
+            assertEquals(frame.fileName, frameFromJson.fileName)
+            assertEquals(frame.lineNumber, frameFromJson.lineNumber)
+        }
+
+        deferred.cancelAndJoin()
+    }
+
+    private suspend fun suspendingMethod() {
+        delay(Long.MAX_VALUE)
+    }
+}