Update initContainer image when deployment rollout (#87)

Signed-off-by: shaoyue.chen <[email protected]>
zilliztech · Mar 8, 2024 · baa4a25 · baa4a25
1 parent 370a9ae
commit baa4a25
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 12 deletions.
diff --git a/apis/milvus.io/v1beta1/components_types.go b/apis/milvus.io/v1beta1/components_types.go
@@ -134,6 +134,7 @@ type MilvusComponents struct {
 	ToolImage string `json:"toolImage,omitempty"`
 
 	// UpdateToolImage when milvus-operator upgraded, whether milvus should restart to update the tool image, too
+	// otherwise, the tool image will be updated when milvus deploy's podTemplate changed
 	// +kubebuilder:validation:Optional
 	UpdateToolImage bool `json:"updateToolImage,omitempty"`
 

diff --git a/pkg/controllers/deployment_updater.go b/pkg/controllers/deployment_updater.go
@@ -105,14 +105,8 @@ func updatePodMeta(template *corev1.PodTemplateSpec, appLabels map[string]string
 func updateInitContainers(template *corev1.PodTemplateSpec, updater deploymentUpdater) {
 	configContainerIdx := GetContainerIndex(template.Spec.InitContainers, configContainerName)
 	spec := updater.GetMilvus().Spec
-	if configContainerIdx < 0 {
-		var container = new(corev1.Container)
-		if len(template.Spec.InitContainers) < 1 {
-			template.Spec.InitContainers = []corev1.Container{}
-		}
-		template.Spec.InitContainers = append(template.Spec.InitContainers, *renderInitContainer(container, spec.Com.ToolImage))
-	} else if spec.Com.UpdateToolImage {
-		renderInitContainer(&template.Spec.InitContainers[configContainerIdx], spec.Com.ToolImage)
+	if configContainerIdx < 0 || spec.Com.UpdateToolImage {
+		updateConfigContainer(template, updater)
 	}
 
 	initContainers := updater.GetInitContainers()
@@ -128,6 +122,20 @@ func updateInitContainers(template *corev1.PodTemplateSpec, updater deploymentUp
 	}
 }
 
+func updateConfigContainer(template *corev1.PodTemplateSpec, updater deploymentUpdater) {
+	configContainerIdx := GetContainerIndex(template.Spec.InitContainers, configContainerName)
+	spec := updater.GetMilvus().Spec
+	if configContainerIdx < 0 {
+		var container = new(corev1.Container)
+		if len(template.Spec.InitContainers) < 1 {
+			template.Spec.InitContainers = []corev1.Container{}
+		}
+		template.Spec.InitContainers = append(template.Spec.InitContainers, *renderInitContainer(container, spec.Com.ToolImage))
+	} else {
+		renderInitContainer(&template.Spec.InitContainers[configContainerIdx], spec.Com.ToolImage)
+	}
+}
+
 func updateScheduleSpec(template *corev1.PodTemplateSpec, updater deploymentUpdater) {
 	mergedComSpec := updater.GetMergedComponentSpec()
 	if len(mergedComSpec.SchedulerName) > 0 {
@@ -228,14 +236,15 @@ func updateMilvusContainer(template *corev1.PodTemplateSpec, updater deploymentU
 }
 
 func updateSomeFieldsOnlyWhenRolling(template *corev1.PodTemplateSpec, updater deploymentUpdater) {
+	// when perform rolling update
+	// we add some other perfered updates
+	updateConfigContainer(template, updater)
 	componentName := updater.GetComponentName()
 	containerIdx := GetContainerIndex(template.Spec.Containers, updater.GetComponentName())
 	container := &template.Spec.Containers[containerIdx]
 	if componentName == ProxyName || componentName == StandaloneName {
 		template.Labels[v1beta1.ServiceLabel] = v1beta1.TrueStr
 	}
-	// will perform rolling update
-	// we add some other perfered updates
 	container.StartupProbe = GetStartupProbe()
 	container.LivenessProbe = GetLivenessProbe()
 	container.ReadinessProbe = GetReadinessProbe()
@@ -256,11 +265,13 @@ func updateSomeFieldsOnlyWhenRolling(template *corev1.PodTemplateSpec, updater d
 			},
 		}
 	}
+	// oneMonthSeconds we set both podtemplate.spec.terminationGracePeriodSeconds &
+	// deployment.spec.progressDeadlineSeconds to one month, to avoid kill -9 on pod automatically.
+	// so that when pod stuck on rolling, the service will still be available.
+	// We'll have enough time to find the root cause and handle it gracefully.
 	template.Spec.TerminationGracePeriodSeconds = int64Ptr(int64(oneMonthSeconds))
 }
 
-// oneMonthSeconds we set both podtemplate.spec.terminationGracePeriodSeconds &
-// deployment.spec.progressDeadlineSeconds to one month, to avoid kill -9 on pod automatically
 const oneMonthSeconds = 24 * 30 * int(time.Hour/time.Second)
 
 func updateSidecars(template *corev1.PodTemplateSpec, updater deploymentUpdater) {

diff --git a/pkg/controllers/deployment_updater_test.go b/pkg/controllers/deployment_updater_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/milvus-io/milvus-operator/pkg/util"
 	"github.com/stretchr/testify/assert"
 	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
 )
 
 func TestMilvus_UpdateDeployment(t *testing.T) {
@@ -30,6 +31,75 @@ func TestMilvus_UpdateDeployment(t *testing.T) {
 		assert.Equal(t, []string{"/milvus/tools/run.sh", "milvus", "run", "mycomponent"}, deployment.Spec.Template.Spec.Containers[0].Args)
 	})
 
+	t.Run("with init container", func(t *testing.T) {
+		inst := env.Inst.DeepCopy()
+		inst.Spec.Com.Standalone.InitContainers = []v1beta1.Values{{}}
+		inst.Spec.GetServiceComponent().Commands = []string{"milvus", "run", "mycomponent"}
+		updater := newMilvusDeploymentUpdater(*inst, env.Reconciler.Scheme, MilvusStandalone)
+		deployment := &appsv1.Deployment{}
+		deployment.Name = "deploy"
+		deployment.Namespace = "ns"
+		err := updateDeployment(deployment, updater)
+		assert.NoError(t, err)
+		assert.Len(t, deployment.Spec.Template.Spec.InitContainers, 2)
+	})
+
+	globalCommonInfo.OperatorImageInfo = DefaultOperatorImageInfo
+	defer func() {
+		globalCommonInfo.OperatorImageInfo = ImageInfo{}
+	}()
+	t.Run("not update configContainer when podTemplate not updated", func(t *testing.T) {
+		inst := env.Inst.DeepCopy()
+		inst.Spec.GetServiceComponent().Commands = []string{"milvus", "run", "mycomponent"}
+		updater := newMilvusDeploymentUpdater(*inst, env.Reconciler.Scheme, MilvusStandalone)
+		deployment := &appsv1.Deployment{}
+		deployment.Name = "deploy"
+		deployment.Namespace = "ns"
+		err := updateDeployment(deployment, updater)
+		assert.NoError(t, err)
+		deployment.Spec.Template.Spec.InitContainers = []corev1.Container{
+			{
+				Name: configContainerName,
+			},
+		}
+		err = updateDeployment(deployment, updater)
+		assert.NoError(t, err)
+		assert.Empty(t, deployment.Spec.Template.Spec.InitContainers[0].Image)
+	})
+
+	t.Run("update configContainer when UpdateToolImage is true", func(t *testing.T) {
+		inst := env.Inst.DeepCopy()
+		inst.Spec.Com.UpdateToolImage = true
+		inst.Spec.GetServiceComponent().Commands = []string{"milvus", "run", "mycomponent"}
+		updater := newMilvusDeploymentUpdater(*inst, env.Reconciler.Scheme, MilvusStandalone)
+		deployment := &appsv1.Deployment{}
+		deployment.Name = "deploy"
+		deployment.Namespace = "ns"
+		err := updateDeployment(deployment, updater)
+		assert.NoError(t, err)
+		deployment.Spec.Template.Spec.InitContainers[0].Image = ""
+		err = updateDeployment(deployment, updater)
+		assert.NoError(t, err)
+		assert.Equal(t, DefaultOperatorImageInfo.Image, deployment.Spec.Template.Spec.InitContainers[0].Image)
+	})
+
+	t.Run("update configContainer when podTemplate updated", func(t *testing.T) {
+		inst := env.Inst.DeepCopy()
+		inst.Spec.GetServiceComponent().Commands = []string{"milvus", "run", "mycomponent"}
+		updater := newMilvusDeploymentUpdater(*inst, env.Reconciler.Scheme, MilvusStandalone)
+		deployment := &appsv1.Deployment{}
+		deployment.Name = "deploy"
+		deployment.Namespace = "ns"
+		deployment.Spec.Template.Spec.InitContainers = []corev1.Container{
+			{
+				Name: configContainerName,
+			},
+		}
+		err := updateDeployment(deployment, updater)
+		assert.NoError(t, err)
+		assert.Equal(t, DefaultOperatorImageInfo.Image, deployment.Spec.Template.Spec.InitContainers[0].Image)
+	})
+
 	t.Run("persistence disabled", func(t *testing.T) {
 		inst := env.Inst.DeepCopy()
 		updater := newMilvusDeploymentUpdater(*inst, env.Reconciler.Scheme, MilvusStandalone)