From 4cddc2636309c94eda86c3099e12f10cd8cab6ea Mon Sep 17 00:00:00 2001
From: Igor Wiedler <iwiedler@gitlab.com>
Date: Fri, 23 Apr 2021 16:32:24 +0200
Subject: [PATCH] Rudimentary stack profiler

Usage:

  JSONNET_STACK_PROFILE=stacks.out jsonnet main.jsonnet
  cat stacks.out | flamegraph.pl --hash > flamegraph.svg
---
 cmd/jsonnet/cmd.go |  3 +++
 interpreter.go     | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/cmd/jsonnet/cmd.go b/cmd/jsonnet/cmd.go
index ca217abec..a7e9de64f 100644
--- a/cmd/jsonnet/cmd.go
+++ b/cmd/jsonnet/cmd.go
@@ -397,6 +397,9 @@ func main() {
 	cmd.StartCPUProfile()
 	defer cmd.StopCPUProfile()
 
+	jsonnet.StartStackProfile()
+	defer jsonnet.StopStackProfile()
+
 	vm := jsonnet.MakeVM()
 	vm.ErrorFormatter.SetColorFormatter(color.New(color.FgRed).Fprintf)
 
diff --git a/interpreter.go b/interpreter.go
index a2ba1e2f3..b948c5dc0 100644
--- a/interpreter.go
+++ b/interpreter.go
@@ -17,13 +17,18 @@ limitations under the License.
 package jsonnet
 
 import (
+	"bufio"
 	"bytes"
 	"fmt"
 	"io"
+	"log"
 	"math"
+	"math/rand"
+	"os"
 	"reflect"
 	"sort"
 	"strconv"
+	"strings"
 
 	"github.com/google/go-jsonnet/ast"
 	"github.com/google/go-jsonnet/astgen"
@@ -968,7 +973,45 @@ func jsonToValue(i *interpreter, v interface{}) (value, error) {
 	}
 }
 
+var (
+	stackProfileOut   *bufio.Writer
+	stackProfileRatio = 0.01
+)
+
+func StartStackProfile() {
+	var err error
+
+	if os.Getenv("JSONNET_STACK_PROFILE") != "" {
+		stackProfileOutFile, err := os.Create(os.Getenv("JSONNET_STACK_PROFILE"))
+		if err != nil {
+			log.Fatal("could not create stack profile: ", err)
+		}
+		stackProfileOut = bufio.NewWriter(stackProfileOutFile)
+	}
+
+	if os.Getenv("JSONNET_STACK_PROFILE_RATIO") != "" {
+		stackProfileRatio, err = strconv.ParseFloat(os.Getenv("JSONNET_STACK_PROFILE_RATIO"), 64)
+		if err != nil {
+			log.Fatal("could not parse stack profile ratio: ", err)
+		}
+	}
+}
+
+func StopStackProfile() {
+	if stackProfileOut != nil {
+		stackProfileOut.Flush()
+	}
+}
+
 func (i *interpreter) EvalInCleanEnv(env *environment, ast ast.Node, trimmable bool) (value, error) {
+	if stackProfileOut != nil && rand.Float64() < stackProfileRatio {
+		stack := []string{}
+		for _, frame := range i.getCurrentStackTrace() {
+			stack = append(stack, frame.Loc.String()+":"+frame.Name)
+		}
+		fmt.Fprintln(stackProfileOut, strings.Join(stack, ";")+" 1")
+	}
+
 	err := i.newCall(*env, trimmable)
 	if err != nil {
 		return nil, err