support tcpCheck in podProbe (openkruise#1474)

Co-authored-by: jicheng.sk <jicheng.sk@alibaba-inc.com>
ppbits · Feb 28, 2024 · e93fca5 · e93fca5
1 parent b7cd616
commit e93fca5
Show file tree

Hide file tree

Showing 20 changed files with 3,382 additions and 138 deletions.
diff --git a/apis/apps/v1alpha1/node_pod_probe_types.go b/apis/apps/v1alpha1/node_pod_probe_types.go
@@ -32,6 +32,8 @@ type PodProbe struct {
 	Namespace string `json:"namespace"`
 	// pod uid
 	UID string `json:"uid"`
+	// pod ip
+	IP string `json:"IP"`
 	// Custom container probe, supports Exec, Tcp, and returns the result to Pod yaml
 	Probes []ContainerProbe `json:"probes,omitempty"`
 }

diff --git a/config/crd/bases/apps.kruise.io_nodepodprobes.yaml b/config/crd/bases/apps.kruise.io_nodepodprobes.yaml
@@ -38,6 +38,9 @@ spec:
               podProbes:
                 items:
                   properties:
+                    IP:
+                      description: pod ip
+                      type: string
                     name:
                       description: pod name
                       type: string
@@ -221,6 +224,7 @@ spec:
                       description: pod uid
                       type: string
                   required:
+                  - IP
                   - name
                   - namespace
                   - uid

diff --git a/pkg/controller/podprobemarker/pod_probe_marker_controller.go b/pkg/controller/podprobemarker/pod_probe_marker_controller.go
@@ -23,17 +23,10 @@ import (
 	"reflect"
 	"strings"
 
-	appsv1alpha1 "github.com/openkruise/kruise/apis/apps/v1alpha1"
-	"github.com/openkruise/kruise/pkg/features"
-	"github.com/openkruise/kruise/pkg/util"
-	utilclient "github.com/openkruise/kruise/pkg/util/client"
-	"github.com/openkruise/kruise/pkg/util/controllerfinder"
-	utildiscovery "github.com/openkruise/kruise/pkg/util/discovery"
-	utilfeature "github.com/openkruise/kruise/pkg/util/feature"
-	"github.com/openkruise/kruise/pkg/util/ratelimiter"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/client-go/util/retry"
 	"k8s.io/klog/v2"
 	kubecontroller "k8s.io/kubernetes/pkg/controller"
@@ -44,6 +37,15 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 	"sigs.k8s.io/controller-runtime/pkg/source"
+
+	appsv1alpha1 "github.com/openkruise/kruise/apis/apps/v1alpha1"
+	"github.com/openkruise/kruise/pkg/features"
+	"github.com/openkruise/kruise/pkg/util"
+	utilclient "github.com/openkruise/kruise/pkg/util/client"
+	"github.com/openkruise/kruise/pkg/util/controllerfinder"
+	utildiscovery "github.com/openkruise/kruise/pkg/util/discovery"
+	utilfeature "github.com/openkruise/kruise/pkg/util/feature"
+	"github.com/openkruise/kruise/pkg/util/ratelimiter"
 )
 
 func init() {
@@ -234,15 +236,47 @@ func (r *ReconcilePodProbeMarker) updateNodePodProbes(ppm *appsv1alpha1.PodProbe
 				exist = true
 				for j := range ppm.Spec.Probes {
 					probe := ppm.Spec.Probes[j]
+					if podProbe.IP == "" {
+						podProbe.IP = pod.Status.PodIP
+					}
+					if probe.Probe.TCPSocket != nil {
+						probe, err = convertTcpSocketProbeCheckPort(probe, pod)
+						if err != nil {
+							klog.Errorf("Failed to convert tcpSocket probe port, err: %v, pod: %v/%v", err, pod.Namespace, pod.Name)
+							continue
+						}
+					}
+					if probe.Probe.HTTPGet != nil {
+						probe, err = convertHttpGetProbeCheckPort(probe, pod)
+						if err != nil {
+							klog.Errorf("Failed to convert httpGet probe port, err: %v, pod: %v/%v", err, pod.Namespace, pod.Name)
+							continue
+						}
+					}
 					setPodContainerProbes(podProbe, probe, ppm.Name)
 				}
 				break
 			}
 		}
 		if !exist {
-			podProbe := appsv1alpha1.PodProbe{Name: pod.Name, Namespace: pod.Namespace, UID: string(pod.UID)}
+			podProbe := appsv1alpha1.PodProbe{Name: pod.Name, Namespace: pod.Namespace, UID: string(pod.UID), IP: pod.Status.PodIP}
 			for j := range ppm.Spec.Probes {
 				probe := ppm.Spec.Probes[j]
+				// look up a port in a container by name & convert container name port
+				if probe.Probe.TCPSocket != nil {
+					probe, err = convertTcpSocketProbeCheckPort(probe, pod)
+					if err != nil {
+						klog.Errorf("Failed to convert tcpSocket probe port, err: %v, pod: %v/%v", err, pod.Namespace, pod.Name)
+						continue
+					}
+				}
+				if probe.Probe.HTTPGet != nil {
+					probe, err = convertHttpGetProbeCheckPort(probe, pod)
+					if err != nil {
+						klog.Errorf("Failed to convert httpGet probe port, err: %v, pod: %v/%v", err, pod.Namespace, pod.Name)
+						continue
+					}
+				}
 				podProbe.Probes = append(podProbe.Probes, appsv1alpha1.ContainerProbe{
 					Name:          fmt.Sprintf("%s#%s", ppm.Name, probe.Name),
 					ContainerName: probe.ContainerName,
@@ -266,6 +300,43 @@ func (r *ReconcilePodProbeMarker) updateNodePodProbes(ppm *appsv1alpha1.PodProbe
 	return nil
 }
 
+func convertHttpGetProbeCheckPort(probe appsv1alpha1.PodContainerProbe, pod *corev1.Pod) (appsv1alpha1.PodContainerProbe, error) {
+	probeNew := probe.DeepCopy()
+	if probe.Probe.HTTPGet.Port.Type == intstr.Int {
+		return *probeNew, nil
+	}
+	container := util.GetPodContainerByName(probe.ContainerName, pod)
+	if container == nil {
+		return *probeNew, fmt.Errorf("Failed to get container by name: %v in pod: %v/%v", probe.ContainerName, pod.Namespace, pod.Name)
+	}
+	portInt, err := util.ExtractPort(probe.Probe.HTTPGet.Port, *container)
+	if err != nil {
+		return *probeNew, fmt.Errorf("Failed to extract port for container: %v in pod: %v/%v", container.Name, pod.Namespace, pod.Name)
+	}
+	// If you need to parse integer values with specific bit sizes, avoid strconv.Atoi,
+	// and instead use strconv.ParseInt or strconv.ParseUint, which also allow specifying the bit size.
+	// https://codeql.github.com/codeql-query-help/go/go-incorrect-integer-conversion/
+	probeNew.Probe.HTTPGet.Port = intstr.FromInt(portInt)
+	return *probeNew, nil
+}
+
+func convertTcpSocketProbeCheckPort(probe appsv1alpha1.PodContainerProbe, pod *corev1.Pod) (appsv1alpha1.PodContainerProbe, error) {
+	probeNew := probe.DeepCopy()
+	if probe.Probe.TCPSocket.Port.Type == intstr.Int {
+		return *probeNew, nil
+	}
+	container := util.GetPodContainerByName(probe.ContainerName, pod)
+	if container == nil {
+		return *probeNew, fmt.Errorf("Failed to get container by name: %v in pod: %v/%v", probe.ContainerName, pod.Namespace, pod.Name)
+	}
+	portInt, err := util.ExtractPort(probe.Probe.TCPSocket.Port, *container)
+	if err != nil {
+		return *probeNew, fmt.Errorf("Failed to extract port for container: %v in pod: %v/%v", container.Name, pod.Namespace, pod.Name)
+	}
+	probeNew.Probe.TCPSocket.Port = intstr.FromInt(portInt)
+	return *probeNew, nil
+}
+
 func setPodContainerProbes(podProbe *appsv1alpha1.PodProbe, probe appsv1alpha1.PodContainerProbe, ppmName string) {
 	newProbe := appsv1alpha1.ContainerProbe{
 		Name:          fmt.Sprintf("%s#%s", ppmName, probe.Name),