Add matchConditions to Stage

pkg/apis/internalversion/stage_types.go

@@ -137,12 +137,20 @@ type StageSelector struct {
 	MatchAnnotations map[string]string
 	// MatchExpressions is a list of label selector requirements. The requirements are ANDed.
 	MatchExpressions []SelectorRequirement
+	// MatchConditions is a list of label selector conditions. The conditions are ANDed.
+	MatchConditions []SelectorCondition
+}
+
+// SelectorCondition is a resource selector condition is a set of conditions that must be true for a match.
+type SelectorCondition struct {
+	// Expression represents the expression which will be evaluated by CEL.
+	Expression string
 }
 
 // SelectorRequirement is a resource selector requirement is a selector that contains values, a key,
 // and an operator that relates the key and values.
 type SelectorRequirement struct {
-	// The name of the scope that the selector applies to.
+	// Key represents the expression which will be evaluated by JQ.
 	Key string
 	// Represents a scope's relationship to a set of values.
 	Operator SelectorOperator

pkg/apis/v1alpha1/stage_types.go

@@ -176,12 +176,20 @@ type StageSelector struct {
 	MatchAnnotations map[string]string `json:"matchAnnotations,omitempty"`
 	// MatchExpressions is a list of label selector requirements. The requirements are ANDed.
 	MatchExpressions []SelectorRequirement `json:"matchExpressions,omitempty"`
+	// MatchConditions is a list of label selector conditions. The conditions are ANDed.
+	MatchConditions []SelectorCondition `json:"matchConditions,omitempty"`
+}
+
+// SelectorCondition is a resource selector condition is a set of conditions that must be true for a match.
+type SelectorCondition struct {
+	// Expression represents the expression which will be evaluated by CEL.
+	Expression string `json:"expression"`
 }
 
 // SelectorRequirement is a resource selector requirement is a selector that contains values, a key,
 // and an operator that relates the key and values.
 type SelectorRequirement struct {
-	// The name of the scope that the selector applies to.
+	// Key represents the expression which will be evaluated by JQ.
 	Key string `json:"key"`
 	// Represents a scope's relationship to a set of values.
 	Operator SelectorOperator `json:"operator"`

pkg/kwok/controllers/controller.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"maps"
 	"os"
 	"strconv"
 	"strings"
@@ -47,6 +48,7 @@ import (
 	"sigs.k8s.io/kwok/pkg/config/resources"
 	"sigs.k8s.io/kwok/pkg/consts"
 	"sigs.k8s.io/kwok/pkg/log"
+	"sigs.k8s.io/kwok/pkg/utils/cel"
 	"sigs.k8s.io/kwok/pkg/utils/client"
 	"sigs.k8s.io/kwok/pkg/utils/gotpl"
 	"sigs.k8s.io/kwok/pkg/utils/informer"
@@ -103,6 +105,7 @@ var (
 // Controller is a fake kubelet implementation that can be used to test
 type Controller struct {
 	conf Config
+	env  *cel.Environment
 
 	stagesManager *StagesManager
 
@@ -198,7 +201,26 @@ func NewController(conf Config) (*Controller, error) {
 		return nil, err
 	}
 
+	types := slices.Clone(cel.DefaultTypes)
+	conversions := slices.Clone(cel.DefaultConversions)
+	funcs := maps.Clone(cel.DefaultFuncs)
+	methods := maps.Clone(cel.FuncsToMethods(cel.DefaultFuncs))
+	vars := map[string]any{
+		"self": nil,
+	}
+	env, err := cel.NewEnvironment(cel.EnvironmentConfig{
+		Types:       types,
+		Conversions: conversions,
+		Methods:     methods,
+		Funcs:       funcs,
+		Vars:        vars,
+	})
+	if err != nil {
+		return nil, err
+	}
+
 	c := &Controller{
+		env:  env,
 		conf: conf,
 	}
 
@@ -581,7 +603,7 @@ func (c *Controller) Start(ctx context.Context) error {
 
 	if len(c.conf.LocalStages) != 0 {
 		for ref, stage := range c.conf.LocalStages {
-			lifecycle, err := NewLifecycle(stage)
+			lifecycle, err := NewLifecycle(stage, c.env)
 			if err != nil {
 				return err
 			}

pkg/kwok/controllers/lifecycle.go

@@ -25,15 +25,16 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 
 	"sigs.k8s.io/kwok/pkg/apis/internalversion"
+	"sigs.k8s.io/kwok/pkg/utils/cel"
 	"sigs.k8s.io/kwok/pkg/utils/expression"
 	"sigs.k8s.io/kwok/pkg/utils/format"
 )
 
 // NewLifecycle returns a new Lifecycle.
-func NewLifecycle(stages []*internalversion.Stage) (Lifecycle, error) {
+func NewLifecycle(stages []*internalversion.Stage, env *cel.Environment) (Lifecycle, error) {
 	lcs := Lifecycle{}
 	for _, stage := range stages {
-		lc, err := NewLifecycleStage(stage)
+		lc, err := NewLifecycleStage(stage, env)
 		if err != nil {
 			return nil, fmt.Errorf("lifecycle stage: %w", err)
 		}
@@ -48,10 +49,10 @@ func NewLifecycle(stages []*internalversion.Stage) (Lifecycle, error) {
 // Lifecycle is a list of lifecycle stage.
 type Lifecycle []*LifecycleStage
 
-func (s Lifecycle) match(label, annotation labels.Set, data interface{}) ([]*LifecycleStage, error) {
+func (s Lifecycle) match(ctx context.Context, label, annotation labels.Set, data, jsonStandard any) ([]*LifecycleStage, error) {
 	out := []*LifecycleStage{}
 	for _, stage := range s {
-		ok, err := stage.match(label, annotation, data)
+		ok, err := stage.match(ctx, label, annotation, data, jsonStandard)
 		if err != nil {
 			return nil, err
 		}
@@ -63,12 +64,8 @@ func (s Lifecycle) match(label, annotation labels.Set, data interface{}) ([]*Lif
 }
 
 // Match returns matched stage.
-func (s Lifecycle) Match(label, annotation labels.Set, data interface{}) (*LifecycleStage, error) {
-	data, err := expression.ToJSONStandard(data)
-	if err != nil {
-		return nil, err
-	}
-	stages, err := s.match(label, annotation, data)
+func (s Lifecycle) Match(ctx context.Context, label, annotation labels.Set, data, jsonStandard any) (*LifecycleStage, error) {
+	stages, err := s.match(ctx, label, annotation, data, jsonStandard)
 	if err != nil {
 		return nil, err
 	}
@@ -102,7 +99,7 @@ func (s Lifecycle) Match(label, annotation labels.Set, data interface{}) (*Lifec
 }
 
 // NewLifecycleStage returns a new LifecycleStage.
-func NewLifecycleStage(s *internalversion.Stage) (*LifecycleStage, error) {
+func NewLifecycleStage(s *internalversion.Stage, env *cel.Environment) (*LifecycleStage, error) {
 	stage := &LifecycleStage{
 		name: s.Name,
 	}
@@ -126,6 +123,15 @@ func NewLifecycleStage(s *internalversion.Stage) (*LifecycleStage, error) {
 			stage.matchExpressions = append(stage.matchExpressions, requirement)
 		}
 	}
+	if selector.MatchConditions != nil {
+		for _, express := range selector.MatchConditions {
+			program, err := env.Compile(express.Expression)
+			if err != nil {
+				return nil, err
+			}
+			stage.matchConditions = append(stage.matchConditions, program)
+		}
+	}
 
 	stage.next = &s.Spec.Next
 	if delay := s.Spec.Delay; delay != nil {
@@ -177,6 +183,7 @@ type LifecycleStage struct {
 	matchLabels      labels.Selector
 	matchAnnotations labels.Selector
 	matchExpressions []*expression.Requirement
+	matchConditions  []cel.Program
 
 	weight int
 	next   *internalversion.StageNext
@@ -187,7 +194,7 @@ type LifecycleStage struct {
 	immediateNextStage bool
 }
 
-func (s *LifecycleStage) match(label, annotation labels.Set, jsonStandard interface{}) (bool, error) {
+func (s *LifecycleStage) match(ctx context.Context, label, annotation labels.Set, data, jsonStandard any) (bool, error) {
 	if s.matchLabels != nil {
 		if !s.matchLabels.Matches(label) {
 			return false, nil
@@ -201,7 +208,7 @@ func (s *LifecycleStage) match(label, annotation labels.Set, jsonStandard interf
 
 	if s.matchExpressions != nil {
 		for _, requirement := range s.matchExpressions {
-			ok, err := requirement.Matches(context.Background(), jsonStandard)
+			ok, err := requirement.Matches(ctx, jsonStandard)
 			if err != nil {
 				return false, err
 			}
@@ -210,6 +217,25 @@ func (s *LifecycleStage) match(label, annotation labels.Set, jsonStandard interf
 			}
 		}
 	}
+
+	if s.matchConditions != nil {
+		for _, program := range s.matchConditions {
+			val, _, err := program.ContextEval(ctx, map[string]any{
+				"self": data,
+			})
+			if err != nil {
+				return false, err
+			}
+			ok, err := cel.AsBool(val)
+			if err != nil {
+				return false, err
+			}
+			if !ok {
+				return false, nil
+			}
+		}
+	}
+
 	return true, nil
 }
 

pkg/kwok/controllers/node_controller.go

@@ -352,7 +352,7 @@ func (c *NodeController) preprocess(ctx context.Context, node *corev1.Node) erro
 	}
 
 	lifecycle := c.lifecycle.Get()
-	stage, err := lifecycle.Match(node.Labels, node.Annotations, data)
+	stage, err := lifecycle.Match(ctx, node.Labels, node.Annotations, node, data)
 	if err != nil {
 		return fmt.Errorf("stage match: %w", err)
 	}

pkg/kwok/controllers/node_controller_test.go

@@ -78,7 +78,7 @@ func TestNodeController(t *testing.T) {
 	nodeInit, _ := config.UnmarshalWithType[*internalversion.Stage](nodefast.DefaultNodeInit)
 	nodeStages := []*internalversion.Stage{nodeInit}
 
-	lifecycle, _ := NewLifecycle(nodeStages)
+	lifecycle, _ := NewLifecycle(nodeStages, nil)
 	nodes, err := NewNodeController(NodeControllerConfig{
 		TypedClient:          clientset,
 		NodeIP:               "10.0.0.1",

pkg/kwok/controllers/pod_controller.go

@@ -248,7 +248,7 @@ func (c *PodController) preprocess(ctx context.Context, pod *corev1.Pod) error {
 	}
 
 	lifecycle := c.lifecycle.Get()
-	stage, err := lifecycle.Match(pod.Labels, pod.Annotations, data)
+	stage, err := lifecycle.Match(ctx, pod.Labels, pod.Annotations, pod, data)
 	if err != nil {
 		return fmt.Errorf("stage match: %w", err)
 	}

pkg/kwok/controllers/pod_controller_test.go

@@ -190,7 +190,7 @@ func TestPodController(t *testing.T) {
 		t.Fatal(fmt.Errorf("failed to watch nodes: %w", err))
 	}
 
-	lifecycle, _ := NewLifecycle(podStages)
+	lifecycle, _ := NewLifecycle(podStages, nil)
 	annotationSelector, _ := labels.Parse("fake=custom")
 	pods, err := NewPodController(PodControllerConfig{
 		TypedClient:                           clientset,