Merge pull request #110733 from psschwei/probe-grace-period-units

Add unit tests for grace period in killContainer func

pkg/kubelet/kuberuntime/kuberuntime_container.go

@@ -673,24 +673,7 @@ func (m *kubeGenericRuntimeManager) killContainer(pod *v1.Pod, containerID kubec
 	}
 
 	// From this point, pod and container must be non-nil.
-	gracePeriod := int64(minimumGracePeriodInSeconds)
+	gracePeriod := setTerminationGracePeriod(pod, containerSpec, containerName, containerID, reason)
-	switch {
-	case pod.DeletionGracePeriodSeconds != nil:
-		gracePeriod = *pod.DeletionGracePeriodSeconds
-	case pod.Spec.TerminationGracePeriodSeconds != nil:
-		gracePeriod = *pod.Spec.TerminationGracePeriodSeconds
-
-		switch reason {
-		case reasonStartupProbe:
-			if containerSpec.StartupProbe != nil && containerSpec.StartupProbe.TerminationGracePeriodSeconds != nil {
-				gracePeriod = *containerSpec.StartupProbe.TerminationGracePeriodSeconds
-			}
-		case reasonLivenessProbe:
-			if containerSpec.LivenessProbe != nil && containerSpec.LivenessProbe.TerminationGracePeriodSeconds != nil {
-				gracePeriod = *containerSpec.LivenessProbe.TerminationGracePeriodSeconds
-			}
-		}
-	}
 
 	if len(message) == 0 {
 		message = fmt.Sprintf("Stopping container %s", containerSpec.Name)
@@ -991,3 +974,35 @@ func (m *kubeGenericRuntimeManager) removeContainerLog(containerID string) error
 func (m *kubeGenericRuntimeManager) DeleteContainer(containerID kubecontainer.ContainerID) error {
 	return m.removeContainer(containerID.ID)
 }
+
+// setTerminationGracePeriod determines the grace period to use when killing a container
+func setTerminationGracePeriod(pod *v1.Pod, containerSpec *v1.Container, containerName string, containerID kubecontainer.ContainerID, reason containerKillReason) int64 {
+	gracePeriod := int64(minimumGracePeriodInSeconds)
+	switch {
+	case pod.DeletionGracePeriodSeconds != nil:
+		return *pod.DeletionGracePeriodSeconds
+	case pod.Spec.TerminationGracePeriodSeconds != nil:
+		switch reason {
+		case reasonStartupProbe:
+			if isProbeTerminationGracePeriodSecondsSet(pod, containerSpec, containerSpec.StartupProbe, containerName, containerID, "StartupProbe") {
+				return *containerSpec.StartupProbe.TerminationGracePeriodSeconds
+			}
+		case reasonLivenessProbe:
+			if isProbeTerminationGracePeriodSecondsSet(pod, containerSpec, containerSpec.LivenessProbe, containerName, containerID, "LivenessProbe") {
+				return *containerSpec.LivenessProbe.TerminationGracePeriodSeconds
+			}
+		}
+		return *pod.Spec.TerminationGracePeriodSeconds
+	}
+	return gracePeriod
+}
+
+func isProbeTerminationGracePeriodSecondsSet(pod *v1.Pod, containerSpec *v1.Container, probe *v1.Probe, containerName string, containerID kubecontainer.ContainerID, probeType string) bool {
+	if probe != nil && probe.TerminationGracePeriodSeconds != nil {
+		if *probe.TerminationGracePeriodSeconds > *pod.Spec.TerminationGracePeriodSeconds {
+			klog.V(4).InfoS("Using probe-level grace period that is greater than the pod-level grace period", "pod", klog.KObj(pod), "pod-uid", pod.UID, "containerName", containerName, "containerID", containerID.String(), "probe-type", probeType, "probe-grace-period", *probe.TerminationGracePeriodSeconds, "pod-grace-period", *pod.Spec.TerminationGracePeriodSeconds)
+		}
+		return true
+	}
+	return false
+}

pkg/kubelet/kuberuntime/kuberuntime_container_test.go

@@ -451,3 +451,167 @@ func TestRestartCountByLogDir(t *testing.T) {
 		assert.Equal(t, count, tc.restartCount, "count %v should equal restartCount %v", count, tc.restartCount)
 	}
 }
