podprobemarker support serverless pod

Signed-off-by: liheng.zms <[email protected]>

apis/apps/v1alpha1/pod_probe_marker_types.go

@@ -21,6 +21,28 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
+const (
+	// PodProbeMarkerAnnotationKey records the Probe Spec, mainly used for serverless Pod scenarios, as follows:
+	//  annotations:
+	//    kruise.io/podprobemarker: |
+	//      [
+	//          {
+	//              "containerName": "minecraft",
+	//              "name": "healthy",
+	//              "podConditionType": "game.kruise.io/healthy",
+	//              "probe": {
+	//                  "exec": {
+	//                      "command": [
+	//                          "bash",
+	//                          "/data/probe.sh"
+	//                      ]
+	//                  }
+	//              }
+	//          }
+	//      ]
+	PodProbeMarkerAnnotationKey = "kruise.io/podprobemarker"
+)
+
 // PodProbeMarkerSpec defines the desired state of PodProbeMarker
 type PodProbeMarkerSpec struct {
 	// Selector is a label query over pods that should exec custom probe

pkg/util/pods.go

@@ -402,3 +402,13 @@ func GetPodContainerByName(cName string, pod *v1.Pod) *v1.Container {
 
 	return nil
 }
+
+// IsRestartableInitContainer returns true if the initContainer has
+// ContainerRestartPolicyAlways.
+func IsRestartableInitContainer(initContainer *v1.Container) bool {
+	if initContainer.RestartPolicy == nil {
+		return false
+	}
+
+	return *initContainer.RestartPolicy == v1.ContainerRestartPolicyAlways
+}

pkg/webhook/pod/mutating/pod_probe_marker.go

@@ -0,0 +1,93 @@
+/*
+Copyright 2024 The Kruise 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 mutating
+
+import (
+	"context"
+	"fmt"
+
+	appsv1alpha1 "github.com/openkruise/kruise/apis/apps/v1alpha1"
+	"github.com/openkruise/kruise/pkg/util"
+	utilclient "github.com/openkruise/kruise/pkg/util/client"
+	admissionv1 "k8s.io/api/admission/v1"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/klog/v2"
+	"k8s.io/kube-openapi/pkg/util/sets"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+)
+
+// mutating relate-pub annotation in pod
+func (h *PodCreateHandler) podProbeMakerMutatingPod(ctx context.Context, req admission.Request, pod *corev1.Pod) (skip bool, err error) {
+	if len(req.AdmissionRequest.SubResource) > 0 || req.AdmissionRequest.Operation != admissionv1.Create ||
+		req.AdmissionRequest.Resource.Resource != "pods" {
+		return true, nil
+	}
+	ppmList := &appsv1alpha1.PodProbeMarkerList{}
+	if err = h.Client.List(context.TODO(), ppmList, &client.ListOptions{Namespace: pod.Namespace}, utilclient.DisableDeepCopy); err != nil {
+		return false, err
+	} else if len(ppmList.Items) == 0 {
+		return true, nil
+	}
+
+	containers := sets.NewString()
+	for _, c := range pod.Spec.Containers {
+		containers.Insert(c.Name)
+	}
+	for _, c := range pod.Spec.InitContainers {
+		if util.IsRestartableInitContainer(&c) {
+			containers.Insert(c.Name)
+		}
+	}
+
+	matchedProbeKey := sets.NewString()
+	matchedConditions := sets.NewString()
+	matchedProbes := make([]appsv1alpha1.PodContainerProbe, 0)
+	for _, obj := range ppmList.Items {
+		// This error is irreversible, so continue
+		labelSelector, err := util.ValidatedLabelSelectorAsSelector(obj.Spec.Selector)
+		if err != nil {
+			continue
+		}
+		// If a PUB with a nil or empty selector creeps in, it should match nothing, not everything.
+		if labelSelector.Empty() || !labelSelector.Matches(labels.Set(pod.Labels)) {
+			continue
+		}
+		for i := range obj.Spec.Probes {
+			probe := obj.Spec.Probes[i]
+			key := fmt.Sprintf("%s/%s", probe.ContainerName, probe.Name)
+			if matchedConditions.Has(probe.PodConditionType) || matchedProbeKey.Has(key) || !containers.Has(probe.ContainerName) || probe.PodConditionType == "" {
+				continue
+			}
+			matchedProbes = append(matchedProbes, probe)
+			matchedProbeKey.Insert(key)
+			matchedConditions.Insert(probe.PodConditionType)
+		}
+	}
+	if len(matchedProbes) == 0 {
+		return true, nil
+	}
+
+	if pod.Annotations == nil {
+		pod.Annotations = map[string]string{}
+	}
+	body := util.DumpJSON(matchedProbes)
+	pod.Annotations[appsv1alpha1.PodProbeMarkerAnnotationKey] = body
+	klog.V(3).InfoS("mutating add pod annotation", "namespace", pod.Namespace, "name", pod.Name, "key", appsv1alpha1.PodProbeMarkerAnnotationKey, "value", body)
+	return false, nil
+}

pkg/webhook/pod/mutating/pod_probe_marker_test.go

@@ -0,0 +1,229 @@
+/*
+Copyright 2024 The Kruise 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 mutating
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	appsv1alpha1 "github.com/openkruise/kruise/apis/apps/v1alpha1"
+	"github.com/openkruise/kruise/pkg/util"
+	admissionv1 "k8s.io/api/admission/v1"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/kubernetes/scheme"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+)
+
+func TestPodProbeMakerMutatingPod(t *testing.T) {
+	cases := []struct {
+		name               string
+		getPod             func() *v1.Pod
+		getPodProbeMarkers func() []*appsv1alpha1.PodProbeMarker
+		expected           map[string]string
+	}{
+		{
+			name: "podprobemarker, selector matched, but no conditionType",
+			getPod: func() *v1.Pod {
+				al := v1.ContainerRestartPolicyAlways
+				return &v1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							"version": "test",
+						},
+						Labels: map[string]string{
+							"app": "web",
+						},
+						Namespace: "test",
+						Name:      "pod-1",
+					},
+					Spec: v1.PodSpec{
+						InitContainers: []v1.Container{
+							{
+								Name: "init-1",
+							},
+							{
+								Name:          "init-2",
+								RestartPolicy: &al,
+							},
+						},
+						Containers: []v1.Container{
+							{
+								Name: "main",
+							},
+							{
+								Name: "envoy",
+							},
+						},
+					},
+				}
+			},
+			getPodProbeMarkers: func() []*appsv1alpha1.PodProbeMarker {
+				obj1 := &appsv1alpha1.PodProbeMarker{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "healthy",
+						Namespace: "test",
+					},
+					Spec: appsv1alpha1.PodProbeMarkerSpec{
+						Selector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								"app": "web",
+							},
+						},
+						Probes: []appsv1alpha1.PodContainerProbe{
+							{
+								ContainerName: "main",
+								Name:          "healthy",
+							},
+							{
+								ContainerName:    "invalid",
+								Name:             "healthy",
+								PodConditionType: "game.kruise.io/healthy",
+							},
+						},
+					},
+				}
+
+				return []*appsv1alpha1.PodProbeMarker{obj1}
+			},
+			expected: map[string]string{
+				"version": "test",
+			},
+		},
+		{
+			name: "podprobemarker, selector matched",
+			getPod: func() *v1.Pod {
+				al := v1.ContainerRestartPolicyAlways
+				return &v1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							"version": "test",
+						},
+						Labels: map[string]string{
+							"app": "web",
+						},
+						Namespace: "test",
+						Name:      "pod-1",
+					},
+					Spec: v1.PodSpec{
+						InitContainers: []v1.Container{
+							{
+								Name: "init-1",
+							},
+							{
+								Name:          "init-2",
+								RestartPolicy: &al,
+							},
+						},
+						Containers: []v1.Container{
+							{
+								Name: "main",
+							},
+							{
+								Name: "envoy",
+							},
+						},
+					},
+				}
+			},
+			getPodProbeMarkers: func() []*appsv1alpha1.PodProbeMarker {
+				obj1 := &appsv1alpha1.PodProbeMarker{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "healthy",
+						Namespace: "test",
+					},
+					Spec: appsv1alpha1.PodProbeMarkerSpec{
+						Selector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								"app": "web",
+							},
+						},
+						Probes: []appsv1alpha1.PodContainerProbe{
+							{
+								ContainerName:    "main",
+								Name:             "healthy",
+								PodConditionType: "game.kruise.io/healthy",
+							},
+							{
+								ContainerName:    "envoy",
+								Name:             "healthy",
+								PodConditionType: "game.kruise.io/healthy",
+							},
+						},
+					},
+				}
+				obj2 := &appsv1alpha1.PodProbeMarker{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "init",
+						Namespace: "test",
+					},
+					Spec: appsv1alpha1.PodProbeMarkerSpec{
+						Selector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								"app": "web",
+							},
+						},
+						Probes: []appsv1alpha1.PodContainerProbe{
+							{
+								ContainerName:    "init-1",
+								Name:             "init",
+								PodConditionType: "game.kruise.io/init",
+							},
+							{
+								ContainerName:    "init-2",
+								Name:             "init",
+								PodConditionType: "game.kruise.io/init",
+							},
+						},
+					},
+				}
+				return []*appsv1alpha1.PodProbeMarker{obj1, obj2}
+			},
+			expected: map[string]string{
+				"version":                                "test",
+				appsv1alpha1.PodProbeMarkerAnnotationKey: `[{"name":"healthy","containerName":"main","probe":{},"podConditionType":"game.kruise.io/healthy"},{"name":"init","containerName":"init-2","probe":{},"podConditionType":"game.kruise.io/init"}]`,
+			},
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			if c.name != "podprobemarker, selector matched" {
+				return
+			}
+			decoder := admission.NewDecoder(scheme.Scheme)
+			builder := fake.NewClientBuilder()
+			for i := range c.getPodProbeMarkers() {
+				obj := c.getPodProbeMarkers()[i]
+				builder.WithObjects(obj)
+			}
+			testClient := builder.Build()
+			podHandler := &PodCreateHandler{Decoder: decoder, Client: testClient}
+			req := newAdmission(admissionv1.Create, runtime.RawExtension{}, runtime.RawExtension{}, "")
+			pod := c.getPod()
+			if _, err := podHandler.podProbeMakerMutatingPod(context.Background(), req, pod); err != nil {
+				t.Fatalf("failed to mutating pod, err: %v", err)
+			}
+			if !reflect.DeepEqual(c.expected, pod.Annotations) {
+				t.Fatalf("expected: %s, got: %s", util.DumpJSON(c.expected), util.DumpJSON(pod.Annotations))
+			}
+		})
+	}
+}