Merge pull request #1162 from wzshiming/feat/weight-from

Add weightFrom to Stages

kustomize/crd/bases/kwok.x-k8s.io_stages.yaml

@@ -273,6 +273,17 @@ spec:
                   a random stage will be matched as the next stage based on the weight.
                 minimum: 0
                 type: integer
+              weightFrom:
+                description: |-
+                  WeightFrom means is the expression used to get the value.
+                  If it is a number type, convert to int.
+                  If it is a string type, the value get will be parsed by strconv.ParseInt.
+                properties:
+                  expressionFrom:
+                    description: ExpressionFrom is the expression used to get the
+                      value.
+                    type: string
+                type: object
             required:
             - next
             - resourceRef

kustomize/stage/node/chaos/node-not-ready.yaml

@@ -24,6 +24,8 @@ spec:
     jitterDurationFrom:
       expressionFrom: '.metadata.annotations["node-not-ready.stage.kwok.x-k8s.io/jitter-delay"]'
   weight: 10000
+  weightFrom:
+    expressionFrom: '.metadata.annotations["node-not-ready.stage.kwok.x-k8s.io/weight"]'
   next:
     statusTemplate: |
       {{ $now := Now }}

kustomize/stage/node/fast/testdata/node-pending.output.yaml

@@ -68,3 +68,4 @@ stages:
     subresource: status
     type: application/merge-patch+json
   stage: node-initialize
+  weight: 0

kustomize/stage/pod/chaos/pod-container-running-failed.yaml

@@ -26,6 +26,8 @@ spec:
     jitterDurationFrom:
       expressionFrom: '.metadata.annotations["pod-container-running-failed.stage.kwok.x-k8s.io/jitter-delay"]'
   weight: 10000
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-container-running-failed.stage.kwok.x-k8s.io/weight"]'
   next:
     statusTemplate: |
       {{ $now := Now }}

kustomize/stage/pod/chaos/pod-init-container-running-failed.yaml

@@ -26,6 +26,8 @@ spec:
     jitterDurationFrom:
       expressionFrom: '.metadata.annotations["pod-init-container-running-failed.stage.kwok.x-k8s.io/jitter-delay"]'
   weight: 10000
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-init-container-running-failed.stage.kwok.x-k8s.io/weight"]'
   next:
     statusTemplate: |
       {{ $now := Now }}

kustomize/stage/pod/fast/testdata/pod-pending.output.yaml

@@ -32,3 +32,4 @@ stages:
     subresource: status
     type: application/merge-patch+json
   stage: pod-ready
+  weight: 0

kustomize/stage/pod/fast/testdata/pod-running-with-job.output.yaml

@@ -22,3 +22,4 @@ stages:
     subresource: status
     type: application/merge-patch+json
   stage: pod-complete
+  weight: 0

kustomize/stage/pod/fast/testdata/pod-terminating.output.yaml

@@ -5,3 +5,4 @@ stages:
 - next:
   - kind: delete
   stage: pod-delete
+  weight: 0

kustomize/stage/pod/general/pod-complete.yaml

@@ -23,6 +23,8 @@ spec:
       values:
       - 'Job'
   weight: 1
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-complete.stage.kwok.x-k8s.io/weight"]'
   delay:
     durationMilliseconds: 1000
     durationFrom:

kustomize/stage/pod/general/pod-create.yaml

@@ -13,6 +13,8 @@ spec:
     - key: '.status.podIP'
       operator: 'DoesNotExist'
   weight: 1
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-create.stage.kwok.x-k8s.io/weight"]'
   delay:
     durationMilliseconds: 1000
     durationFrom:

kustomize/stage/pod/general/pod-delete.yaml

@@ -13,6 +13,8 @@ spec:
     - key: '.metadata.finalizers'
       operator: 'DoesNotExist'
   weight: 1
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-delete.stage.kwok.x-k8s.io/weight"]'
   delay:
     durationMilliseconds: 1000
     durationFrom:

kustomize/stage/pod/general/pod-init-container-completed.yaml

@@ -17,6 +17,8 @@ spec:
     - key: '.status.initContainerStatuses.[].state.running.startedAt'
       operator: 'Exists'
   weight: 1
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-init-container-completed.stage.kwok.x-k8s.io/weight"]'
   delay:
     durationMilliseconds: 1000
     durationFrom:

kustomize/stage/pod/general/pod-init-container-running.yaml

@@ -21,6 +21,8 @@ spec:
     - key: '.status.initContainerStatuses.[].state.waiting.reason'
       operator: 'Exists'
   weight: 1
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-init-container-running.stage.kwok.x-k8s.io/weight"]'
   delay:
     durationMilliseconds: 1000
     durationFrom:

kustomize/stage/pod/general/pod-ready.yaml

@@ -17,6 +17,8 @@ spec:
     - key: '.status.containerStatuses.[].state.running.startedAt'
       operator: 'DoesNotExist'
   weight: 1
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-ready.stage.kwok.x-k8s.io/weight"]'
   delay:
     durationMilliseconds: 1000
     durationFrom:

