feat: improve jsonnet performance through caching

jsonnetsecure/cmd.go

@@ -22,22 +22,8 @@ func NewJsonnetCmd() *cobra.Command {
 			if err := params.DecodeFrom(cmd.InOrStdin()); err != nil {
 				return err
 			}
-			vm := MakeSecureVM()
 
-			for _, it := range params.ExtCodes {
-				vm.ExtCode(it.Key, it.Value)
-			}
-			for _, it := range params.ExtVars {
-				vm.ExtVar(it.Key, it.Value)
-			}
-			for _, it := range params.TLACodes {
-				vm.TLACode(it.Key, it.Value)
-			}
-			for _, it := range params.TLAVars {
-				vm.TLAVar(it.Key, it.Value)
-			}
-
-			result, err := vm.EvaluateAnonymousSnippet(params.Filename, params.Snippet)
+			result, err := evaluateJsonnetSnippet(&params)
 			if err != nil {
 				return errors.Wrap(err, "failed to evaluate snippet")
 			}
@@ -52,3 +38,22 @@ func NewJsonnetCmd() *cobra.Command {
 
 	return cmd
 }
+
+func evaluateJsonnetSnippet(params *processParameters) (string, error) {
+	vm := MakeSecureVM()
+
+	for _, it := range params.ExtCodes {
+		vm.ExtCode(it.Key, it.Value)
+	}
+	for _, it := range params.ExtVars {
+		vm.ExtVar(it.Key, it.Value)
+	}
+	for _, it := range params.TLACodes {
+		vm.TLACode(it.Key, it.Value)
+	}
+	for _, it := range params.TLAVars {
+		vm.TLAVar(it.Key, it.Value)
+	}
+
+	return vm.EvaluateAnonymousSnippet(params.Filename, params.Snippet)
+}

jsonnetsecure/jsonnet.go

@@ -10,6 +10,7 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/dgraph-io/ristretto"
 	"github.com/google/go-jsonnet"
 )
 
@@ -31,17 +32,19 @@ type (
 	}
 
 	ProcessVM struct {
-		ctx    context.Context
-		path   string
-		args   []string
-		params processParameters
+		ctx          context.Context
+		path         string
+		args         []string
+		params       processParameters
+		snippetCache *ristretto.Cache
 	}
 
 	vmOptions struct {
 		useProcessVM      bool
 		jsonnetBinaryPath string
 		args              []string
 		ctx               context.Context
+		snippetCache      *ristretto.Cache
 	}
 
 	Option func(o *vmOptions)
@@ -74,6 +77,12 @@ func WithProcessArgs(args ...string) Option {
 	}
 }
 
+func WithSnippetCache(cache *ristretto.Cache) Option {
+	return func(o *vmOptions) {
+		o.snippetCache = cache
+	}
+}
+
 func MakeSecureVM(opts ...Option) VM {
 	options := newVMOptions()
 	for _, o := range opts {

jsonnetsecure/jsonnet_processvm.go

@@ -25,9 +25,10 @@ import (
 
 func NewProcessVM(opts *vmOptions) VM {
 	return &ProcessVM{
-		path: opts.jsonnetBinaryPath,
-		args: opts.args,
-		ctx:  opts.ctx,
+		path:         opts.jsonnetBinaryPath,
+		args:         opts.args,
+		ctx:          opts.ctx,
+		snippetCache: opts.snippetCache,
 	}
 }
 
@@ -36,6 +37,18 @@ func (p *ProcessVM) EvaluateAnonymousSnippet(filename string, snippet string) (_
 	ctx, span := tracer.Start(p.ctx, "jsonnetsecure.ProcessVM.EvaluateAnonymousSnippet", trace.WithAttributes(attribute.String("filename", filename)))
 	defer otelx.End(span, &err)
 
+	if p.snippetCache != nil {
+		if _, ok := p.snippetCache.Get(snippet); ok {
+			_, span := tracer.Start(p.ctx, "jsonnetsecure.ProcessVM.EvaluateAnonymousSnippet.InMemory", trace.WithAttributes(attribute.String("filename", filename)))
+			defer otelx.End(span, &err)
+			result, err := evaluateJsonnetSnippet(&p.params)
+			if err != nil {
+				return "", errors.Wrap(err, "failed to evaluate snippet")
+			}
+			return result, nil
+		}
+	}
+
 	// We retry the process creation, because it sometimes times out.
 	const processVMTimeout = 1 * time.Second
 	return backoff.RetryWithData(func() (_ string, err error) {
@@ -71,6 +84,10 @@ func (p *ProcessVM) EvaluateAnonymousSnippet(filename string, snippet string) (_
 			return "", fmt.Errorf("jsonnetsecure: %w (stdout=%q stderr=%q)", err, stdout.String(), stderr.String())
 		}
 
+		if p.snippetCache != nil {
+			p.snippetCache.Set(snippet, nil, 1)
+		}
+
 		return stdout.String(), nil
 	}, backoff.WithContext(backoff.NewExponentialBackOff(), ctx))
 }

jsonnetsecure/jsonnet_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/dgraph-io/ristretto"
 	"github.com/google/go-jsonnet"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -110,6 +111,24 @@ func TestSecureVM(t *testing.T) {
 		})
 	})
 
+	t.Run("case=caching", func(t *testing.T) {
+		snippet := "{a:1}"
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+		cache, err := ristretto.NewCache(&ristretto.Config{NumCounters: 1000, MaxCost: 1000, BufferItems: 1000})
+		require.NoError(t, err)
+		vm := MakeSecureVM(
+			WithProcessIsolatedVM(ctx),
+			WithSnippetCache(cache),
+			WithJsonnetBinary(testBinary),
+		)
+
+		for i := 0; i < 10; i++ {
+			_, err = vm.EvaluateAnonymousSnippet("test", snippet)
+			require.NoError(t, err)
+		}
+	})
+
 	t.Run("case=process isolation", func(t *testing.T) {
 		snippet := "local f(x) = if x == 0 then [] else [f(x - 1), f(x - 1)]; f(100)"
 		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)

jsonnetsecure/provider.go

@@ -7,6 +7,8 @@ import (
 	"context"
 	"os"
 	"testing"
+
+	"github.com/dgraph-io/ristretto"
 )
 
 type (
@@ -27,6 +29,7 @@ type (
 	// running the current binary with the provided subcommand.
 	DefaultProvider struct {
 		Subcommand string
+		Cache      *ristretto.Cache
 	}
 )
 
@@ -50,5 +53,6 @@ func (p *DefaultProvider) JsonnetVM(ctx context.Context) (VM, error) {
 		WithProcessIsolatedVM(ctx),
 		WithJsonnetBinary(self),
 		WithProcessArgs(p.Subcommand),
+		WithSnippetCache(p.Cache),
 	), nil
 }