Hide the Jsonnet Go implementation behind an interface

Tanka currently evals jsonnet using the Go native code.
However, some other implementations have come up in the past years that could be worth using (ex: https://github.com/CertainLach/jrsonnet, which is much faster)

In this PR is the first step: I create an interface where all the jsonnet eval code happens. The Go Jsonnet implementation is now hidden behind this interface.

The setting can either be passed as a global flag or as an env spec attribute to be used when exporting (`spec.exportJsonnetImplementation`)

cmd/tk/flags.go

@@ -58,12 +58,14 @@ func labelSelectorFlag(fs *pflag.FlagSet) func() labels.Selector {
 func jsonnetFlags(fs *pflag.FlagSet) func() tanka.JsonnetOpts {
 	getExtCode, getTLACode := cliCodeParser(fs)
 	maxStack := fs.Int("max-stack", 0, "Jsonnet VM max stack. The default value is the value set in the go-jsonnet library. Increase this if you get: max stack frames exceeded")
+	jsonnetImplementation := fs.String("jsonnet-implementation", "go", "Only go is supported for now.")
 
 	return func() tanka.JsonnetOpts {
 		return tanka.JsonnetOpts{
-			MaxStack: *maxStack,
-			ExtCode:  getExtCode(),
-			TLACode:  getTLACode(),
+			MaxStack:              *maxStack,
+			ExtCode:               getExtCode(),
+			TLACode:               getTLACode(),
+			JsonnetImplementation: *jsonnetImplementation,
 		}
 	}
 }

pkg/jsonnet/eval.go

@@ -9,8 +9,10 @@ import (
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
 
+	"github.com/grafana/tanka/pkg/jsonnet/implementation"
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/goimpl"
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/types"
 	"github.com/grafana/tanka/pkg/jsonnet/jpath"
-	"github.com/grafana/tanka/pkg/jsonnet/native"
 )
 
 // Modifier allows to set optional parameters on the Jsonnet VM.
@@ -31,12 +33,13 @@ func (i *InjectedCode) Set(key, value string) {
 
 // Opts are additional properties for the Jsonnet VM
 type Opts struct {
-	MaxStack    int
-	ExtCode     InjectedCode
-	TLACode     InjectedCode
-	ImportPaths []string
-	EvalScript  string
-	CachePath   string
+	JsonnetImplementation string
+	MaxStack              int
+	ExtCode               InjectedCode
+	TLACode               InjectedCode
+	ImportPaths           []string
+	EvalScript            string
+	CachePath             string
 
 	CachePathRegexes []*regexp.Regexp
 }
@@ -75,37 +78,11 @@ func (o Opts) Clone() Opts {
 	}
 }
 
-// MakeVM returns a Jsonnet VM with some extensions of Tanka, including:
-// - extended importer
-// - extCode and tlaCode applied
-// - native functions registered
-func MakeVM(opts Opts) *jsonnet.VM {
-	vm := jsonnet.MakeVM()
-	vm.Importer(NewExtendedImporter(opts.ImportPaths))
-
-	for k, v := range opts.ExtCode {
-		vm.ExtCode(k, v)
-	}
-	for k, v := range opts.TLACode {
-		vm.TLACode(k, v)
-	}
-
-	for _, nf := range native.Funcs() {
-		vm.NativeFunction(nf)
-	}
-
-	if opts.MaxStack > 0 {
-		vm.MaxStack = opts.MaxStack
-	}
-
-	return vm
-}
-
 // EvaluateFile evaluates the Jsonnet code in the given file and returns the
 // result in JSON form. It disregards opts.ImportPaths in favor of automatically
 // resolving these according to the specified file.
 func EvaluateFile(jsonnetFile string, opts Opts) (string, error) {
-	evalFunc := func(vm *jsonnet.VM) (string, error) {
+	evalFunc := func(vm types.JsonnetVM) (string, error) {
 		return vm.EvaluateFile(jsonnetFile)
 	}
 	data, err := os.ReadFile(jsonnetFile)
@@ -119,13 +96,13 @@ func EvaluateFile(jsonnetFile string, opts Opts) (string, error) {
 // If cache options are given, a hash from the data will be computed and
 // the resulting string will be cached for future retrieval
 func Evaluate(path, data string, opts Opts) (string, error) {
-	evalFunc := func(vm *jsonnet.VM) (string, error) {
+	evalFunc := func(vm types.JsonnetVM) (string, error) {
 		return vm.EvaluateAnonymousSnippet(path, data)
 	}
 	return evaluateSnippet(evalFunc, path, data, opts)
 }
 
-type evalFunc func(vm *jsonnet.VM) (string, error)
+type evalFunc func(vm types.JsonnetVM) (string, error)
 
 func evaluateSnippet(evalFunc evalFunc, path, data string, opts Opts) (string, error) {
 	var cache *FileEvalCache
@@ -134,17 +111,22 @@ func evaluateSnippet(evalFunc evalFunc, path, data string, opts Opts) (string, e
 	}
 
 	// Create VM
+	jsonnetImpl, err := implementation.Get(opts.JsonnetImplementation)
+	if err != nil {
+		return "", err
+	}
 	jpath, _, _, err := jpath.Resolve(path, false)
 	if err != nil {
 		return "", errors.Wrap(err, "resolving import paths")
 	}
 	opts.ImportPaths = jpath
-	vm := MakeVM(opts)
+	vm := jsonnetImpl.MakeVM(opts.ImportPaths, opts.ExtCode, opts.TLACode, opts.MaxStack)
+	importVM := goimpl.MakeRawVM(opts.ImportPaths, opts.ExtCode, opts.TLACode, opts.MaxStack) // TODO: use interface
 
 	var hash string
 	if cache != nil {
 		startTime := time.Now()
-		if hash, err = getSnippetHash(vm, path, data); err != nil {
+		if hash, err = getSnippetHash(importVM, path, data); err != nil {
 			return "", err
 		}
 		cacheLog := log.Debug().Str("path", path).Str("hash", hash).Dur("duration_ms", time.Since(startTime))

pkg/jsonnet/implementation/goimpl/impl.go

@@ -0,0 +1,26 @@
+package goimpl
+
+import (
+	"github.com/google/go-jsonnet"
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/types"
+)
+
+type JsonnetGoVM struct {
+	vm *jsonnet.VM
+}
+
+func (vm *JsonnetGoVM) EvaluateAnonymousSnippet(filename, snippet string) (string, error) {
+	return vm.vm.EvaluateAnonymousSnippet(filename, snippet)
+}
+
+func (vm *JsonnetGoVM) EvaluateFile(filename string) (string, error) {
+	return vm.vm.EvaluateFile(filename)
+}
+
+type JsonnetGoImplementation struct{}
+
+func (i *JsonnetGoImplementation) MakeVM(importPaths []string, extCode map[string]string, tlaCode map[string]string, maxStack int) types.JsonnetVM {
+	return &JsonnetGoVM{
+		vm: MakeRawVM(importPaths, extCode, tlaCode, maxStack),
+	}
+}

pkg/jsonnet/importer.go → pkg/jsonnet/implementation/goimpl/importer.go

@@ -1,4 +1,4 @@
-package jsonnet
+package goimpl
 
 import (
 	"path/filepath"
@@ -8,10 +8,10 @@ import (
 
 const locationInternal = "<internal>"
 
-// ExtendedImporter wraps jsonnet.FileImporter to add additional functionality:
+// extendedImporter wraps jsonnet.FileImporter to add additional functionality:
 // - `import "file.yaml"`
 // - `import "tk"`
-type ExtendedImporter struct {
+type extendedImporter struct {
 	loaders    []importLoader    // for loading jsonnet from somewhere. First one that returns non-nil is used
 	processors []importProcessor // for post-processing (e.g. yaml -> json)
 }
@@ -24,10 +24,10 @@ type importLoader func(importedFrom, importedPath string) (c *jsonnet.Contents,
 // further
 type importProcessor func(contents, foundAt string) (c *jsonnet.Contents, err error)
 
-// NewExtendedImporter returns a new instance of ExtendedImporter with the
+// newExtendedImporter returns a new instance of ExtendedImporter with the
 // correct jpaths set up
-func NewExtendedImporter(jpath []string) *ExtendedImporter {
-	return &ExtendedImporter{
+func newExtendedImporter(jpath []string) *extendedImporter {
+	return &extendedImporter{
 		loaders: []importLoader{
 			tkLoader,
 			newFileLoader(&jsonnet.FileImporter{
@@ -38,7 +38,7 @@ func NewExtendedImporter(jpath []string) *ExtendedImporter {
 }
 
 // Import implements the functionality offered by the ExtendedImporter
-func (i *ExtendedImporter) Import(importedFrom, importedPath string) (contents jsonnet.Contents, foundAt string, err error) {
+func (i *extendedImporter) Import(importedFrom, importedPath string) (contents jsonnet.Contents, foundAt string, err error) {
 	// load using loader
 	for _, loader := range i.loaders {
 		c, f, err := loader(importedFrom, importedPath)

pkg/jsonnet/tk.libsonnet.go → pkg/jsonnet/implementation/goimpl/tk.libsonnet.go

@@ -1,4 +1,4 @@
-package jsonnet
+package goimpl
 
 import jsonnet "github.com/google/go-jsonnet"
 

pkg/jsonnet/implementation/goimpl/vm.go

@@ -0,0 +1,33 @@
+package goimpl
+
+import (
+	"github.com/google/go-jsonnet"
+	"github.com/grafana/tanka/pkg/jsonnet/native"
+)
+
+// MakeRawVM returns a Jsonnet VM with some extensions of Tanka, including:
+// - extended importer
+// - extCode and tlaCode applied
+// - native functions registered
+// This is exposed because Go is used for advanced use cases, like finding transitive imports or linting.
+func MakeRawVM(importPaths []string, extCode map[string]string, tlaCode map[string]string, maxStack int) *jsonnet.VM {
+	vm := jsonnet.MakeVM()
+	vm.Importer(newExtendedImporter(importPaths))
+
+	for k, v := range extCode {
+		vm.ExtCode(k, v)
+	}
+	for k, v := range tlaCode {
+		vm.TLACode(k, v)
+	}
+
+	for _, nf := range native.Funcs() {
+		vm.NativeFunction(nf)
+	}
+
+	if maxStack > 0 {
+		vm.MaxStack = maxStack
+	}
+
+	return vm
+}

pkg/jsonnet/implementation/implementation.go

@@ -0,0 +1,16 @@
+package implementation
+
+import (
+	"fmt"
+
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/goimpl"
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/types"
+)
+
+func Get(name string) (types.JsonnetImplementation, error) {
+	if name == "go" || name == "" {
+		return &goimpl.JsonnetGoImplementation{}, nil
+	}
+
+	return nil, fmt.Errorf("unknown jsonnet implementation: %s", name)
+}

pkg/jsonnet/implementation/types/types.go

@@ -0,0 +1,10 @@
+package types
+
+type JsonnetVM interface {
+	EvaluateAnonymousSnippet(filename, snippet string) (string, error)
+	EvaluateFile(filename string) (string, error)
+}
+
+type JsonnetImplementation interface {
+	MakeVM(importPaths []string, extCode map[string]string, tlaCode map[string]string, maxStack int) JsonnetVM
+}

pkg/jsonnet/imports.go

@@ -15,8 +15,8 @@ import (
 	"github.com/google/go-jsonnet/toolutils"
 	"github.com/pkg/errors"
 
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/goimpl"
 	"github.com/grafana/tanka/pkg/jsonnet/jpath"
-	"github.com/grafana/tanka/pkg/jsonnet/native"
 )
 
 var importsRegexp = regexp.MustCompile(`import(str)?\s+['"]([^'"%()]+)['"]`)
@@ -48,12 +48,7 @@ func TransitiveImports(dir string) ([]string, error) {
 		return nil, errors.Wrap(err, "resolving JPATH")
 	}
 
-	vm := jsonnet.MakeVM()
-	vm.Importer(NewExtendedImporter(jpath))
-	for _, nf := range native.Funcs() {
-		vm.NativeFunction(nf)
-	}
-
+	vm := goimpl.MakeRawVM(jpath, nil, nil, 0)
 	node, err := jsonnet.SnippetToAST(filepath.Base(entrypoint), string(sonnet))
 	if err != nil {
 		return nil, errors.Wrap(err, "creating Jsonnet AST")

pkg/jsonnet/imports_test.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/goimpl"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -53,7 +54,7 @@ func BenchmarkGetSnippetHash(b *testing.B) {
 			// Create a VM. It's important to reuse the same VM
 			// While there is a caching mechanism that normally shouldn't be shared in a benchmark iteration,
 			// it's useful to evaluate its impact here, because the caching will also improve the evaluation performance afterwards.
-			vm := MakeVM(Opts{ImportPaths: []string{tempDir}})
+			vm := goimpl.MakeRawVM([]string{tempDir}, nil, nil, 0)
 			content, err := os.ReadFile(filepath.Join(tempDir, "main.jsonnet"))
 			require.NoError(b, err)
 

pkg/jsonnet/lint.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/gobwas/glob"
 	"github.com/google/go-jsonnet/linter"
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/goimpl"
 	"github.com/grafana/tanka/pkg/jsonnet/jpath"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
@@ -106,14 +107,13 @@ func lintWithRecover(file string) (buf bytes.Buffer, success bool) {
 		return
 	}
 
-	vm := MakeVM(Opts{})
 	jpaths, _, _, err := jpath.Resolve(file, true)
 	if err != nil {
 		fmt.Fprintf(&buf, "got an error getting jpath for %s: %v\n\n", file, err)
 		return
 	}
+	vm := goimpl.MakeRawVM(jpaths, nil, nil, 0)
 
-	vm.Importer(NewExtendedImporter(jpaths))
 	failed := linter.LintSnippet(vm, &buf, []linter.Snippet{{FileName: file, Code: string(content)}})
 	return buf, !failed
 }

pkg/process/data_test.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"path/filepath"
 
-	"github.com/grafana/tanka/pkg/jsonnet"
+	"github.com/grafana/tanka/pkg/jsonnet/implementation/goimpl"
 	"github.com/grafana/tanka/pkg/kubernetes/manifest"
 )
 
@@ -18,9 +18,7 @@ type testData struct {
 func loadFixture(name string) testData {
 	filename := filepath.Join("./testdata", name)
 
-	vm := jsonnet.MakeVM(jsonnet.Opts{
-		ImportPaths: []string{"./testdata"},
-	})
+	vm := goimpl.MakeRawVM([]string{"./testdata"}, nil, nil, 0)
 
 	data, err := vm.EvaluateFile(filename)
 	if err != nil {

pkg/spec/v1alpha1/environment.go

@@ -57,14 +57,15 @@ func (m Metadata) NameLabel() string {
 
 // Spec defines Kubernetes properties
 type Spec struct {
-	APIServer        string           `json:"apiServer,omitempty"`
-	ContextNames     []string         `json:"contextNames,omitempty"`
-	Namespace        string           `json:"namespace"`
-	DiffStrategy     string           `json:"diffStrategy,omitempty"`
-	ApplyStrategy    string           `json:"applyStrategy,omitempty"`
-	InjectLabels     bool             `json:"injectLabels,omitempty"`
-	ResourceDefaults ResourceDefaults `json:"resourceDefaults"`
-	ExpectVersions   ExpectVersions   `json:"expectVersions"`
+	APIServer                   string           `json:"apiServer,omitempty"`
+	ContextNames                []string         `json:"contextNames,omitempty"`
+	Namespace                   string           `json:"namespace"`
+	DiffStrategy                string           `json:"diffStrategy,omitempty"`
+	ApplyStrategy               string           `json:"applyStrategy,omitempty"`
+	InjectLabels                bool             `json:"injectLabels,omitempty"`
+	ResourceDefaults            ResourceDefaults `json:"resourceDefaults"`
+	ExpectVersions              ExpectVersions   `json:"expectVersions"`
+	ExportJsonnetImplementation string           `json:"exportJsonnetImplementation,omitempty"`
 }
 
 // ExpectVersions holds semantic version constraints

pkg/tanka/parallel.go

@@ -49,6 +49,10 @@ func parallelLoadEnvironments(envs []*v1alpha1.Environment, opts parallelOpts) (
 		// to Tanka workflow thus being able to handle such cases
 		o.JsonnetOpts = o.JsonnetOpts.Clone()
 
+		if o.JsonnetOpts.JsonnetImplementation == "" {
+			o.JsonnetOpts.JsonnetImplementation = env.Spec.ExportJsonnetImplementation
+		}
+
 		o.Name = env.Metadata.Name
 		path := env.Metadata.Namespace
 		rootDir, err := jpath.FindRoot(path)