plugins: Reduce locks during decision logging (#6797)

Signed-off-by: Magnus Jungsbluth <magnus.jungsbluth@zalando.de>
Co-authored-by: Johan Fylling <johan.dev@fylling.se>

plugins/logs/plugin.go

@@ -415,21 +415,42 @@ func (c *Config) validateAndInjectDefaults(services []string, pluginsList []stri
 
 // Plugin implements decision log buffering and uploading.
 type Plugin struct {
-	manager   *plugins.Manager
-	config    Config
-	buffer    *logBuffer
-	enc       *chunkEncoder
-	mtx       sync.Mutex
-	stop      chan chan struct{}
-	reconfig  chan reconfigure
-	mask      *rego.PreparedEvalQuery
-	maskMutex sync.Mutex
-	drop      *rego.PreparedEvalQuery
-	dropMutex sync.Mutex
-	limiter   *rate.Limiter
-	metrics   metrics.Metrics
-	logger    logging.Logger
-	status    *lstat.Status
+	manager      *plugins.Manager
+	config       Config
+	buffer       *logBuffer
+	enc          *chunkEncoder
+	mtx          sync.Mutex
+	stop         chan chan struct{}
+	reconfig     chan reconfigure
+	preparedMask prepareOnce
+	preparedDrop prepareOnce
+	limiter      *rate.Limiter
+	metrics      metrics.Metrics
+	logger       logging.Logger
+	status       *lstat.Status
+}
+
+type prepareOnce struct {
+	once          *sync.Once
+	preparedQuery *rego.PreparedEvalQuery
+	err           error
+}
+
+func newPrepareOnce() *prepareOnce {
+	return &prepareOnce{
+		once: new(sync.Once),
+	}
+}
+
+func (po *prepareOnce) drop() {
+	po.once = new(sync.Once)
+}
+
+func (po *prepareOnce) prepareOnce(f func() (*rego.PreparedEvalQuery, error)) (*rego.PreparedEvalQuery, error) {
+	po.once.Do(func() {
+		po.preparedQuery, po.err = f()
+	})
+	return po.preparedQuery, po.err
 }
 
 type reconfigure struct {
@@ -513,14 +534,16 @@ func (b *ConfigBuilder) Parse() (*Config, error) {
 func New(parsedConfig *Config, manager *plugins.Manager) *Plugin {
 
 	plugin := &Plugin{
-		manager:  manager,
-		config:   *parsedConfig,
-		stop:     make(chan chan struct{}),
-		buffer:   newLogBuffer(*parsedConfig.Reporting.BufferSizeLimitBytes),
-		enc:      newChunkEncoder(*parsedConfig.Reporting.UploadSizeLimitBytes),
-		reconfig: make(chan reconfigure),
-		logger:   manager.Logger().WithFields(map[string]interface{}{"plugin": Name}),
-		status:   &lstat.Status{},
+		manager:      manager,
+		config:       *parsedConfig,
+		stop:         make(chan chan struct{}),
+		buffer:       newLogBuffer(*parsedConfig.Reporting.BufferSizeLimitBytes),
+		enc:          newChunkEncoder(*parsedConfig.Reporting.UploadSizeLimitBytes),
+		reconfig:     make(chan reconfigure),
+		logger:       manager.Logger().WithFields(map[string]interface{}{"plugin": Name}),
+		status:       &lstat.Status{},
+		preparedDrop: *newPrepareOnce(),
+		preparedMask: *newPrepareOnce(),
 	}
 
 	if parsedConfig.Reporting.MaxDecisionsPerSecond != nil {
@@ -713,13 +736,8 @@ func (p *Plugin) Reconfigure(_ context.Context, config interface{}) {
 	done := make(chan struct{})
 	p.reconfig <- reconfigure{config: config, done: done}
 
-	p.maskMutex.Lock()
-	defer p.maskMutex.Unlock()
-	p.mask = nil
-
-	p.dropMutex.Lock()
-	defer p.dropMutex.Unlock()
-	p.drop = nil
+	p.preparedMask.drop()
+	p.preparedDrop.drop()
 
 	<-done
 }
@@ -753,13 +771,8 @@ func (p *Plugin) Trigger(ctx context.Context) error {
 // fires. This indicates a new compiler instance is available. The decision
 // logger needs to prepare a new masking query.
 func (p *Plugin) compilerUpdated(storage.Transaction) {
-	p.maskMutex.Lock()
-	defer p.maskMutex.Unlock()
-	p.mask = nil
-
-	p.dropMutex.Lock()
-	defer p.dropMutex.Unlock()
-	p.drop = nil
+	p.preparedMask.drop()
+	p.preparedDrop.drop()
 }
 
 func (p *Plugin) loop() {
@@ -975,42 +988,33 @@ func (p *Plugin) bufferChunk(buffer *logBuffer, bs []byte) {
 }
 
 func (p *Plugin) maskEvent(ctx context.Context, txn storage.Transaction, input ast.Value, event *EventV1) error {
+	pq, err := p.preparedMask.prepareOnce(func() (*rego.PreparedEvalQuery, error) {
+		var pq rego.PreparedEvalQuery
 
-	mask, err := func() (rego.PreparedEvalQuery, error) {
+		query := ast.NewBody(ast.NewExpr(ast.NewTerm(p.config.maskDecisionRef)))
 
-		p.maskMutex.Lock()
-		defer p.maskMutex.Unlock()
+		r := rego.New(
+			rego.ParsedQuery(query),
+			rego.Compiler(p.manager.GetCompiler()),
+			rego.Store(p.manager.Store),
+			rego.Transaction(txn),
+			rego.Runtime(p.manager.Info),
+			rego.EnablePrintStatements(p.manager.EnablePrintStatements()),
+			rego.PrintHook(p.manager.PrintHook()),
+		)
 
-		if p.mask == nil {
-
-			query := ast.NewBody(ast.NewExpr(ast.NewTerm(p.config.maskDecisionRef)))
-
-			r := rego.New(
-				rego.ParsedQuery(query),
-				rego.Compiler(p.manager.GetCompiler()),
-				rego.Store(p.manager.Store),
-				rego.Transaction(txn),
-				rego.Runtime(p.manager.Info),
-				rego.EnablePrintStatements(p.manager.EnablePrintStatements()),
-				rego.PrintHook(p.manager.PrintHook()),
-			)
-
-			pq, err := r.PrepareForEval(context.Background())
-			if err != nil {
-				return rego.PreparedEvalQuery{}, err
-			}
-
-			p.mask = &pq
+		pq, err := r.PrepareForEval(context.Background())
+		if err != nil {
+			return nil, err
 		}
-
-		return *p.mask, nil
-	}()
+		return &pq, nil
+	})
 
 	if err != nil {
 		return err
 	}
 
-	rs, err := mask.Eval(
+	rs, err := pq.Eval(
 		ctx,
 		rego.EvalParsedInput(input),
 		rego.EvalTransaction(txn),
@@ -1038,40 +1042,34 @@ func (p *Plugin) maskEvent(ctx context.Context, txn storage.Transaction, input a
 }
 
 func (p *Plugin) dropEvent(ctx context.Context, txn storage.Transaction, input ast.Value) (bool, error) {
+	var err error
 
-	drop, err := func() (rego.PreparedEvalQuery, error) {
+	pq, err := p.preparedDrop.prepareOnce(func() (*rego.PreparedEvalQuery, error) {
+		var pq rego.PreparedEvalQuery
 
-		p.dropMutex.Lock()
-		defer p.dropMutex.Unlock()
+		query := ast.NewBody(ast.NewExpr(ast.NewTerm(p.config.dropDecisionRef)))
+		r := rego.New(
+			rego.ParsedQuery(query),
+			rego.Compiler(p.manager.GetCompiler()),
+			rego.Store(p.manager.Store),
+			rego.Transaction(txn),
+			rego.Runtime(p.manager.Info),
+			rego.EnablePrintStatements(p.manager.EnablePrintStatements()),
+			rego.PrintHook(p.manager.PrintHook()),
+		)
 
-		if p.drop == nil {
-			query := ast.NewBody(ast.NewExpr(ast.NewTerm(p.config.dropDecisionRef)))
-			r := rego.New(
-				rego.ParsedQuery(query),
-				rego.Compiler(p.manager.GetCompiler()),
-				rego.Store(p.manager.Store),
-				rego.Transaction(txn),
-				rego.Runtime(p.manager.Info),
-				rego.EnablePrintStatements(p.manager.EnablePrintStatements()),
-				rego.PrintHook(p.manager.PrintHook()),
-			)
-
-			pq, err := r.PrepareForEval(context.Background())
-			if err != nil {
-				return rego.PreparedEvalQuery{}, err
-			}
-
-			p.drop = &pq
+		pq, err := r.PrepareForEval(context.Background())
+		if err != nil {
+			return nil, err
 		}
-
-		return *p.drop, nil
-	}()
+		return &pq, nil
+	})
 
 	if err != nil {
 		return false, err
 	}
 
-	rs, err := drop.Eval(
+	rs, err := pq.Eval(
 		ctx,
 		rego.EvalParsedInput(input),
 		rego.EvalTransaction(txn),

plugins/logs/plugin_test.go

@@ -2350,6 +2350,152 @@ func TestPluginDrop(t *testing.T) {
 	}
 }
 
+func TestPluginMaskErrorHandling(t *testing.T) {
+	rawPolicy := []byte(`
+			package system.log
+			drop {
+				endswith(input.path, "bar")
+			}`)
+	event := &EventV1{Path: "foo/bar"}
+
+	// Setup fixture. Populate store with simple drop policy.
+	ctx := context.Background()
+	store := inmem.New()
+
+	//checks if raw policy is valid and stores policy in store
+	err := storage.Txn(ctx, store, storage.WriteParams, func(txn storage.Transaction) error {
+		if err := store.UpsertPolicy(ctx, txn, "test.rego", rawPolicy); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var output []string
+
+	// Create and start manager. Start is required so that stored policies
+	// get compiled and made available to the plugin.
+	manager, err := plugins.New(
+		nil,
+		"test",
+		store,
+		plugins.EnablePrintStatements(true),
+		plugins.PrintHook(appendingPrintHook{printed: &output}),
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := manager.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	// Instantiate the plugin.
+	cfg := &Config{Service: "svc"}
+	trigger := plugins.DefaultTriggerMode
+	cfg.validateAndInjectDefaults([]string{"svc"}, nil, &trigger)
+
+	plugin := New(cfg, manager)
+
+	if err := plugin.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+	input, err := event.AST()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	type badTransaction struct {
+		storage.Transaction
+	}
+
+	expErr := "storage_invalid_txn_error: unexpected transaction type *logs.badTransaction"
+	err = plugin.maskEvent(ctx, &badTransaction{}, input, event)
+	if err.Error() != expErr {
+		t.Fatalf("Expected error %v got %v", expErr, err)
+	}
+
+	// We expect the same error on a second call, even though the mask query failed to prepare and won't be prepared again.
+	err = plugin.maskEvent(ctx, nil, input, event)
+	if err.Error() != expErr {
+		t.Fatalf("Expected error %v got %v", expErr, err)
+	}
+}
+
+func TestPluginDropErrorHandling(t *testing.T) {
+	rawPolicy := []byte(`
+			package system.log
+			drop {
+				endswith(input.path, "bar")
+			}`)
+	event := &EventV1{Path: "foo/bar"}
+
+	// Setup fixture. Populate store with simple drop policy.
+	ctx := context.Background()
+	store := inmem.New()
+
+	//checks if raw policy is valid and stores policy in store
+	err := storage.Txn(ctx, store, storage.WriteParams, func(txn storage.Transaction) error {
+		if err := store.UpsertPolicy(ctx, txn, "test.rego", rawPolicy); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var output []string
+
+	// Create and start manager. Start is required so that stored policies
+	// get compiled and made available to the plugin.
+	manager, err := plugins.New(
+		nil,
+		"test",
+		store,
+		plugins.EnablePrintStatements(true),
+		plugins.PrintHook(appendingPrintHook{printed: &output}),
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := manager.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	// Instantiate the plugin.
+	cfg := &Config{Service: "svc"}
+	trigger := plugins.DefaultTriggerMode
+	cfg.validateAndInjectDefaults([]string{"svc"}, nil, &trigger)
+
+	plugin := New(cfg, manager)
+
+	if err := plugin.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+	input, err := event.AST()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	type badTransaction struct {
+		storage.Transaction
+	}
+
+	expErr := "storage_invalid_txn_error: unexpected transaction type *logs.badTransaction"
+	_, err = plugin.dropEvent(ctx, &badTransaction{}, input)
+	if err.Error() != expErr {
+		t.Fatalf("Expected error %v got %v", expErr, err)
+	}
+
+	// We expect the same error on a second call, even though the drop query failed to prepare and won't be prepared again.
+	_, err = plugin.dropEvent(ctx, nil, input)
+	if err.Error() != expErr {
+		t.Fatalf("Expected error %v got %v", expErr, err)
+	}
+}
+
 type testFixtureOptions struct {
 	ConsoleLogger                  *test.Logger
 	ReportingUploadSizeLimitBytes  int64