+
+func TestKillContainerGracePeriod(t *testing.T) {
+
+	shortGracePeriod := int64(10)
+	mediumGracePeriod := int64(30)
+	longGracePeriod := int64(60)
+
+	tests := []struct {
+		name                string
+		pod                 *v1.Pod
+		reason              containerKillReason
+		expectedGracePeriod int64
+	}{
+		{
+			name: "default termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{Containers: []v1.Container{{Name: "foo"}}},
+			},
+			reason:              reasonUnknown,
+			expectedGracePeriod: int64(2),
+		},
+		{
+			name: "use pod termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers:                    []v1.Container{{Name: "foo"}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonUnknown,
+			expectedGracePeriod: longGracePeriod,
+		},
+		{
+			name: "liveness probe overrides pod termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name: "foo", LivenessProbe: &v1.Probe{TerminationGracePeriodSeconds: &shortGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonLivenessProbe,
+			expectedGracePeriod: shortGracePeriod,
+		},
+		{
+			name: "startup probe overrides pod termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name: "foo", StartupProbe: &v1.Probe{TerminationGracePeriodSeconds: &shortGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonStartupProbe,
+			expectedGracePeriod: shortGracePeriod,
+		},
+		{
+			name: "startup probe overrides pod termination grace period, probe period > pod period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name: "foo", StartupProbe: &v1.Probe{TerminationGracePeriodSeconds: &longGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &shortGracePeriod,
+				},
+			},
+			reason:              reasonStartupProbe,
+			expectedGracePeriod: longGracePeriod,
+		},
+		{
+			name: "liveness probe overrides pod termination grace period, probe period > pod period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name: "foo", LivenessProbe: &v1.Probe{TerminationGracePeriodSeconds: &longGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &shortGracePeriod,
+				},
+			},
+			reason:              reasonLivenessProbe,
+			expectedGracePeriod: longGracePeriod,
+		},
+		{
+			name: "non-liveness probe failure, use pod termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name: "foo", LivenessProbe: &v1.Probe{TerminationGracePeriodSeconds: &shortGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonUnknown,
+			expectedGracePeriod: longGracePeriod,
+		},
+		{
+			name: "non-startup probe failure, use pod termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name: "foo", StartupProbe: &v1.Probe{TerminationGracePeriodSeconds: &shortGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonUnknown,
+			expectedGracePeriod: longGracePeriod,
+		},
+		{
+			name: "all three grace periods set, use pod termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name:          "foo",
+						StartupProbe:  &v1.Probe{TerminationGracePeriodSeconds: &shortGracePeriod},
+						LivenessProbe: &v1.Probe{TerminationGracePeriodSeconds: &mediumGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonUnknown,
+			expectedGracePeriod: longGracePeriod,
+		},
+		{
+			name: "all three grace periods set, use startup termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name:          "foo",
+						StartupProbe:  &v1.Probe{TerminationGracePeriodSeconds: &shortGracePeriod},
+						LivenessProbe: &v1.Probe{TerminationGracePeriodSeconds: &mediumGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonStartupProbe,
+			expectedGracePeriod: shortGracePeriod,
+		},
+		{
+			name: "all three grace periods set, use liveness termination grace period",
+			pod: &v1.Pod{
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{{
+						Name:          "foo",
+						StartupProbe:  &v1.Probe{TerminationGracePeriodSeconds: &shortGracePeriod},
+						LivenessProbe: &v1.Probe{TerminationGracePeriodSeconds: &mediumGracePeriod},
+					}},
+					TerminationGracePeriodSeconds: &longGracePeriod,
+				},
+			},
+			reason:              reasonLivenessProbe,
+			expectedGracePeriod: mediumGracePeriod,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			actualGracePeriod := setTerminationGracePeriod(test.pod, &test.pod.Spec.Containers[0], "", kubecontainer.ContainerID{}, test.reason)
+			require.Equal(t, test.expectedGracePeriod, actualGracePeriod)
+		})
+	}
+}