bundle: Load bundle once (#1136)

* bundle: Load bundle once Previously, loading of the bundle happened many times, each time it takes around 20ms. This improves the performance of the language server by loading the bundle once and reusing it since the bundle is never changed once a binary is running. Signed-off-by: Charlie Egan <charlie@styra.com> * bundle: use pkg var Follows PR review comment Signed-off-by: Charlie Egan <charlie@styra.com> --------- Signed-off-by: Charlie Egan <charlie@styra.com>
StyraInc · Sep 25, 2024 · 6ecb36c · 6ecb36c
1 parent 464a7bc
commit 6ecb36c
Show file tree

Hide file tree

Showing 8 changed files with 31 additions and 48 deletions.
diff --git a/bundle/bundle.go b/bundle/bundle.go
@@ -2,10 +2,17 @@ package bundle
 
 import (
 	"embed"
+
+	rio "github.com/styrainc/regal/internal/io"
 )
 
 // Bundle FS will include the tests as well, but since that has negligible impact on the size of the binary,
 // it's preferable to filter them out from the bundle than to e.g. create a separate directory for tests
 //
 //go:embed *
 var Bundle embed.FS
+
+// LoadedBundle contains the loaded contents of the Bundle.
+//
+//nolint:gochecknoglobals
+var LoadedBundle = rio.MustLoadRegalBundleFS(Bundle)
diff --git a/cmd/lint.go b/cmd/lint.go
@@ -249,13 +249,8 @@ func lint(args []string, params *lintCommandParams) (report.Report, error) {
 		m.Timer(regalmetrics.RegalConfigSearch).Stop()
 	}
 
-	// regal rules are loaded here and passed to the linter separately
-	// as the configuration is also used to determine feature toggles
-	// and the defaults from the data.yaml here.
-	regalRules := rio.MustLoadRegalBundleFS(rbundle.Bundle)
-
 	regal := linter.NewEmptyLinter().
-		WithAddedBundle(regalRules).
+		WithAddedBundle(&rbundle.LoadedBundle).
 		WithDisableAll(params.disableAll).
 		WithDisabledCategories(params.disableCategory.v...).
 		WithDisabledRules(params.disable.v...).
@@ -324,7 +319,7 @@ func lint(args []string, params *lintCommandParams) (report.Report, error) {
 		m.Timer(regalmetrics.RegalConfigParse).Stop()
 	}
 
-	go updateCheckAndWarn(params, regalRules, &userConfig)
+	go updateCheckAndWarn(params, &rbundle.LoadedBundle, &userConfig)
 
 	result, err := regal.Lint(ctx)
 	if err != nil {
@@ -339,8 +334,8 @@ func lint(args []string, params *lintCommandParams) (report.Report, error) {
 	return result, rep.Publish(ctx, result) //nolint:wrapcheck
 }
 
-func updateCheckAndWarn(params *lintCommandParams, regalRules bundle.Bundle, userConfig *config.Config) {
-	mergedConfig, err := config.LoadConfigWithDefaultsFromBundle(&regalRules, userConfig)
+func updateCheckAndWarn(params *lintCommandParams, regalRules *bundle.Bundle, userConfig *config.Config) {
+	mergedConfig, err := config.LoadConfigWithDefaultsFromBundle(regalRules, userConfig)
 	if err != nil {
 		if params.debug {
 			log.Printf("failed to merge user config with default config when checking version: %v", err)

diff --git a/cmd/test.go b/cmd/test.go
@@ -25,8 +25,8 @@ import (
 	"github.com/open-policy-agent/opa/util"
 	"github.com/open-policy-agent/opa/version"
 
+	rbundle "github.com/styrainc/regal/bundle"
 	"github.com/styrainc/regal/internal/compile"
-	"github.com/styrainc/regal/internal/embeds"
 	rio "github.com/styrainc/regal/internal/io"
 	"github.com/styrainc/regal/pkg/builtins"
 	"github.com/styrainc/regal/pkg/config"
@@ -136,18 +136,20 @@ func opaTest(args []string) int {
 		return 1
 	}
 
-	regalBundle := rio.MustLoadRegalBundleFS(embeds.EmbedBundleFS)
-
 	txn, err := store.NewTransaction(ctx, storage.WriteParams)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
 
 		return 1
 	}
 
-	if err := store.Write(ctx, txn, storage.AddOp,
+	if err := store.Write(
+		ctx,
+		txn,
+		storage.AddOp,
 		storage.MustParsePath("/regal"),
-		regalBundle.Data["regal"]); err != nil {
+		rbundle.LoadedBundle.Data["regal"],
+	); err != nil {
 		panic(err)
 	}
 
@@ -168,7 +170,7 @@ func opaTest(args []string) int {
 		WithEnablePrintStatements(!testParams.benchmark).
 		WithSchemas(compile.RegalSchemaSet()).
 		WithUseTypeCheckAnnotations(true).
-		WithModuleLoader(moduleLoader(regalBundle)).
+		WithModuleLoader(moduleLoader(&rbundle.LoadedBundle)).
 		WithRewriteTestRules(testParams.varValues)
 
 	if testParams.threshold > 0 && !testParams.coverage {
@@ -213,12 +215,12 @@ func opaTest(args []string) int {
 	return 0
 }
 
-func moduleLoader(regal bundle.Bundle) ast.ModuleLoader {
+func moduleLoader(regalRules *bundle.Bundle) ast.ModuleLoader {
 	// We use the package declarations to know which modules we still need, and return
 	// those from the embedded regal bundle.
-	extra := make(map[string]struct{}, len(regal.Modules))
+	extra := make(map[string]struct{}, len(regalRules.Modules))
 
-	for _, mod := range regal.Modules {
+	for _, mod := range regalRules.Modules {
 		extra[mod.Parsed.Package.Path.String()] = struct{}{}
 	}
 
@@ -229,7 +231,7 @@ func moduleLoader(regal bundle.Bundle) ast.ModuleLoader {
 
 		extraMods := map[string]*ast.Module{}
 
-		for id, mod := range regal.ParsedModules("bundle") {
+		for id, mod := range regalRules.ParsedModules("bundle") {
 			if _, ok := extra[mod.Package.Path.String()]; ok {
 				extraMods[id] = mod
 			}

diff --git a/internal/lsp/completions/providers/policy.go b/internal/lsp/completions/providers/policy.go
@@ -9,7 +9,6 @@ import (
 	"os"
 
 	"github.com/open-policy-agent/opa/ast"
-	"github.com/open-policy-agent/opa/bundle"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/topdown"
@@ -28,13 +27,6 @@ type Policy struct {
 	pq rego.PreparedEvalQuery
 }
 
-//nolint:gochecknoglobals
-var regalRules = func() bundle.Bundle {
-	regalRules := rio.MustLoadRegalBundleFS(rbundle.Bundle)
-
-	return regalRules
-}()
-
 // NewPolicy creates a new Policy provider. This provider is distinctly different from the other providers
 // as it acts like the entrypoint for all Rego-based providers, and not a single provider "function" like
 // the Go providers do.
@@ -151,7 +143,7 @@ func prepareRegoArgs(store storage.Store, query ast.Body) []func(*rego.Rego) {
 	return []func(*rego.Rego){
 		rego.Store(store),
 		rego.ParsedQuery(query),
-		rego.ParsedBundle("regal", &regalRules),
+		rego.ParsedBundle("regal", &rbundle.LoadedBundle),
 		rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule),
 		rego.Function1(builtins.RegalLastMeta, builtins.RegalLast),
 		// TODO: remove later

diff --git a/internal/lsp/completions/refs/used.go b/internal/lsp/completions/refs/used.go
@@ -32,8 +32,6 @@ var pqInitOnce sync.Once
 // This function is only used by language server code paths and so init() is not
 // used.
 func initialize() {
-	regalRules := rio.MustLoadRegalBundleFS(rbundle.Bundle)
-
 	dataBundle := bundle.Bundle{
 		Manifest: bundle.Manifest{
 			Roots:    &[]string{"internal"},
@@ -49,7 +47,7 @@ func initialize() {
 	}
 
 	regoArgs := []func(*rego.Rego){
-		rego.ParsedBundle("regal", &regalRules),
+		rego.ParsedBundle("regal", &rbundle.LoadedBundle),
 		rego.ParsedBundle("internal", &dataBundle),
 		rego.Query(`data.regal.lsp.completion.ref_names`),
 		rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule),

diff --git a/internal/lsp/rego/rego.go b/internal/lsp/rego/rego.go
@@ -112,11 +112,9 @@ type policy struct {
 }
 
 func initialize() {
-	regalRules := rio.MustLoadRegalBundleFS(rbundle.Bundle)
-
 	createArgs := func(args ...func(*rego.Rego)) []func(*rego.Rego) {
 		return append([]func(*rego.Rego){
-			rego.ParsedBundle("regal", &regalRules),
+			rego.ParsedBundle("regal", &rbundle.LoadedBundle),
 			rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule),
 			rego.Function1(builtins.RegalLastMeta, builtins.RegalLast),
 		}, args...)

diff --git a/internal/lsp/server.go b/internal/lsp/server.go
@@ -24,7 +24,7 @@ import (
 	"github.com/open-policy-agent/opa/format"
 	"github.com/open-policy-agent/opa/storage"
 
-	"github.com/styrainc/regal/bundle"
+	rbundle "github.com/styrainc/regal/bundle"
 	"github.com/styrainc/regal/internal/capabilities"
 	rio "github.com/styrainc/regal/internal/io"
 	"github.com/styrainc/regal/internal/lsp/bundles"
@@ -334,13 +334,6 @@ func (l *LanguageServer) StartConfigWorker(ctx context.Context) {
 		return
 	}
 
-	regalRules, err := rio.LoadRegalBundleFS(bundle.Bundle)
-	if err != nil {
-		l.logError(fmt.Errorf("failed to load regal bundle for defaulting of user config: %w", err))
-
-		return
-	}
-
 	for {
 		select {
 		case <-ctx.Done():
@@ -362,7 +355,7 @@ func (l *LanguageServer) StartConfigWorker(ctx context.Context) {
 				return
 			}
 
-			mergedConfig, err := config.LoadConfigWithDefaultsFromBundle(&regalRules, &userConfig)
+			mergedConfig, err := config.LoadConfigWithDefaultsFromBundle(&rbundle.LoadedBundle, &userConfig)
 			if err != nil {
 				l.logError(fmt.Errorf("failed to load config: %w", err))
 

diff --git a/pkg/linter/linter.go b/pkg/linter/linter.go
@@ -73,10 +73,8 @@ var (
 
 // NewLinter creates a new Regal linter.
 func NewLinter() Linter {
-	regalRules := rio.MustLoadRegalBundleFS(rbundle.Bundle)
-
 	return Linter{
-		ruleBundles: []*bundle.Bundle{&regalRules},
+		ruleBundles: []*bundle.Bundle{&rbundle.LoadedBundle},
 	}
 }
 
@@ -102,8 +100,8 @@ func (l Linter) WithInputModules(input *rules.Input) Linter {
 }
 
 // WithAddedBundle adds a bundle of rules and data to include in evaluation.
-func (l Linter) WithAddedBundle(b bundle.Bundle) Linter {
-	l.ruleBundles = append(l.ruleBundles, &b)
+func (l Linter) WithAddedBundle(b *bundle.Bundle) Linter {
+	l.ruleBundles = append(l.ruleBundles, b)
 
 	return l
 }