feat: add more generalized Stats

constraint/pkg/client/client.go

@@ -16,6 +16,7 @@ import (
 	clienterrors "github.com/open-policy-agent/frameworks/constraint/pkg/client/errors"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/handler"
+	"github.com/open-policy-agent/frameworks/constraint/pkg/instrumentation"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/types"
 	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
@@ -670,7 +671,7 @@ func (c *Client) Review(ctx context.Context, obj interface{}, opts ...drivers.Qu
 	for target, review := range reviews {
 		constraints := constraintsByTarget[target]
 
-		resp, err := c.review(ctx, target, constraints, review, opts...)
+		resp, stats, err := c.review(ctx, target, constraints, review, opts...)
 		if err != nil {
 			errMap.Add(target, err)
 			continue
@@ -684,6 +685,18 @@ func (c *Client) Review(ctx context.Context, obj interface{}, opts ...drivers.Qu
 		resp.Sort()
 
 		responses.ByTarget[target] = resp
+		if stats != nil {
+			// add the target label to these stats for future collation.
+			targetLabel := &instrumentation.Label{Name: "target", Value: target}
+			for _, stat := range stats {
+				if stat.Labels == nil || len(stat.Labels) == 0 {
+					stat.Labels = []*instrumentation.Label{targetLabel}
+				} else {
+					stat.Labels = append(stat.Labels, targetLabel)
+				}
+			}
+			responses.StatsEntries = append(responses.StatsEntries, stats...)
+		}
 	}
 
 	if len(errMap) == 0 {
@@ -693,8 +706,9 @@ func (c *Client) Review(ctx context.Context, obj interface{}, opts ...drivers.Qu
 	return responses, &errMap
 }
 
-func (c *Client) review(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*types.Response, error) {
+func (c *Client) review(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*types.Response, []*instrumentation.StatsEntry, error) {
 	var results []*types.Result
+	var stats []*instrumentation.StatsEntry
 	var tracesBuilder strings.Builder
 	errs := &errors.ErrorMap{}
 
@@ -703,11 +717,11 @@ func (c *Client) review(ctx context.Context, target string, constraints []*unstr
 	for _, constraint := range constraints {
 		template, ok := c.templates[strings.ToLower(constraint.GetObjectKind().GroupVersionKind().Kind)]
 		if !ok {
-			return nil, fmt.Errorf("%w: while loading driver for constraint %s", ErrMissingConstraintTemplate, constraint.GetName())
+			return nil, nil, fmt.Errorf("%w: while loading driver for constraint %s", ErrMissingConstraintTemplate, constraint.GetName())
 		}
 		driver := c.driverForTemplate(template.template)
 		if driver == "" {
-			return nil, fmt.Errorf("%w: while loading driver for constraint %s", clienterrors.ErrNoDriver, constraint.GetName())
+			return nil, nil, fmt.Errorf("%w: while loading driver for constraint %s", clienterrors.ErrNoDriver, constraint.GetName())
 		}
 		driverToConstraints[driver] = append(driverToConstraints[driver], constraint)
 	}
@@ -716,17 +730,21 @@ func (c *Client) review(ctx context.Context, target string, constraints []*unstr
 		if len(driverToConstraints[driverName]) == 0 {
 			continue
 		}
-		driverResults, trace, err := driver.Query(ctx, target, driverToConstraints[driverName], review, opts...)
+		qr, err := driver.Query(ctx, target, driverToConstraints[driverName], review, opts...)
 		if err != nil {
 			errs.Add(driverName, err)
 			continue
 		}
-		results = append(results, driverResults...)
+		if qr != nil {
+			results = append(results, qr.Results...)
 
-		if trace != nil {
-			tracesBuilder.WriteString(fmt.Sprintf("DRIVER %s:\n\n", driverName))
-			tracesBuilder.WriteString(*trace)
-			tracesBuilder.WriteString("\n\n")
+			stats = append(stats, qr.StatsEntries...)
+
+			if qr.Trace != nil {
+				tracesBuilder.WriteString(fmt.Sprintf("DRIVER %s:\n\n", driverName))
+				tracesBuilder.WriteString(*qr.Trace)
+				tracesBuilder.WriteString("\n\n")
+			}
 		}
 	}
 
@@ -749,7 +767,7 @@ func (c *Client) review(ctx context.Context, target string, constraints []*unstr
 		Trace:   trace,
 		Target:  target,
 		Results: results,
-	}, errRet
+	}, stats, errRet
 }
 
 // Dump dumps the state of OPA to aid in debugging.
@@ -767,6 +785,25 @@ func (c *Client) Dump(ctx context.Context) (string, error) {
 	return dumpBuilder.String(), nil
 }
 
+func (c *Client) GetDescriptionForStat(source instrumentation.Source, statName string) string {
+	if source.Type != instrumentation.EngineSourceType {
+		// only handle engine source for now
+		return instrumentation.UnknownDescription
+	}
+
+	driver, ok := c.drivers[source.Value]
+	if !ok {
+		return instrumentation.UnknownDescription
+	}
+
+	desc, err := driver.GetDescriptionForStat(statName)
+	if err != nil {
+		return instrumentation.UnknownDescription
+	}
+
+	return desc
+}
+
 // knownTargets returns a sorted list of known target names.
 func (c *Client) knownTargets() []string {
 	var knownTargets []string

constraint/pkg/client/drivers/fake/fake.go

@@ -197,7 +197,7 @@ func (d *Driver) RemoveData(ctx context.Context, target string, path storage.Pat
 	return nil
 }
 
-func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) ([]*types.Result, *string, error) {
+func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*drivers.QueryResponse, error) {
 	results := []*types.Result{}
 	for i := range constraints {
 		constraint := constraints[i]
@@ -209,9 +209,14 @@ func (d *Driver) Query(ctx context.Context, target string, constraints []*unstru
 		}
 		results = append(results, result)
 	}
-	return results, nil, nil
+
+	return &drivers.QueryResponse{Results: results}, nil
 }
 
 func (d *Driver) Dump(ctx context.Context) (string, error) {
 	return "", nil
 }
+
+func (d *Driver) GetDescriptionForStat(_ string) (string, error) {
+	return "", fmt.Errorf("unknown stat name")
+}

constraint/pkg/client/drivers/interface.go

@@ -4,7 +4,6 @@ import (
 	"context"
 
 	"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
-	"github.com/open-policy-agent/frameworks/constraint/pkg/types"
 	"github.com/open-policy-agent/opa/storage"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
@@ -42,15 +41,17 @@ type Driver interface {
 	RemoveData(ctx context.Context, target string, path storage.Path) error
 
 	// Query runs the passed target's Constraints against review.
-	//
-	// Returns results for each violated Constraint.
-	// Returns a trace if specified in query options or enabled at Driver creation.
+	// Returns a QueryResponse type.
 	// Returns an error if there was a problem executing the Query.
-	Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...QueryOpt) ([]*types.Result, *string, error)
+	Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...QueryOpt) (*QueryResponse, error)
 
 	// Dump outputs the entire state of compiled Templates, added Constraints, and
 	// cached data used for referential Constraints.
 	Dump(ctx context.Context) (string, error)
+
+	// GetDescriptionForStat returns the description for a given stat name
+	// or errors out for an unknown stat.
+	GetDescriptionForStat(statName string) (string, error)
 }
 
 // ConstraintKey uniquely identifies a Constraint.

constraint/pkg/client/drivers/query_opts.go

@@ -2,9 +2,10 @@ package drivers
 
 type QueryCfg struct {
 	TracingEnabled bool
+	StatsEnabled   bool
 }
 
-// QueryOpt specifies optional arguments for Rego queries.
+// QueryOpt specifies optional arguments for Query driver calls.
 type QueryOpt func(*QueryCfg)
 
 // Tracing enables Rego tracing for a single query.
@@ -14,3 +15,12 @@ func Tracing(enabled bool) QueryOpt {
 		cfg.TracingEnabled = enabled
 	}
 }
+
+// Stats(true) enables the driver to return evaluation stats for a single
+// query. If stats is enabled for the Driver at construction time, then
+// Stats(false) does not disable Stats for this single query.
+func Stats(enabled bool) QueryOpt {
+	return func(cfg *QueryCfg) {
+		cfg.StatsEnabled = enabled
+	}
+}

constraint/pkg/client/drivers/rego/args.go

@@ -154,6 +154,16 @@ func Externs(externs ...string) Arg {
 	}
 }
 
+// GatherStats starts collecting various stats around the
+// underlying engine's calls.
+func GatherStats() Arg {
+	return func(driver *Driver) error {
+		driver.gatherStats = true
+
+		return nil
+	}
+}
+
 // Currently rules should only access data.inventory.
 var validDataFields = map[string]bool{
 	"inventory": true,

constraint/pkg/client/drivers/rego/driver.go

@@ -17,6 +17,7 @@ import (
 	clienterrors "github.com/open-policy-agent/frameworks/constraint/pkg/client/errors"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/externaldata"
+	"github.com/open-policy-agent/frameworks/constraint/pkg/instrumentation"
 	"github.com/open-policy-agent/frameworks/constraint/pkg/types"
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
@@ -31,6 +32,15 @@ import (
 const (
 	libRoot   = "data.lib"
 	violation = "violation"
+
+	templateRunTimeNS     = "templateRunTimeNS"
+	templateRunTimeNsDesc = "the number of nanoseconds it took to evaluate all constraints for a template"
+
+	constraintCountName        = "constraintCount"
+	constraintCountDescription = "the number of constraints that were evaluated for the given constraint kind"
+
+	tracingEnabledLabelName = "TracingEnabled"
+	printEnabledLabelName   = "PrintEnabled"
 )
 
 var _ drivers.Driver = &Driver{}
@@ -72,14 +82,9 @@ type Driver struct {
 
 	// clientCertWatcher is a watcher for the TLS certificate used to communicate with providers.
 	clientCertWatcher *certwatcher.CertWatcher
-}
 
-// EvaluationMeta has rego specific metadata from evaluation.
-type EvaluationMeta struct {
-	// TemplateRunTime is the number of milliseconds it took to evaluate all constraints for a template.
-	TemplateRunTime float64 `json:"templateRunTime"`
-	// ConstraintCount indicates how many constraints were evaluated for an underlying rego engine eval call.
-	ConstraintCount uint `json:"constraintCount"`
+	// gatherStats controls whether the driver gathers any stats around its API calls.
+	gatherStats bool
 }
 
 // Name returns the name of the driver.
@@ -233,9 +238,9 @@ func (d *Driver) eval(ctx context.Context, compiler *ast.Compiler, target string
 	return res, t, err
 }
 
-func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) ([]*types.Result, *string, error) {
+func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*drivers.QueryResponse, error) {
 	if len(constraints) == 0 {
-		return nil, nil, nil
+		return nil, nil
 	}
 
 	constraintsByKind := toConstraintsByKind(constraints)
@@ -250,27 +255,34 @@ func (d *Driver) Query(ctx context.Context, target string, constraints []*unstru
 	// once per call to Query instead of once per compiler.
 	reviewMap, err := toInterfaceMap(review)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	d.mtx.RLock()
 	defer d.mtx.RUnlock()
 
+	cfg := &drivers.QueryCfg{}
+	for _, opt := range opts {
+		opt(cfg)
+	}
+
+	var statsEntries []*instrumentation.StatsEntry
+
 	for kind, kindConstraints := range constraintsByKind {
 		evalStartTime := time.Now()
 		compiler := d.compilers.getCompiler(target, kind)
 		if compiler == nil {
 			// The Template was just removed, so the Driver is in an inconsistent
 			// state with Client. Raise this as an error rather than attempting to
 			// continue.
-			return nil, nil, fmt.Errorf("missing Template %q for target %q", kind, target)
+			return nil, fmt.Errorf("missing Template %q for target %q", kind, target)
 		}
 
 		// Parse input into an ast.Value to avoid round-tripping through JSON when
 		// possible.
 		parsedInput, err := toParsedInput(target, kindConstraints, reviewMap)
 		if err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 
 		resultSet, trace, err := d.eval(ctx, compiler, target, path, parsedInput, opts...)
@@ -297,25 +309,54 @@ func (d *Driver) Query(ctx context.Context, target string, constraints []*unstru
 
 		kindResults, err := drivers.ToResults(constraintsMap, resultSet)
 		if err != nil {
-			return nil, nil, err
-		}
-
-		for _, result := range kindResults {
-			result.EvaluationMeta = EvaluationMeta{
-				TemplateRunTime: float64(evalEndTime.Nanoseconds()) / 1000000,
-				ConstraintCount: uint(len(kindResults)),
-			}
+			return nil, err
 		}
 
 		results = append(results, kindResults...)
+
+		if d.gatherStats || (cfg != nil && cfg.StatsEnabled) {
+			statsEntries = append(statsEntries,
+				&instrumentation.StatsEntry{
+					Scope:    instrumentation.TemplateScope,
+					StatsFor: kind,
+					Stats: []*instrumentation.Stat{
+						{
+							Name:  templateRunTimeNS,
+							Value: uint64(evalEndTime.Nanoseconds()),
+							Source: instrumentation.Source{
+								Type:  instrumentation.EngineSourceType,
+								Value: schema.Name,
+							},
+						},
+						{
+							Name:  constraintCountName,
+							Value: len(kindConstraints),
+							Source: instrumentation.Source{
+								Type:  instrumentation.EngineSourceType,
+								Value: schema.Name,
+							},
+						},
+					},
+					Labels: []*instrumentation.Label{
+						{
+							Name:  tracingEnabledLabelName,
+							Value: d.traceEnabled || cfg.TracingEnabled,
+						},
+						{
+							Name:  printEnabledLabelName,
+							Value: d.printEnabled,
+						},
+					},
+				})
+		}
 	}
 
 	traceString := traceBuilder.String()
 	if len(traceString) != 0 {
-		return results, &traceString, nil
+		return &drivers.QueryResponse{Results: results, Trace: &traceString, StatsEntries: statsEntries}, nil
 	}
 
-	return results, nil, nil
+	return &drivers.QueryResponse{Results: results, StatsEntries: statsEntries}, nil
 }
 
 func (d *Driver) Dump(ctx context.Context) (string, error) {
@@ -358,6 +399,17 @@ func (d *Driver) Dump(ctx context.Context) (string, error) {
 	return string(b), nil
 }
 
+func (d *Driver) GetDescriptionForStat(statName string) (string, error) {
+	switch statName {
+	case templateRunTimeNS:
+		return templateRunTimeNsDesc, nil
+	case constraintCountName:
+		return constraintCountDescription, nil
+	default:
+		return "", fmt.Errorf("unknown stat name")
+	}
+}
+
 func (d *Driver) getTLSCertificate() (*tls.Certificate, error) {
 	if !d.enableExternalDataClientAuth {
 		return nil, nil