From 21aa8d4a13155e57554130fd0cbc078d5f9c79e1 Mon Sep 17 00:00:00 2001
From: Anders Eknert <anders@styra.com>
Date: Mon, 11 Nov 2024 22:16:19 +0100
Subject: [PATCH] Add `--instrument` flag to `regal lint` (#1254)

Same as for `opa eval`

Unrelated change is increasing the granularity of the memory allocation
collectors, as I found that helped a lot when tracking down allocations

Signed-off-by: Anders Eknert <anders@styra.com>
---
 cmd/lint.go          |  7 +++++++
 cmd/profiling.go     |  4 ++--
 pkg/linter/linter.go | 16 ++++++++++++++++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/cmd/lint.go b/cmd/lint.go
index f4c4b80d..dc58cccb 100644
--- a/cmd/lint.go
+++ b/cmd/lint.go
@@ -49,6 +49,7 @@ type lintCommandParams struct {
 	enablePrint     bool
 	metrics         bool
 	profile         bool
+	instrument      bool
 	disableAll      bool
 	enableAll       bool
 }
@@ -167,6 +168,8 @@ func init() {
 		"enable metrics reporting (currently supported only for JSON output format)")
 	lintCommand.Flags().BoolVar(&params.profile, "profile", false,
 		"enable profiling metrics to be added to reporting (currently supported only for JSON output format)")
+	lintCommand.Flags().BoolVar(&params.instrument, "instrument", false,
+		"enable instrumentation metrics to be added to reporting (currently supported only for JSON output format)")
 
 	lintCommand.Flags().VarP(&params.disable, "disable", "d",
 		"disable specific rule(s). This flag can be repeated.")
@@ -285,6 +288,10 @@ func lint(args []string, params *lintCommandParams) (report.Report, error) {
 		regal = regal.WithProfiling(true)
 	}
 
+	if params.instrument {
+		regal = regal.WithInstrumentation(true)
+	}
+
 	var userConfig config.Config
 
 	userConfigFile, err := readUserConfig(params, regalDir)
diff --git a/cmd/profiling.go b/cmd/profiling.go
index 3d3215a1..cfb44e79 100644
--- a/cmd/profiling.go
+++ b/cmd/profiling.go
@@ -28,9 +28,9 @@ func wrapProfiling(f func([]string) error) func(*cobra.Command, []string) error
 		case "clock":
 			opts = append(opts, profile.ClockProfile)
 		case "mem_heap":
-			opts = append(opts, profile.MemProfileHeap)
+			opts = append(opts, profile.MemProfileHeap, profile.MemProfileRate(1024))
 		case "mem_allocs":
-			opts = append(opts, profile.MemProfileAllocs)
+			opts = append(opts, profile.MemProfileAllocs, profile.MemProfileRate(1024))
 		case "trace":
 			opts = append(opts, profile.TraceProfile)
 		case "goroutine":
diff --git a/pkg/linter/linter.go b/pkg/linter/linter.go
index a330b644..1fb3d8db 100644
--- a/pkg/linter/linter.go
+++ b/pkg/linter/linter.go
@@ -60,6 +60,7 @@ type Linter struct {
 	disableAll           bool
 	enableAll            bool
 	profiling            bool
+	instrumentation      bool
 }
 
 //nolint:gochecknoglobals
@@ -199,6 +200,13 @@ func (l Linter) WithProfiling(enabled bool) Linter {
 	return l
 }
 
+// WithInstrumentation enables instrumentation metrics.
+func (l Linter) WithInstrumentation(enabled bool) Linter {
+	l.instrumentation = enabled
+
+	return l
+}
+
 // WithPathPrefix sets the root path prefix for the linter.
 // A root directory prefix can be used to resolve relative paths
 // referenced in the linter configuration with absolute file paths or URIs.
@@ -719,6 +727,10 @@ func (l Linter) prepareRegoArgs(query ast.Body) ([]func(*rego.Rego), error) {
 		)
 	}
 
+	if l.instrumentation {
+		regoArgs = append(regoArgs, rego.Instrument(true))
+	}
+
 	if l.dataBundle != nil {
 		regoArgs = append(regoArgs, rego.ParsedBundle("internal", l.dataBundle))
 	}
@@ -864,6 +876,10 @@ func (l Linter) lintWithRegoRules(
 				evalArgs = append(evalArgs, rego.EvalQueryTracer(prof))
 			}
 
+			if l.instrumentation {
+				evalArgs = append(evalArgs, rego.EvalInstrument(true))
+			}
+
 			resultSet, err := pq.Eval(ctx, evalArgs...)
 			if err != nil {
 				errCh <- fmt.Errorf("error encountered in query evaluation %w", err)