Upgrade masters last when upgrading ES clusters

pkg/controller/elasticsearch/driver/upgrade.go

@@ -207,6 +207,12 @@ func isVersionUpgrade(es esv1.Elasticsearch) (bool, error) {
 	if err != nil {
 		return false, err
 	}
+
+	// If status version is empty, this is a new cluster, not an upgrade
+	if es.Status.Version == "" {
+		return false, nil
+	}
+
 	statusVersion, err := version.Parse(es.Status.Version)
 	if err != nil {
 		return false, err

pkg/controller/elasticsearch/driver/upscale.go

@@ -15,6 +15,8 @@ import (
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/expectations"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/metadata"
 	sset "github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/statefulset"
+	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/version"
+	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/elasticsearch/label"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/elasticsearch/nodespec"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/elasticsearch/reconcile"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/elasticsearch/settings"
@@ -66,33 +68,72 @@ func HandleUpscaleAndSpecChanges(
 	if err != nil {
 		return results, fmt.Errorf("adjust resources: %w", err)
 	}
-	// reconcile all resources
+
+	// Check if this is a version upgrade
+	isVersionUpgrade, err := isVersionUpgrade(ctx.es)
+	if err != nil {
+		return results, fmt.Errorf("while checking for version upgrade: %w", err)
+	}
+
+	// If this is not a version upgrade, process all resources normally and return
+	if !isVersionUpgrade {
+		actualStatefulSets, requeue, err := reconcileResources(ctx, actualStatefulSets, adjusted)
+		if err != nil {
+			return results, fmt.Errorf("while reconciling resources: %w", err)
+		}
+		results.Requeue = requeue
+		results.ActualStatefulSets = actualStatefulSets
+		return results, nil
+	}
+
+	// Version upgrade: separate master and non-master StatefulSets
+	var masterResources, nonMasterResources []nodespec.Resources
 	for _, res := range adjusted {
-		res := res
-		if err := settings.ReconcileConfig(ctx.parentCtx, ctx.k8sClient, ctx.es, res.StatefulSet.Name, res.Config, ctx.meta); err != nil {
-			return results, fmt.Errorf("reconcile config: %w", err)
+		if label.IsMasterNodeSet(res.StatefulSet) {
+			masterResources = append(masterResources, res)
+		} else {
+			nonMasterResources = append(nonMasterResources, res)
 		}
-		if _, err := common.ReconcileService(ctx.parentCtx, ctx.k8sClient, &res.HeadlessService, &ctx.es); err != nil {
-			return results, fmt.Errorf("reconcile service: %w", err)
+	}
+
+	// First, reconcile all non-master resources
+	actualStatefulSets, requeue, err := reconcileResources(ctx, actualStatefulSets, nonMasterResources)
+	if err != nil {
+		return results, fmt.Errorf("while reconciling non-master resources: %w", err)
+	}
+	if requeue {
+		results.Requeue = true
+		results.ActualStatefulSets = actualStatefulSets
+		return results, nil
+	}
+
+	targetVersion, err := version.Parse(ctx.es.Spec.Version)
+	if err != nil {
+		return results, fmt.Errorf("while parsing Elasticsearch upgrade target version: %w", err)
+	}
+
+	// Check if all non-master StatefulSets have completed their upgrades before proceeding with master StatefulSets
+	if len(masterResources) > 0 {
+		allNonMastersUpgraded, err := areAllNonMasterStatefulSetsUpgraded(ctx.k8sClient, actualStatefulSets, targetVersion)
+		if err != nil {
+			return results, fmt.Errorf("while checking non-master upgrade status: %w", err)
 		}
-		if actualSset, exists := actualStatefulSets.GetByName(res.StatefulSet.Name); exists {
-			recreateSset, err := handleVolumeExpansion(ctx.parentCtx, ctx.k8sClient, ctx.es, res.StatefulSet, actualSset, ctx.validateStorageClass)
-			if err != nil {
-				return results, fmt.Errorf("handle volume expansion: %w", err)
-			}
-			if recreateSset {
-				// The StatefulSet is scheduled for recreation: let's requeue before attempting any further spec change.
-				results.Requeue = true
-				continue
-			}
+
+		if !allNonMastersUpgraded {
+			// Non-master StatefulSets are still upgrading, skipping master StatefulSets temporarily.
+			// This will cause a requeue, and master StatefulSets will attempt to be processed in the next reconciliation
+			results.ActualStatefulSets = actualStatefulSets
+			results.Requeue = true
+			return results, nil
 		}
-		reconciled, err := es_sset.ReconcileStatefulSet(ctx.parentCtx, ctx.k8sClient, ctx.es, res.StatefulSet, ctx.expectations)
+
+		// All non-master StatefulSets are upgraded, now process master StatefulSets
+		actualStatefulSets, results.Requeue, err = reconcileResources(ctx, actualStatefulSets, masterResources)
 		if err != nil {
-			return results, fmt.Errorf("reconcile StatefulSet: %w", err)
+			return results, fmt.Errorf("while reconciling master resources: %w", err)
 		}
-		// update actual with the reconciled ones for next steps to work with up-to-date information
-		actualStatefulSets = actualStatefulSets.WithStatefulSet(reconciled)
 	}
+
 	results.ActualStatefulSets = actualStatefulSets
 	return results, nil
 }
@@ -166,3 +207,88 @@ func adjustStatefulSetReplicas(
 
 	return expected, nil
 }
+
+// reconcileResources handles the common StatefulSet reconciliation logic
+// It returns:
+// - the updated StatefulSets
+// - whether a requeue is needed
+// - any errors that occurred
+func reconcileResources(
+	ctx upscaleCtx,
+	actualStatefulSets es_sset.StatefulSetList,
+	resources []nodespec.Resources,
+) (es_sset.StatefulSetList, bool, error) {
+	requeue := false
+	for _, res := range resources {
+		res := res
+		if err := settings.ReconcileConfig(ctx.parentCtx, ctx.k8sClient, ctx.es, res.StatefulSet.Name, res.Config, ctx.meta); err != nil {
+			return actualStatefulSets, false, fmt.Errorf("reconcile config: %w", err)
+		}
+		if _, err := common.ReconcileService(ctx.parentCtx, ctx.k8sClient, &res.HeadlessService, &ctx.es); err != nil {
+			return actualStatefulSets, false, fmt.Errorf("reconcile service: %w", err)
+		}
+		if actualSset, exists := actualStatefulSets.GetByName(res.StatefulSet.Name); exists {
+			recreateSset, err := handleVolumeExpansion(ctx.parentCtx, ctx.k8sClient, ctx.es, res.StatefulSet, actualSset, ctx.validateStorageClass)
+			if err != nil {
+				return actualStatefulSets, false, fmt.Errorf("handle volume expansion: %w", err)
+			}
+			if recreateSset {
+				// The StatefulSet is scheduled for recreation: let's requeue before attempting any further spec change.
+				requeue = true
+				continue
+			}
+		}
+		reconciled, err := es_sset.ReconcileStatefulSet(ctx.parentCtx, ctx.k8sClient, ctx.es, res.StatefulSet, ctx.expectations)
+		if err != nil {
+			return actualStatefulSets, false, fmt.Errorf("reconcile StatefulSet: %w", err)
+		}
+		// update actual with the reconciled ones for next steps to work with up-to-date information
+		actualStatefulSets = actualStatefulSets.WithStatefulSet(reconciled)
+	}
+	return actualStatefulSets, requeue, nil
+}
+
+// areAllNonMasterStatefulSetsUpgraded checks if all non-master StatefulSets have completed their upgrades
+func areAllNonMasterStatefulSetsUpgraded(
+	client k8s.Client,
+	actualStatefulSets es_sset.StatefulSetList,
+	targetVersion version.Version,
+) (bool, error) {
+	for _, statefulSet := range actualStatefulSets {
+		// Skip master StatefulSets
+		if label.IsMasterNodeSet(statefulSet) {
+			continue
+		}
+
+		// If the StatefulSet is not at the target version, it is not upgraded
+		// so don't even bother looking at the state/status of the StatefulSet.
+		actualVersion, err := es_sset.GetESVersion(statefulSet)
+		if err != nil {
+			return false, err
+		}
+		if actualVersion.LT(targetVersion) {
+			return false, nil
+		}
+
+		// Check if this StatefulSet has pending updates
+		if statefulSet.Status.UpdatedReplicas != statefulSet.Status.Replicas {
+			return false, nil
+		}
+
+		// Check if there are any pods that need to be upgraded
+		pods, err := es_sset.GetActualPodsForStatefulSet(client, k8s.ExtractNamespacedName(&statefulSet))
+		if err != nil {
+			return false, err
+		}
+
+		for _, pod := range pods {
+			// Check if pod revision matches StatefulSet update revision
+			if statefulSet.Status.UpdateRevision != "" && sset.PodRevision(pod) != statefulSet.Status.UpdateRevision {
+				// This pod still needs to be upgraded
+				return false, nil
+			}
+		}
+	}
+
+	return true, nil
+}