From f577de40faaba383d966a1649f956929dd012091 Mon Sep 17 00:00:00 2001 From: Henning Perl Date: Mon, 22 Jan 2024 12:11:13 +0100 Subject: [PATCH] feat: improve jsonnet performance through caching --- jsonnetsecure/cmd.go | 35 +++++++++++++++++------------- jsonnetsecure/jsonnet.go | 17 +++++++++++---- jsonnetsecure/jsonnet_processvm.go | 23 +++++++++++++++++--- jsonnetsecure/jsonnet_test.go | 19 ++++++++++++++++ jsonnetsecure/provider.go | 4 ++++ 5 files changed, 76 insertions(+), 22 deletions(-) diff --git a/jsonnetsecure/cmd.go b/jsonnetsecure/cmd.go index d8112535..f7c6e28b 100644 --- a/jsonnetsecure/cmd.go +++ b/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(¶ms) 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) +} diff --git a/jsonnetsecure/jsonnet.go b/jsonnetsecure/jsonnet.go index 69efc9fa..072c0edb 100644 --- a/jsonnetsecure/jsonnet.go +++ b/jsonnetsecure/jsonnet.go @@ -10,6 +10,7 @@ import ( "runtime" "testing" + "github.com/dgraph-io/ristretto" "github.com/google/go-jsonnet" ) @@ -31,10 +32,11 @@ 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 { @@ -42,6 +44,7 @@ type ( 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 { diff --git a/jsonnetsecure/jsonnet_processvm.go b/jsonnetsecure/jsonnet_processvm.go index f057156f..1a386470 100644 --- a/jsonnetsecure/jsonnet_processvm.go +++ b/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)) } diff --git a/jsonnetsecure/jsonnet_test.go b/jsonnetsecure/jsonnet_test.go index 1904d055..4b697506 100644 --- a/jsonnetsecure/jsonnet_test.go +++ b/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) diff --git a/jsonnetsecure/provider.go b/jsonnetsecure/provider.go index 5990e421..8763584d 100644 --- a/jsonnetsecure/provider.go +++ b/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 }