databricks · stephenamar-db · Mar 6, 2026 · Mar 6, 2026
diff --git a/sjsonnet/src-jvm-native/sjsonnet/Config.scala b/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

diff --git a/sjsonnet/src-jvm-native/sjsonnet/SjsonnetMainBase.scala b/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
   }
 
   /**

diff --git a/sjsonnet/src/sjsonnet/Evaluator.scala b/sjsonnet/src/sjsonnet/Evaluator.scala
@@ -30,15 +30,32 @@ 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
     if (stackDepth > maxStack)
       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] =
@@ -228,7 +245,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 =
@@ -244,7 +261,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 =
@@ -259,7 +276,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 =
@@ -275,7 +292,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 =
@@ -293,7 +310,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 =
@@ -314,7 +331,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
@@ -323,7 +340,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))
@@ -335,7 +352,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))
@@ -347,7 +364,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(
@@ -361,7 +378,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(
@@ -389,7 +406,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
@@ -502,7 +519,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
@@ -730,7 +747,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()
       }
@@ -921,9 +938,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()
             }
@@ -936,9 +954,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()
             }
@@ -980,7 +999,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)

diff --git a/sjsonnet/src/sjsonnet/FlameGraphProfiler.scala b/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()
+  }
+}
diff --git a/sjsonnet/src/sjsonnet/Interpreter.scala b/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