Skip to content

Commit b9a4ef2

Browse files
Allow overriding readinessProbe and LivenessProbe (#1776)
There can only be 1 probe and without this change, specifying an exec readinessProbe in the spec.override would add it on top of the default TCP probe. As for livenessProbe, since there's currently no default, everything works without this change. However, we could add the default livenessProbe in the future and allowing this override doesn't break anything right now, so it's just future-proofing. Co-authored-by: joey <[email protected]>
1 parent e23610b commit b9a4ef2

File tree

2 files changed

+87
-0
lines changed

2 files changed

+87
-0
lines changed

internal/resource/statefulset.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,15 @@ func patchPodSpec(podSpec, podSpecOverride *corev1.PodSpec) (corev1.PodSpec, err
286286
sortVolumeMounts(patchedPodSpec.Containers[0].VolumeMounts)
287287
}
288288

289+
// allow overriding the liveness and readiness probes
290+
// note: as of 2024, we don't set a default LivenessProbe
291+
if rmqContainer.LivenessProbe != nil {
292+
patchedPodSpec.Containers[0].LivenessProbe = rmqContainer.LivenessProbe
293+
}
294+
if rmqContainer.ReadinessProbe != nil {
295+
patchedPodSpec.Containers[0].ReadinessProbe = rmqContainer.ReadinessProbe
296+
}
297+
289298
// A user may wish to override the controller-set securityContext for the RabbitMQ & init containers so that the
290299
// container runtime can override them. If the securityContext has been set to an empty struct, `strategicpatch.StrategicMergePatch`
291300
// won't pick this up, so manually override it here.

internal/resource/statefulset_test.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2126,6 +2126,84 @@ default_pass = {{ .Data.data.password }}
21262126

21272127
})
21282128

2129+
It("can replace the default readinessProbe", func() {
2130+
instance.Spec.Override.StatefulSet = &rabbitmqv1beta1.StatefulSet{
2131+
Spec: &rabbitmqv1beta1.StatefulSetSpec{
2132+
Template: &rabbitmqv1beta1.PodTemplateSpec{
2133+
Spec: &corev1.PodSpec{
2134+
Containers: []corev1.Container{
2135+
{
2136+
Name: "rabbitmq",
2137+
ReadinessProbe: &corev1.Probe{
2138+
ProbeHandler: corev1.ProbeHandler{
2139+
Exec: &corev1.ExecAction{
2140+
Command: []string{"custom-readiness-probe", "arg1"},
2141+
},
2142+
},
2143+
},
2144+
},
2145+
},
2146+
},
2147+
},
2148+
},
2149+
}
2150+
2151+
builder = &resource.RabbitmqResourceBuilder{
2152+
Instance: &instance,
2153+
Scheme: scheme,
2154+
}
2155+
stsBuilder := builder.StatefulSet()
2156+
Expect(stsBuilder.Update(statefulSet)).To(Succeed())
2157+
2158+
Expect(statefulSet.Spec.Template.Spec.Containers[0].ReadinessProbe.ProbeHandler).To(Equal(
2159+
corev1.ProbeHandler{
2160+
Exec: &corev1.ExecAction{
2161+
Command: []string{"custom-readiness-probe", "arg1"},
2162+
},
2163+
},
2164+
))
2165+
})
2166+
2167+
It("can add/replace a LivenessProbe", func() {
2168+
// note: we currently don't have a default LivenessProbe
2169+
instance.Spec.Override.StatefulSet = &rabbitmqv1beta1.StatefulSet{
2170+
Spec: &rabbitmqv1beta1.StatefulSetSpec{
2171+
Template: &rabbitmqv1beta1.PodTemplateSpec{
2172+
Spec: &corev1.PodSpec{
2173+
Containers: []corev1.Container{
2174+
{
2175+
Name: "rabbitmq",
2176+
LivenessProbe: &corev1.Probe{
2177+
ProbeHandler: corev1.ProbeHandler{
2178+
Exec: &corev1.ExecAction{
2179+
Command: []string{"custom-liveness-probe", "arg1"},
2180+
},
2181+
},
2182+
},
2183+
},
2184+
},
2185+
},
2186+
},
2187+
},
2188+
}
2189+
2190+
builder = &resource.RabbitmqResourceBuilder{
2191+
Instance: &instance,
2192+
Scheme: scheme,
2193+
}
2194+
stsBuilder := builder.StatefulSet()
2195+
Expect(stsBuilder.Update(statefulSet)).To(Succeed())
2196+
2197+
Expect(statefulSet.Spec.Template.Spec.Containers[0].LivenessProbe.ProbeHandler).To(Equal(
2198+
corev1.ProbeHandler{
2199+
Exec: &corev1.ExecAction{
2200+
Command: []string{"custom-liveness-probe", "arg1"},
2201+
},
2202+
},
2203+
))
2204+
2205+
})
2206+
21292207
Context("TopologySpreadConstraints composition", func() {
21302208
BeforeEach(func() {
21312209
instance.Spec.Override.StatefulSet = &rabbitmqv1beta1.StatefulSet{

0 commit comments

Comments
 (0)