kustomize/stage/pod/general/pod-remove-finalizer.yaml

@@ -15,6 +15,8 @@ spec:
       values:
       - 'kwok.x-k8s.io/fake'
   weight: 1
+  weightFrom:
+    expressionFrom: '.metadata.annotations["pod-remove-finalizer.stage.kwok.x-k8s.io/weight"]'
   delay:
     durationMilliseconds: 1000
     durationFrom:

pkg/apis/internalversion/stage_types.go

@@ -38,6 +38,10 @@ type StageSpec struct {
 	// Weight means when multiple stages share the same ResourceRef and Selector,
 	// a random stage will be matched as the next stage based on the weight.
 	Weight int
+	// WeightFrom means is the expression used to get the value.
+	// If it is a number type, convert to int.
+	// If it is a string type, the value get will be parsed by strconv.ParseInt.
+	WeightFrom *ExpressionFromSource
 	// Delay means there is a delay in this stage.
 	Delay *StageDelay
 	// Next indicates that this stage will be moved to.

pkg/apis/v1alpha1/stage_types.go

@@ -69,6 +69,10 @@ type StageSpec struct {
 	// +kubebuilder:default=0
 	// +kubebuilder:validation:Minimum=0
 	Weight int `json:"weight,omitempty"`
+	// WeightFrom means is the expression used to get the value.
+	// If it is a number type, convert to int.
+	// If it is a string type, the value get will be parsed by strconv.ParseInt.
+	WeightFrom *ExpressionFromSource `json:"weightFrom,omitempty"`
 	// Delay means there is a delay in this stage.
 	Delay *StageDelay `json:"delay,omitempty"`
 	// Next indicates that this stage will be moved to.

pkg/kwok/controllers/node_controller.go

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

pkg/kwok/controllers/pod_controller.go

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

pkg/kwok/controllers/stage_controller.go

@@ -196,7 +196,7 @@ func (c *StageController) preprocess(ctx context.Context, resource *unstructured
 	}
 
 	lc := c.lifecycle.Get()
-	stage, err := lc.Match(resource.GetLabels(), resource.GetAnnotations(), data)
+	stage, err := lc.Match(ctx, resource.GetLabels(), resource.GetAnnotations(), data)
 	if err != nil {
 		return fmt.Errorf("stage match: %w", err)
 	}

pkg/tools/stage/stage.go

@@ -65,7 +65,7 @@ func TestingStages(ctx context.Context, target any, stages []*internalversion.St
 		return nil, err
 	}
 
-	lcstages, err := lc.ListAllPossible(testTarget.GetLabels(), testTarget.GetAnnotations(), testTarget)
+	lcstages, err := lc.ListAllPossible(ctx, testTarget.GetLabels(), testTarget.GetAnnotations(), testTarget)
 	if err != nil {
 		return nil, err
 	}
@@ -94,8 +94,8 @@ func testingStage(ctx context.Context, testTarget Obj, stage *lifecycle.Stage) (
 		meta["delay"] = delay
 	}
 
-	weight := stage.Weight()
-	if weight != 0 {
+	weight, ok := stage.Weight(ctx, testTarget)
+	if ok {
 		meta["weight"] = weight
 	}
 

pkg/utils/expression/value_int_from.go

@@ -0,0 +1,94 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package expression
+
+import (
+	"context"
+	"strconv"
+)
+
+// IntGetter is a interface that can be used to get a int64 value.
+type IntGetter interface {
+	// Get returns a int64 value.
+	Get(ctx context.Context, v interface{}) (int64, bool)
+}
+
+type int64From struct {
+	value *int64
+	query *Query
+}
+
+// NewIntFrom returns a new IntGetter.
+func NewIntFrom(value *int64, src *string) (IntGetter, error) {
+	if value == nil && src == nil {
+		return int64Noop{}, nil
+	}
+	if src == nil {
+		return intWrapper(*value), nil
+	}
+	query, err := NewQuery(*src)
+	if err != nil {
+		return nil, err
+	}
+	return &int64From{
+		value: value,
+		query: query,
+	}, nil
+}
+
+func (d *int64From) Get(ctx context.Context, v interface{}) (int64, bool) {
+	out, err := d.query.Execute(ctx, v)
+	if err != nil {
+		return 0, false
+	}
+	if len(out) == 0 {
+		if d.value != nil {
+			return *d.value, true
+		}
+		return 0, false
+	}
+	switch t := out[0].(type) {
+	case string:
+		if t == "" {
+			return 0, false
+		}
+		n, err := strconv.ParseInt(t, 0, 0)
+		if err == nil {
+			return n, true
+		}
+		return 0, false
+	case float64: // TODO: don't use float
+		return int64(t), true
+	}
+	if d.value != nil {
+		return *d.value, true
+	}
+	return 0, false
+}
+
+type int64Noop struct {
+}
+
+func (int64Noop) Get(ctx context.Context, v interface{}) (int64, bool) {
+	return 0, false
+}
+
+type intWrapper int64
+
+func (i intWrapper) Get(ctx context.Context, v interface{}) (int64, bool) {
+	return int64(i), true
+}