From 0519ca0e088b9c38b4ae6d200ecb6c0b7895df12 Mon Sep 17 00:00:00 2001 From: Stephen Amar Date: Fri, 6 Mar 2026 01:48:28 +0000 Subject: [PATCH] Add --flamegraph flag for folded stack profile output Introduces a FlameGraphProfiler that records Jsonnet call stacks in Brendan Gregg's folded stack format, suitable for generating flame graphs. The profiler hooks into the existing stack depth tracking in the Evaluator, pushing/popping frame names alongside depth checks. Usage: sjsonnet --flamegraph profile.txt input.jsonnet Then: flamegraph.pl profile.txt > profile.svg --- sjsonnet/src-jvm-native/sjsonnet/Config.scala | 6 +++ .../sjsonnet/SjsonnetMainBase.scala | 25 ++++++--- sjsonnet/src/sjsonnet/Evaluator.scala | 53 +++++++++++++------ .../src/sjsonnet/FlameGraphProfiler.scala | 52 ++++++++++++++++++ sjsonnet/src/sjsonnet/Interpreter.scala | 2 +- 5 files changed, 114 insertions(+), 24 deletions(-) create mode 100644 sjsonnet/src/sjsonnet/FlameGraphProfiler.scala diff --git a/sjsonnet/src-jvm-native/sjsonnet/Config.scala b/sjsonnet/src-jvm-native/sjsonnet/Config.scala index 6138aa3d..8f3948da 100644 --- 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 index c32a7023..a0568a6b 100644 --- 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 index 5b1cf0bd..bbd12519 100644 --- a/sjsonnet/src/sjsonnet/Evaluator.scala +++ b/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 @@ -37,8 +38,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] = @@ -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 new file mode 100644 index 00000000..018e7687 --- /dev/null +++ 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 index 7fe4825f..0fe5a58f 100644 --- 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