Skip to content

Commit 3f237e2

Browse files
acpanaacpana
authored
feat: add more generalized Stats (#283)
* 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>
1 parent fc3dce7 commit 3f237e2

File tree

11 files changed

+797
-110
lines changed

11 files changed

+797
-110
lines changed

‎constraint/pkg/client/client.go‎

Lines changed: 48 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
clienterrors "github.com/open-policy-agent/frameworks/constraint/pkg/client/errors"
1717
"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
1818
"github.com/open-policy-agent/frameworks/constraint/pkg/handler"
19+
"github.com/open-policy-agent/frameworks/constraint/pkg/instrumentation"
1920
"github.com/open-policy-agent/frameworks/constraint/pkg/types"
2021
"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions"
2122
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
@@ -670,7 +671,7 @@ func (c *Client) Review(ctx context.Context, obj interface{}, opts ...drivers.Qu
670671
for target, review := range reviews {
671672
constraints := constraintsByTarget[target]
672673

673-
resp, err := c.review(ctx, target, constraints, review, opts...)
674+
resp, stats, err := c.review(ctx, target, constraints, review, opts...)
674675
if err != nil {
675676
errMap.Add(target, err)
676677
continue
@@ -684,6 +685,18 @@ func (c *Client) Review(ctx context.Context, obj interface{}, opts ...drivers.Qu
684685
resp.Sort()
685686

686687
responses.ByTarget[target] = resp
688+
if stats != nil {
689+
// add the target label to these stats for future collation.
690+
targetLabel := &instrumentation.Label{Name: "target", Value: target}
691+
for _, stat := range stats {
692+
if stat.Labels == nil || len(stat.Labels) == 0 {
693+
stat.Labels = []*instrumentation.Label{targetLabel}
694+
} else {
695+
stat.Labels = append(stat.Labels, targetLabel)
696+
}
697+
}
698+
responses.StatsEntries = append(responses.StatsEntries, stats...)
699+
}
687700
}
688701

689702
if len(errMap) == 0 {
@@ -693,8 +706,9 @@ func (c *Client) Review(ctx context.Context, obj interface{}, opts ...drivers.Qu
693706
return responses, &errMap
694707
}
695708

696-
func (c *Client) review(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*types.Response, error) {
709+
func (c *Client) review(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*types.Response, []*instrumentation.StatsEntry, error) {
697710
var results []*types.Result
711+
var stats []*instrumentation.StatsEntry
698712
var tracesBuilder strings.Builder
699713
errs := &errors.ErrorMap{}
700714

@@ -703,11 +717,11 @@ func (c *Client) review(ctx context.Context, target string, constraints []*unstr
703717
for _, constraint := range constraints {
704718
template, ok := c.templates[strings.ToLower(constraint.GetObjectKind().GroupVersionKind().Kind)]
705719
if !ok {
706-
return nil, fmt.Errorf("%w: while loading driver for constraint %s", ErrMissingConstraintTemplate, constraint.GetName())
720+
return nil, nil, fmt.Errorf("%w: while loading driver for constraint %s", ErrMissingConstraintTemplate, constraint.GetName())
707721
}
708722
driver := c.driverForTemplate(template.template)
709723
if driver == "" {
710-
return nil, fmt.Errorf("%w: while loading driver for constraint %s", clienterrors.ErrNoDriver, constraint.GetName())
724+
return nil, nil, fmt.Errorf("%w: while loading driver for constraint %s", clienterrors.ErrNoDriver, constraint.GetName())
711725
}
712726
driverToConstraints[driver] = append(driverToConstraints[driver], constraint)
713727
}
@@ -716,17 +730,21 @@ func (c *Client) review(ctx context.Context, target string, constraints []*unstr
716730
if len(driverToConstraints[driverName]) == 0 {
717731
continue
718732
}
719-
driverResults, trace, err := driver.Query(ctx, target, driverToConstraints[driverName], review, opts...)
733+
qr, err := driver.Query(ctx, target, driverToConstraints[driverName], review, opts...)
720734
if err != nil {
721735
errs.Add(driverName, err)
722736
continue
723737
}
724-
results = append(results, driverResults...)
738+
if qr != nil {
739+
results = append(results, qr.Results...)
725740

726-
if trace != nil {
727-
tracesBuilder.WriteString(fmt.Sprintf("DRIVER %s:\n\n", driverName))
728-
tracesBuilder.WriteString(*trace)
729-
tracesBuilder.WriteString("\n\n")
741+
stats = append(stats, qr.StatsEntries...)
742+
743+
if qr.Trace != nil {
744+
tracesBuilder.WriteString(fmt.Sprintf("DRIVER %s:\n\n", driverName))
745+
tracesBuilder.WriteString(*qr.Trace)
746+
tracesBuilder.WriteString("\n\n")
747+
}
730748
}
731749
}
732750

@@ -749,7 +767,7 @@ func (c *Client) review(ctx context.Context, target string, constraints []*unstr
749767
Trace: trace,
750768
Target: target,
751769
Results: results,
752-
}, errRet
770+
}, stats, errRet
753771
}
754772

755773
// Dump dumps the state of OPA to aid in debugging.
@@ -767,6 +785,25 @@ func (c *Client) Dump(ctx context.Context) (string, error) {
767785
return dumpBuilder.String(), nil
768786
}
769787

788+
func (c *Client) GetDescriptionForStat(source instrumentation.Source, statName string) string {
789+
if source.Type != instrumentation.EngineSourceType {
790+
// only handle engine source for now
791+
return instrumentation.UnknownDescription
792+
}
793+
794+
driver, ok := c.drivers[source.Value]
795+
if !ok {
796+
return instrumentation.UnknownDescription
797+
}
798+
799+
desc, err := driver.GetDescriptionForStat(statName)
800+
if err != nil {
801+
return instrumentation.UnknownDescription
802+
}
803+
804+
return desc
805+
}
806+
770807
// knownTargets returns a sorted list of known target names.
771808
func (c *Client) knownTargets() []string {
772809
var knownTargets []string

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ func (d *Driver) RemoveData(ctx context.Context, target string, path storage.Pat
197197
return nil
198198
}
199199

200-
func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) ([]*types.Result, *string, error) {
200+
func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*drivers.QueryResponse, error) {
201201
results := []*types.Result{}
202202
for i := range constraints {
203203
constraint := constraints[i]
@@ -209,9 +209,14 @@ func (d *Driver) Query(ctx context.Context, target string, constraints []*unstru
209209
}
210210
results = append(results, result)
211211
}
212-
return results, nil, nil
212+
213+
return &drivers.QueryResponse{Results: results}, nil
213214
}
214215

215216
func (d *Driver) Dump(ctx context.Context) (string, error) {
216217
return "", nil
217218
}
219+
220+
func (d *Driver) GetDescriptionForStat(_ string) (string, error) {
221+
return "", fmt.Errorf("unknown stat name")
222+
}

‎constraint/pkg/client/drivers/interface.go‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55

66
"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
7-
"github.com/open-policy-agent/frameworks/constraint/pkg/types"
87
"github.com/open-policy-agent/opa/storage"
98
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
109
)
@@ -42,15 +41,17 @@ type Driver interface {
4241
RemoveData(ctx context.Context, target string, path storage.Path) error
4342

4443
// Query runs the passed target's Constraints against review.
45-
//
46-
// Returns results for each violated Constraint.
47-
// Returns a trace if specified in query options or enabled at Driver creation.
44+
// Returns a QueryResponse type.
4845
// Returns an error if there was a problem executing the Query.
49-
Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...QueryOpt) ([]*types.Result, *string, error)
46+
Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...QueryOpt) (*QueryResponse, error)
5047

5148
// Dump outputs the entire state of compiled Templates, added Constraints, and
5249
// cached data used for referential Constraints.
5350
Dump(ctx context.Context) (string, error)
51+
52+
// GetDescriptionForStat returns the description for a given stat name
53+
// or errors out for an unknown stat.
54+
GetDescriptionForStat(statName string) (string, error)
5455
}
5556

5657
// ConstraintKey uniquely identifies a Constraint.

‎constraint/pkg/client/drivers/query_opts.go‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@ package drivers
22

33
type QueryCfg struct {
44
TracingEnabled bool
5+
StatsEnabled bool
56
}
67

7-
// QueryOpt specifies optional arguments for Rego queries.
8+
// QueryOpt specifies optional arguments for Query driver calls.
89
type QueryOpt func(*QueryCfg)
910

1011
// Tracing enables Rego tracing for a single query.
@@ -14,3 +15,12 @@ func Tracing(enabled bool) QueryOpt {
1415
cfg.TracingEnabled = enabled
1516
}
1617
}
18+
19+
// Stats(true) enables the driver to return evaluation stats for a single
20+
// query. If stats is enabled for the Driver at construction time, then
21+
// Stats(false) does not disable Stats for this single query.
22+
func Stats(enabled bool) QueryOpt {
23+
return func(cfg *QueryCfg) {
24+
cfg.StatsEnabled = enabled
25+
}
26+
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,16 @@ func Externs(externs ...string) Arg {
154154
}
155155
}
156156

157+
// GatherStats starts collecting various stats around the
158+
// underlying engine's calls.
159+
func GatherStats() Arg {
160+
return func(driver *Driver) error {
161+
driver.gatherStats = true
162+
163+
return nil
164+
}
165+
}
166+
157167
// Currently rules should only access data.inventory.
158168
var validDataFields = map[string]bool{
159169
"inventory": true,

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

Lines changed: 74 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
clienterrors "github.com/open-policy-agent/frameworks/constraint/pkg/client/errors"
1818
"github.com/open-policy-agent/frameworks/constraint/pkg/core/templates"
1919
"github.com/open-policy-agent/frameworks/constraint/pkg/externaldata"
20+
"github.com/open-policy-agent/frameworks/constraint/pkg/instrumentation"
2021
"github.com/open-policy-agent/frameworks/constraint/pkg/types"
2122
"github.com/open-policy-agent/opa/ast"
2223
"github.com/open-policy-agent/opa/rego"
@@ -31,6 +32,15 @@ import (
3132
const (
3233
libRoot = "data.lib"
3334
violation = "violation"
35+
36+
templateRunTimeNS = "templateRunTimeNS"
37+
templateRunTimeNsDesc = "the number of nanoseconds it took to evaluate all constraints for a template"
38+
39+
constraintCountName = "constraintCount"
40+
constraintCountDescription = "the number of constraints that were evaluated for the given constraint kind"
41+
42+
tracingEnabledLabelName = "TracingEnabled"
43+
printEnabledLabelName = "PrintEnabled"
3444
)
3545

3646
var _ drivers.Driver = &Driver{}
@@ -72,14 +82,9 @@ type Driver struct {
7282

7383
// clientCertWatcher is a watcher for the TLS certificate used to communicate with providers.
7484
clientCertWatcher *certwatcher.CertWatcher
75-
}
7685

77-
// EvaluationMeta has rego specific metadata from evaluation.
78-
type EvaluationMeta struct {
79-
// TemplateRunTime is the number of milliseconds it took to evaluate all constraints for a template.
80-
TemplateRunTime float64 `json:"templateRunTime"`
81-
// ConstraintCount indicates how many constraints were evaluated for an underlying rego engine eval call.
82-
ConstraintCount uint `json:"constraintCount"`
86+
// gatherStats controls whether the driver gathers any stats around its API calls.
87+
gatherStats bool
8388
}
8489

8590
// Name returns the name of the driver.
@@ -233,9 +238,9 @@ func (d *Driver) eval(ctx context.Context, compiler *ast.Compiler, target string
233238
return res, t, err
234239
}
235240

236-
func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) ([]*types.Result, *string, error) {
241+
func (d *Driver) Query(ctx context.Context, target string, constraints []*unstructured.Unstructured, review interface{}, opts ...drivers.QueryOpt) (*drivers.QueryResponse, error) {
237242
if len(constraints) == 0 {
238-
return nil, nil, nil
243+
return nil, nil
239244
}
240245

241246
constraintsByKind := toConstraintsByKind(constraints)
@@ -250,27 +255,34 @@ func (d *Driver) Query(ctx context.Context, target string, constraints []*unstru
250255
// once per call to Query instead of once per compiler.
251256
reviewMap, err := toInterfaceMap(review)
252257
if err != nil {
253-
return nil, nil, err
258+
return nil, err
254259
}
255260

256261
d.mtx.RLock()
257262
defer d.mtx.RUnlock()
258263

264+
cfg := &drivers.QueryCfg{}
265+
for _, opt := range opts {
266+
opt(cfg)
267+
}
268+
269+
var statsEntries []*instrumentation.StatsEntry
270+
259271
for kind, kindConstraints := range constraintsByKind {
260272
evalStartTime := time.Now()
261273
compiler := d.compilers.getCompiler(target, kind)
262274
if compiler == nil {
263275
// The Template was just removed, so the Driver is in an inconsistent
264276
// state with Client. Raise this as an error rather than attempting to
265277
// continue.
266-
return nil, nil, fmt.Errorf("missing Template %q for target %q", kind, target)
278+
return nil, fmt.Errorf("missing Template %q for target %q", kind, target)
267279
}
268280

269281
// Parse input into an ast.Value to avoid round-tripping through JSON when
270282
// possible.
271283
parsedInput, err := toParsedInput(target, kindConstraints, reviewMap)
272284
if err != nil {
273-
return nil, nil, err
285+
return nil, err
274286
}
275287

276288
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
297309

298310
kindResults, err := drivers.ToResults(constraintsMap, resultSet)
299311
if err != nil {
300-
return nil, nil, err
301-
}
302-
303-
for _, result := range kindResults {
304-
result.EvaluationMeta = EvaluationMeta{
305-
TemplateRunTime: float64(evalEndTime.Nanoseconds()) / 1000000,
306-
ConstraintCount: uint(len(kindResults)),
307-
}
312+
return nil, err
308313
}
309314

310315
results = append(results, kindResults...)
316+
317+
if d.gatherStats || (cfg != nil && cfg.StatsEnabled) {
318+
statsEntries = append(statsEntries,
319+
&instrumentation.StatsEntry{
320+
Scope: instrumentation.TemplateScope,
321+
StatsFor: kind,
322+
Stats: []*instrumentation.Stat{
323+
{
324+
Name: templateRunTimeNS,
325+
Value: uint64(evalEndTime.Nanoseconds()),
326+
Source: instrumentation.Source{
327+
Type: instrumentation.EngineSourceType,
328+
Value: schema.Name,
329+
},
330+
},
331+
{
332+
Name: constraintCountName,
333+
Value: len(kindConstraints),
334+
Source: instrumentation.Source{
335+
Type: instrumentation.EngineSourceType,
336+
Value: schema.Name,
337+
},
338+
},
339+
},
340+
Labels: []*instrumentation.Label{
341+
{
342+
Name: tracingEnabledLabelName,
343+
Value: d.traceEnabled || cfg.TracingEnabled,
344+
},
345+
{
346+
Name: printEnabledLabelName,
347+
Value: d.printEnabled,
348+
},
349+
},
350+
})
351+
}
311352
}
312353

313354
traceString := traceBuilder.String()
314355
if len(traceString) != 0 {
315-
return results, &traceString, nil
356+
return &drivers.QueryResponse{Results: results, Trace: &traceString, StatsEntries: statsEntries}, nil
316357
}
317358

318-
return results, nil, nil
359+
return &drivers.QueryResponse{Results: results, StatsEntries: statsEntries}, nil
319360
}
320361

321362
func (d *Driver) Dump(ctx context.Context) (string, error) {
@@ -358,6 +399,17 @@ func (d *Driver) Dump(ctx context.Context) (string, error) {
358399
return string(b), nil
359400
}
360401

402+
func (d *Driver) GetDescriptionForStat(statName string) (string, error) {
403+
switch statName {
404+
case templateRunTimeNS:
405+
return templateRunTimeNsDesc, nil
406+
case constraintCountName:
407+
return constraintCountDescription, nil
408+
default:
409+
return "", fmt.Errorf("unknown stat name")
410+
}
411+
}
412+
361413
func (d *Driver) getTLSCertificate() (*tls.Certificate, error) {
362414
if !d.enableExternalDataClientAuth {
363415
return nil, nil

0 commit comments

Comments
 (0)