Skip to content

Commit

Permalink
feat: add more generalized Stats (#283)
Browse files Browse the repository at this point in the history 
* feat: add instrumentation primitives

Signed-off-by: Alex Pana <8968914+acpana@users.noreply.github.com>

* review feedback

Signed-off-by: Alex Pana <8968914+acpana@users.noreply.github.com>

* move rego tests in the driver package

Signed-off-by: Alex Pana <8968914+acpana@users.noreply.github.com>

* review: make label names const

Signed-off-by: Alex Pana <8968914+acpana@users.noreply.github.com>

* review: add e2e test for GetDescriptionForStat

Signed-off-by: Alex Pana <8968914+acpana@users.noreply.github.com>

* trailing fmt debug stmt

Signed-off-by: Alex Pana <8968914+acpana@users.noreply.github.com>

---------

Signed-off-by: Alex Pana <8968914+acpana@users.noreply.github.com>
  • Loading branch information
@acpana
acpana authored Apr 11, 2023
1 parent fc3dce7 commit 3f237e2
Show file tree
Hide file tree
Showing 11 changed files with 797 additions and 110 deletions.
59 changes: 48 additions & 11 deletions constraint/pkg/client/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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{}

Expand All @@ -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)
}
Expand All @@ -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")
}
}
}

Expand All @@ -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.
Expand All @@ -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
Expand Down
9 changes: 7 additions & 2 deletions constraint/pkg/client/drivers/fake/fake.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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")
}
11 changes: 6 additions & 5 deletions constraint/pkg/client/drivers/interface.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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.
Expand Down
12 changes: 11 additions & 1 deletion constraint/pkg/client/drivers/query_opts.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
}
}
10 changes: 10 additions & 0 deletions constraint/pkg/client/drivers/rego/args.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
96 changes: 74 additions & 22 deletions constraint/pkg/client/drivers/rego/driver.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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{}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand All @@ -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...)
Expand All @@ -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) {
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 3f237e2

Please sign in to comment.