From b5791a1edd6cc47b415c15ea7130b2a28afa87f5 Mon Sep 17 00:00:00 2001
From: "kevin.qiao" <kevin.qiao@vesoft.com>
Date: Tue, 31 Aug 2021 19:43:23 +0800
Subject: [PATCH] update image limit

---
 pkg/controller/component/graphd_cluster.go      | 10 ++++++++++
 pkg/controller/component/metad_cluster.go       | 10 ++++++++++
 pkg/controller/component/storaged_cluster.go    | 10 ++++++++++
 .../nebulacluster/nebula_cluster_control.go     | 17 +++++++++++++++++
 4 files changed, 47 insertions(+)

diff --git a/pkg/controller/component/graphd_cluster.go b/pkg/controller/component/graphd_cluster.go
index b8650843..809dc9a7 100644
--- a/pkg/controller/component/graphd_cluster.go
+++ b/pkg/controller/component/graphd_cluster.go
@@ -93,6 +93,16 @@ func (c *graphdCluster) syncGraphdWorkload(nc *v1alpha1.NebulaCluster) error {
 		return err
 	}
 
+	if nc.Status.Graphd.Version != "" {
+		oldImage := nc.Spec.Graphd.Image + ":" + nc.Status.Graphd.Version
+		if err := extender.SetContainerImage(
+			newWorkload,
+			nc.GraphdComponent().Type().String(),
+			oldImage); err != nil {
+			return err
+		}
+	}
+
 	if err := extender.SetTemplateAnnotations(
 		newWorkload,
 		map[string]string{annotation.AnnPodConfigMapHash: cmHash}); err != nil {
diff --git a/pkg/controller/component/metad_cluster.go b/pkg/controller/component/metad_cluster.go
index a2b27378..ce2dfba3 100644
--- a/pkg/controller/component/metad_cluster.go
+++ b/pkg/controller/component/metad_cluster.go
@@ -99,6 +99,16 @@ func (c *metadCluster) syncMetadWorkload(nc *v1alpha1.NebulaCluster) error {
 		return err
 	}
 
+	if nc.Status.Metad.Version != "" {
+		oldImage := nc.Spec.Metad.Image + ":" + nc.Status.Metad.Version
+		if err := extender.SetContainerImage(
+			newWorkload,
+			nc.MetadComponent().Type().String(),
+			oldImage); err != nil {
+			return err
+		}
+	}
+
 	if err := extender.SetTemplateAnnotations(
 		newWorkload,
 		map[string]string{annotation.AnnPodConfigMapHash: cmHash}); err != nil {
diff --git a/pkg/controller/component/storaged_cluster.go b/pkg/controller/component/storaged_cluster.go
index 2e60e778..ce38ec76 100644
--- a/pkg/controller/component/storaged_cluster.go
+++ b/pkg/controller/component/storaged_cluster.go
@@ -102,6 +102,16 @@ func (c *storagedCluster) syncStoragedWorkload(nc *v1alpha1.NebulaCluster) error
 		return err
 	}
 
+	if nc.Status.Storaged.Version != "" {
+		oldImage := nc.Spec.Storaged.Image + ":" + nc.Status.Storaged.Version
+		if err := extender.SetContainerImage(
+			newWorkload,
+			nc.StoragedComponent().Type().String(),
+			oldImage); err != nil {
+			return err
+		}
+	}
+
 	if err := extender.SetTemplateAnnotations(
 		newWorkload,
 		map[string]string{annotation.AnnPodConfigMapHash: cmHash}); err != nil {
diff --git a/pkg/controller/nebulacluster/nebula_cluster_control.go b/pkg/controller/nebulacluster/nebula_cluster_control.go
index ef6ba77c..81d09c4c 100644
--- a/pkg/controller/nebulacluster/nebula_cluster_control.go
+++ b/pkg/controller/nebulacluster/nebula_cluster_control.go
@@ -65,6 +65,11 @@ func (c *defaultNebulaClusterControl) UpdateNebulaCluster(nc *v1alpha1.NebulaClu
 	var errs []error
 	oldStatus := nc.Status.DeepCopy()
 
+	c.resetImage(nc)
+	if err := c.nebulaClient.UpdateNebulaCluster(nc.DeepCopy()); err != nil {
+		errs = append(errs, err)
+	}
+
 	if err := c.updateNebulaCluster(nc); err != nil {
 		errs = append(errs, err)
 	}
@@ -112,6 +117,18 @@ func (c *defaultNebulaClusterControl) updateNebulaCluster(nc *v1alpha1.NebulaClu
 	return nil
 }
 
+func (c *defaultNebulaClusterControl) resetImage(nc *v1alpha1.NebulaCluster) {
+	if nc.Status.Graphd.Version != "" {
+		nc.Spec.Graphd.Version = nc.Status.Graphd.Version
+	}
+	if nc.Status.Metad.Version != "" {
+		nc.Spec.Metad.Version = nc.Status.Metad.Version
+	}
+	if nc.Status.Storaged.Version != "" {
+		nc.Spec.Storaged.Version = nc.Status.Storaged.Version
+	}
+}
+
 type FakeClusterControl struct {
 	err error
 }