Enable concurrent evaluations of matched expressions

expr.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/google/cel-go/common/operators"
 	"github.com/google/uuid"
+	"golang.org/x/sync/semaphore"
 )
 
 var (
@@ -70,11 +71,17 @@ func NewAggregateEvaluator(
 	parser TreeParser,
 	eval ExpressionEvaluator,
 	evalLoader EvaluableLoader,
+	concurrency int64,
 ) AggregateEvaluator {
+	if concurrency <= 0 {
+		concurrency = 1
+	}
+
 	return &aggregator{
 		eval:   eval,
 		parser: parser,
 		loader: evalLoader,
+		sem:    semaphore.NewWeighted(concurrency),
 		engines: map[EngineType]MatchingEngine{
 			EngineTypeStringHash: newStringEqualityMatcher(),
 			EngineTypeNullMatch:  newNullMatcher(),
@@ -92,6 +99,8 @@ type aggregator struct {
 	// engines records all engines
 	engines map[EngineType]MatchingEngine
 
+	sem *semaphore.Weighted
+
 	// lock prevents concurrent updates of data
 	lock *sync.RWMutex
 	// len stores the current len of aggregable expressions.
@@ -123,6 +132,7 @@ func (a *aggregator) Evaluate(ctx context.Context, data map[string]any) ([]Evalu
 		err     error
 		matched = int32(0)
 		result  = []Evaluable{}
+		wg      sync.WaitGroup
 	)
 
 	// TODO: Concurrently match constant expressions using a semaphore for capacity.
@@ -131,25 +141,42 @@ func (a *aggregator) Evaluate(ctx context.Context, data map[string]any) ([]Evalu
 	if err != nil {
 		return nil, 0, err
 	}
-	for _, expr := range constantEvals {
-		atomic.AddInt32(&matched, 1)
 
-		if expr.GetExpression() == "" {
-			result = append(result, expr)
-			continue
+	for _, item := range constantEvals {
+		if err := a.sem.Acquire(ctx, 1); err != nil {
+			return result, matched, err
 		}
 
-		// NOTE: We don't need to add lifted expression variables,
-		// because match.Parsed.Evaluable() returns the original expression
-		// string.
-		ok, evalerr := a.eval(ctx, expr, data)
-		if evalerr != nil {
-			err = errors.Join(err, evalerr)
-			continue
-		}
-		if ok {
-			result = append(result, expr)
-		}
+		expr := item
+		wg.Add(1)
+		go func() {
+			defer a.sem.Release(1)
+			defer wg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					err = errors.Join(err, fmt.Errorf("recovered from panic in evaluate: %v", r))
+				}
+			}()
+
+			atomic.AddInt32(&matched, 1)
+
+			if expr.GetExpression() == "" {
+				result = append(result, expr)
+				return
+			}
+
+			// NOTE: We don't need to add lifted expression variables,
+			// because match.Parsed.Evaluable() returns the original expression
+			// string.
+			ok, evalerr := a.eval(ctx, expr, data)
+			if evalerr != nil {
+				err = errors.Join(err, evalerr)
+				return
+			}
+			if ok {
+				result = append(result, expr)
+			}
+		}()
 	}
 
 	matches, merr := a.AggregateMatch(ctx, data)
@@ -178,23 +205,41 @@ func (a *aggregator) Evaluate(ctx context.Context, data map[string]any) ([]Evalu
 			continue
 		}
 
-		atomic.AddInt32(&matched, 1)
-		// NOTE: We don't need to add lifted expression variables,
-		// because match.Parsed.Evaluable() returns the original expression
-		// string.
-		ok, evalerr := a.eval(ctx, match, data)
-
-		seen[match.GetID()] = struct{}{}
-
-		if evalerr != nil {
-			err = errors.Join(err, evalerr)
-			continue
-		}
-		if ok {
-			result = append(result, match)
+		if err := a.sem.Acquire(ctx, 1); err != nil {
+			return result, matched, err
 		}
+
+		expr := match
+		wg.Add(1)
+		go func() {
+			defer a.sem.Release(1)
+			defer wg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					err = errors.Join(err, fmt.Errorf("recovered from panic in evaluate: %v", r))
+				}
+			}()
+
+			atomic.AddInt32(&matched, 1)
+			// NOTE: We don't need to add lifted expression variables,
+			// because match.Parsed.Evaluable() returns the original expression
+			// string.
+			ok, evalerr := a.eval(ctx, expr, data)
+
+			seen[expr.GetID()] = struct{}{}
+
+			if evalerr != nil {
+				err = errors.Join(err, evalerr)
+				return
+			}
+			if ok {
+				result = append(result, expr)
+			}
+		}()
 	}
 
+	wg.Wait()
+
 	return result, matched, err
 }