Add --flamegraph flag for folded stack profile output (#640)

## Summary
- Adds a `--flamegraph <file>` CLI flag that writes a flame graph
profile in Brendan Gregg's folded stack format
- Introduces `FlameGraphProfiler` that hooks into the existing
`Evaluator` stack depth tracking, pushing/popping frame names at every
`Apply`, `ApplyBuiltin`, import, comprehension, and object field
evaluation
- Profile output can be fed directly into
[FlameGraph](https://github.com/brendangregg/FlameGraph) tools:
`flamegraph.pl profile.txt > profile.svg`

## Test plan
- [x] `sjsonnet.jvm[3.3.7].test` passes (140/140)
- [x] `sjsonnet.jvm[2.13.18].compile` passes
- [ ] Manual testing: `sjsonnet --flamegraph /tmp/profile.txt
<input>.jsonnet` produces valid folded stack output

Author: stephenamar-db

sjsonnet/src-jvm-native/sjsonnet/Config.scala

@@ -167,6 +167,12 @@ final case class Config(
       doc = "Number of allowed stack frames (default 500)"
     )
     maxStack: Int = 500,
+    @arg(
+      name = "flamegraph",
+      doc =
+        "Write a flame graph profile in folded stack format to the given file. Use with https://github.com/brendangregg/FlameGraph"
+    )
+    flamegraph: Option[String] = None,
     @arg(
       doc = "The jsonnet file you wish to evaluate",
       positional = true

sjsonnet/src-jvm-native/sjsonnet/SjsonnetMainBase.scala

@@ -169,7 +169,8 @@ object SjsonnetMainBase {
         },
         warn,
         std,
-        debugStats = debugStats
+        debugStats = debugStats,
+        flamegraphFile = config.flamegraph
       )
       res <- {
         if (hasWarnings && config.fatalWarnings.value) Left("")
@@ -318,7 +319,8 @@ object SjsonnetMainBase {
       warnLogger: Evaluator.Logger,
       std: Val.Obj,
       evaluatorOverride: Option[Evaluator] = None,
-      debugStats: DebugStats = null): Either[String, String] = {
+      debugStats: DebugStats = null,
+      flamegraphFile: Option[String] = None): Either[String, String] = {
 
     val (jsonnetCode, path) =
       if (config.exec.value) (file, wd / Util.wrapInLessThanGreaterThan("exec"))
@@ -348,6 +350,7 @@ object SjsonnetMainBase {
     )
 
     var currentPos: Position = null
+    var profiler: FlameGraphProfiler = null
     val interp = new Interpreter(
       queryExtVar = (key: String) => extBinding.get(key).map(ExternalVariable.code),
       queryTlaVar = (key: String) => tlaBinding.get(key).map(ExternalVariable.code),
@@ -365,13 +368,19 @@ object SjsonnetMainBase {
           resolver: CachedResolver,
           extVars: String => Option[Expr],
           wd: Path,
-          settings: Settings): Evaluator =
-        evaluatorOverride.getOrElse(
+          settings: Settings): Evaluator = {
+        val ev = evaluatorOverride.getOrElse(
           super.createEvaluator(resolver, extVars, wd, settings)
         )
+        if (flamegraphFile.isDefined) {
+          profiler = new FlameGraphProfiler
+          ev.flameGraphProfiler = profiler
+        }
+        ev
+      }
     }
 
-    (config.multi, config.yamlStream.value) match {
+    val result = (config.multi, config.yamlStream.value) match {
       case (Some(multiPath), _) =>
         val trailingNewline = !config.noTrailingNewline.value
         interp.interpret(jsonnetCode, OsPath(path)).flatMap {
@@ -437,8 +446,12 @@ object SjsonnetMainBase {
           case _ => renderNormal(config, interp, jsonnetCode, path, wd, () => currentPos)
         }
       case _ => renderNormal(config, interp, jsonnetCode, path, wd, () => currentPos)
-
     }
+
+    if (profiler != null)
+      flamegraphFile.foreach(profiler.writeTo)
+
+    result
   }
 
   /**

sjsonnet/src/sjsonnet/Evaluator.scala

@@ -30,6 +30,7 @@ class Evaluator(
 
   private[this] var stackDepth: Int = 0
   private[this] val maxStack: Int = settings.maxStack
+  private[sjsonnet] var flameGraphProfiler: FlameGraphProfiler = _
 
   @inline private[sjsonnet] final def checkStackDepth(pos: Position): Unit = {
     stackDepth += 1
@@ -39,8 +40,24 @@ class Evaluator(
       Error.fail("Max stack frames exceeded.", pos)
   }
 
-  @inline private[sjsonnet] final def decrementStackDepth(): Unit =
+  @inline private[sjsonnet] final def checkStackDepth(pos: Position, expr: Expr): Unit = {
+    stackDepth += 1
+    if (flameGraphProfiler != null) flameGraphProfiler.push(expr.exprErrorString)
+    if (stackDepth > maxStack)
+      Error.fail("Max stack frames exceeded.", pos)
+  }
+
+  @inline private[sjsonnet] final def checkStackDepth(pos: Position, name: String): Unit = {
+    stackDepth += 1
+    if (flameGraphProfiler != null) flameGraphProfiler.push(name)
+    if (stackDepth > maxStack)
+      Error.fail("Max stack frames exceeded.", pos)
+  }
+
+  @inline private[sjsonnet] final def decrementStackDepth(): Unit = {
     stackDepth -= 1
+    if (flameGraphProfiler != null) flameGraphProfiler.pop()
+  }
 
   def materialize(v: Val): Value = Materializer.apply(v)
   val cachedImports: collection.mutable.HashMap[Path, Val] =
@@ -230,7 +247,7 @@ class Evaluator(
    */
   protected def visitApply(e: Apply)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.functionCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
@@ -246,7 +263,7 @@ class Evaluator(
 
   protected def visitApply0(e: Apply0)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.functionCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
@@ -261,7 +278,7 @@ class Evaluator(
 
   protected def visitApply1(e: Apply1)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.functionCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
@@ -277,7 +294,7 @@ class Evaluator(
 
   protected def visitApply2(e: Apply2)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.functionCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
@@ -295,7 +312,7 @@ class Evaluator(
 
   protected def visitApply3(e: Apply3)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.functionCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       val lhs = visitExpr(e.value)
       implicit val tailstrictMode: TailstrictMode =
@@ -316,7 +333,7 @@ class Evaluator(
 
   protected def visitApplyBuiltin0(e: ApplyBuiltin0): Val = {
     if (debugStats != null) debugStats.builtinCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       val result = e.func.evalRhs(this, e.pos)
       if (e.tailstrict) TailCall.resolve(result) else result
@@ -325,7 +342,7 @@ class Evaluator(
 
   protected def visitApplyBuiltin1(e: ApplyBuiltin1)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.builtinCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       if (e.tailstrict) {
         TailCall.resolve(e.func.evalRhs(visitExpr(e.a1), this, e.pos))
@@ -337,7 +354,7 @@ class Evaluator(
 
   protected def visitApplyBuiltin2(e: ApplyBuiltin2)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.builtinCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       if (e.tailstrict) {
         TailCall.resolve(e.func.evalRhs(visitExpr(e.a1), visitExpr(e.a2), this, e.pos))
@@ -349,7 +366,7 @@ class Evaluator(
 
   protected def visitApplyBuiltin3(e: ApplyBuiltin3)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.builtinCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       if (e.tailstrict) {
         TailCall.resolve(
@@ -363,7 +380,7 @@ class Evaluator(
 
   protected def visitApplyBuiltin4(e: ApplyBuiltin4)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.builtinCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       if (e.tailstrict) {
         TailCall.resolve(
@@ -391,7 +408,7 @@ class Evaluator(
 
   protected def visitApplyBuiltin(e: ApplyBuiltin)(implicit scope: ValScope): Val = {
     if (debugStats != null) debugStats.builtinCalls += 1
-    checkStackDepth(e.pos)
+    checkStackDepth(e.pos, e)
     try {
       val arr = new Array[Eval](e.argExprs.length)
       var idx = 0
@@ -504,7 +521,7 @@ class Evaluator(
     cachedImports.getOrElseUpdate(
       p, {
         if (debugStats != null) debugStats.importCalls += 1
-        checkStackDepth(e.pos)
+        checkStackDepth(e.pos, e)
         try {
           val doc = resolver.parse(p, str) match {
             case Right((expr, _)) => expr
@@ -732,7 +749,7 @@ class Evaluator(
       def evalRhs(vs: ValScope, es: EvalScope, fs: FileScope, pos: Position): Val =
         visitExprWithTailCallSupport(rhs)(vs)
       override def evalDefault(expr: Expr, vs: ValScope, es: EvalScope): Val = {
-        checkStackDepth(expr.pos)
+        checkStackDepth(expr.pos, "default")
         try visitExpr(expr)(vs)
         finally decrementStackDepth()
       }
@@ -923,9 +940,10 @@ class Evaluator(
       case Member.Field(offset, fieldName, plus, null, sep, rhs) =>
         val k = visitFieldName(fieldName, offset)
         if (k != null) {
+          val fieldKey = k
           val v = new Val.Obj.Member(plus, sep) {
             def invoke(self: Val.Obj, sup: Val.Obj, fs: FileScope, ev: EvalScope): Val = {
-              checkStackDepth(rhs.pos)
+              checkStackDepth(rhs.pos, fieldKey)
               try visitExpr(rhs)(makeNewScope(self, sup))
               finally decrementStackDepth()
             }
@@ -938,9 +956,10 @@ class Evaluator(
       case Member.Field(offset, fieldName, false, argSpec, sep, rhs) =>
         val k = visitFieldName(fieldName, offset)
         if (k != null) {
+          val fieldKey = k
           val v = new Val.Obj.Member(false, sep) {
             def invoke(self: Val.Obj, sup: Val.Obj, fs: FileScope, ev: EvalScope): Val = {
-              checkStackDepth(rhs.pos)
+              checkStackDepth(rhs.pos, fieldKey)
               try visitMethod(rhs, argSpec, offset)(makeNewScope(self, sup))
               finally decrementStackDepth()
             }
@@ -982,7 +1001,7 @@ class Evaluator(
             k,
             new Val.Obj.Member(e.plus, Visibility.Normal, deprecatedSkipAsserts = true) {
               def invoke(self: Val.Obj, sup: Val.Obj, fs: FileScope, ev: EvalScope): Val = {
-                checkStackDepth(e.value.pos)
+                checkStackDepth(e.value.pos, "object comprehension")
                 try {
                   lazy val newScope: ValScope = s.extend(newBindings, self, sup)
                   lazy val newBindings = visitBindings(binds, newScope)

sjsonnet/src/sjsonnet/FlameGraphProfiler.scala

@@ -0,0 +1,52 @@
+package sjsonnet
+
+import java.io.{BufferedWriter, FileWriter}
+
+/**
+ * Collects stack samples during Jsonnet evaluation and writes them in Brendan Gregg's folded stack
+ * format, suitable for generating flame graphs with https://github.com/brendangregg/FlameGraph.
+ *
+ * Each call to [[push]] records a new frame on the current stack. Each call to [[pop]] removes the
+ * top frame. A sample (incrementing the count for the current stack) is taken on every [[push]], so
+ * deeper call trees contribute proportionally more samples.
+ */
+final class FlameGraphProfiler {
+  private val stack = new java.util.ArrayDeque[String]()
+  private val counts = new java.util.HashMap[String, java.lang.Long]()
+
+  def push(name: String): Unit = {
+    stack.push(name)
+    val key = foldedStack()
+    val prev = counts.get(key)
+    counts.put(key, if (prev == null) 1L else prev + 1L)
+  }
+
+  def pop(): Unit =
+    if (!stack.isEmpty) stack.pop()
+
+  private def foldedStack(): String = {
+    val sb = new StringBuilder
+    val it = stack.descendingIterator()
+    var first = true
+    while (it.hasNext) {
+      if (!first) sb.append(';')
+      sb.append(it.next())
+      first = false
+    }
+    sb.toString
+  }
+
+  def writeTo(path: String): Unit = {
+    val w = new BufferedWriter(new FileWriter(path))
+    try {
+      val it = counts.entrySet().iterator()
+      while (it.hasNext) {
+        val e = it.next()
+        w.write(e.getKey)
+        w.write(' ')
+        w.write(e.getValue.toString)
+        w.newLine()
+      }
+    } finally w.close()
+  }
+}

sjsonnet/src/sjsonnet/Interpreter.scala

@@ -251,7 +251,7 @@ class Interpreter(
               f.evalRhs(vs, es, fs, pos)
 
             override def evalDefault(expr: Expr, vs: ValScope, es: EvalScope): Val = {
-              evaluator.checkStackDepth(expr.pos)
+              evaluator.checkStackDepth(expr.pos, "default")
               try
                 evaluator.visitExpr(expr)(
                   if (tlaExpressions.exists(_ eq expr)) ValScope.empty else vs