Protect Pods from disruptions during upgrades

pkg/cluster/majorversionupgrade.go

@@ -106,6 +106,34 @@ func (c *Cluster) removeFailuresAnnotation() error {
 	return nil
 }
 
+func (c *Cluster) criticalOperationLabel(pods []v1.Pod, remove bool) error {
+	var action string
+	var metadataReq map[string]map[string]map[string]*string
+
+	if remove {
+		action = "remove"
+		metadataReq = map[string]map[string]map[string]*string{"metadata": {"labels": {"critical-operaton": nil}}}
+
+	} else {
+		action = "assign"
+		val := "true"
+		metadataReq = map[string]map[string]map[string]*string{"metadata": {"labels": {"critical-operaton": &val}}}
+	}
+
+	patchReq, err := json.Marshal(metadataReq)
+	if err != nil {
+		return fmt.Errorf("could not marshal ObjectMeta to %s critical operation label: %v", action, err)
+	}
+	for _, pod := range pods {
+		_, err = c.KubeClient.Pods(c.Namespace).Patch(context.TODO(), pod.Name, types.StrategicMergePatchType, patchReq, metav1.PatchOptions{})
+		if err != nil {
+			c.logger.Errorf("failed to %s critical operation label for pod %s: %v", action, pod.Name, err)
+			return err
+		}
+	}
+	return nil
+}
+
 /*
 Execute upgrade when mode is set to manual or full or when the owning team is allowed for upgrade (and mode is "off").
 
@@ -219,6 +247,16 @@ func (c *Cluster) majorVersionUpgrade() error {
 	if allRunning && masterPod != nil {
 		c.logger.Infof("healthy cluster ready to upgrade, current: %d desired: %d", c.currentMajorVersion, desiredVersion)
 		if c.currentMajorVersion < desiredVersion {
+			defer func() error {
+				if err = c.criticalOperationLabel(pods, true); err != nil {
+					return fmt.Errorf("failed to remove critical-operation label: %s", err)
+				}
+				return nil
+			}()
+			if err = c.criticalOperationLabel(pods, false); err != nil {
+				return fmt.Errorf("failed to assign critical-operation label: %s", err)
+			}
+
 			podName := &spec.NamespacedName{Namespace: masterPod.Namespace, Name: masterPod.Name}
 			c.logger.Infof("triggering major version upgrade on pod %s of %d pods", masterPod.Name, numberOfPods)
 			c.eventRecorder.Eventf(c.GetReference(), v1.EventTypeNormal, "Major Version Upgrade", "starting major version upgrade on pod %s of %d pods", masterPod.Name, numberOfPods)