Kubebuilder Operator for StatefulSet Zero-Downtime Updates

// api/v1alpha1/clustereddatabase_types.go

package v1alpha1

import (
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// ClusteredDatabaseSpec defines the desired state of ClusteredDatabase
type ClusteredDatabaseSpec struct {
	// +kubebuilder:validation:Minimum=1
	// Number of desired pods. Must be at least 1.
	Replicas *int32 `json:"replicas"`

	// Image is the container image to run for the database.
	Image string `json:"image"`

	// Resources defines the CPU and memory requests and limits for the database container.
	Resources corev1.ResourceRequirements `json:"resources,omitempty"`

	// Storage defines the persistent storage configuration for the database.
	Storage StorageSpec `json:"storage"`
}

// StorageSpec defines the storage configuration for each pod in the cluster.
type StorageSpec struct {
	// StorageClassName is the name of the StorageClass to use for the PersistentVolumeClaims.
	StorageClassName string `json:"storageClassName"`

	// Size is the size of the storage volume, e.g., "10Gi".
	Size string `json:"size"`
}

// api/v1alpha1/clustereddatabase_types.go (continued)

// ClusteredDatabaseStatus defines the observed state of ClusteredDatabase
type ClusteredDatabaseStatus struct {
	// ReadyReplicas is the number of pods in the cluster that are in a Ready state.
	ReadyReplicas int32 `json:"readyReplicas"`

	// CurrentPrimary identifies the pod that is currently acting as the primary.
	CurrentPrimary string `json:"currentPrimary,omitempty"`

	// Conditions represent the latest available observations of the ClusteredDatabase's state.
	// +optional
	Conditions []metav1.Condition `json:"conditions,omitempty"`
}

// These are the valid conditions for a ClusteredDatabase.
const (
    // AvailableCondition indicates that the database cluster is fully operational and available.
    AvailableCondition string = "Available"
    // UpdatingCondition indicates that the cluster is undergoing an update.
    UpdatingCondition string = "Updating"
    // DegradedCondition indicates that the cluster is not fully operational.
    DegradedCondition string = "Degraded"
)

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// +kubebuilder:printcolumn:name="Replicas",type="integer",JSONPath=".spec.replicas"
// +kubebuilder:printcolumn:name="Ready",type="integer",JSONPath=".status.readyReplicas"
// +kubebuilder:printcolumn:name="Primary",type="string",JSONPath=".status.currentPrimary"
// +kubebuilder:printcolumn:name="Status",type="string",JSONPath=".status.conditions[?(@.type==\"Available\")].status"
// +kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"

// ClusteredDatabase is the Schema for the clustereddatabases API
type ClusteredDatabase struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	Spec   ClusteredDatabaseSpec   `json:"spec,omitempty"`
	Status ClusteredDatabaseStatus `json:"status,omitempty"`
}

//+kubebuilder:object:root=true

// ClusteredDatabaseList contains a list of ClusteredDatabase
type ClusteredDatabaseList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []ClusteredDatabase `json:"items"`
}

func init() {
	SchemeBuilder.Register(&ClusteredDatabase{}, &ClusteredDatabaseList{})
}

// controllers/clustereddatabase_controller.go

import (
	// ... other imports
	"time"

	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	dbv1alpha1 "github.com/your-repo/clustereddb-operator/api/v1alpha1"
)

const clusteredDBFinalizer = "db.example.com/finalizer"

func (r *ClusteredDatabaseReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := r.Log.WithValues("clustereddatabase", req.NamespacedName)

	// Fetch the ClusteredDatabase instance
	db := &dbv1alpha1.ClusteredDatabase{}
	err := r.Get(ctx, req.NamespacedName, db)
	if err != nil {
		if errors.IsNotFound(err) {
			log.Info("ClusteredDatabase resource not found. Ignoring since object must be deleted")
			return ctrl.Result{}, nil
		}
		log.Error(err, "Failed to get ClusteredDatabase")
		return ctrl.Result{}, err
	}

	// Check if the instance is being deleted
	isMarkedForDeletion := db.GetDeletionTimestamp() != nil
	if isMarkedForDeletion {
		if controllerutil.ContainsFinalizer(db, clusteredDBFinalizer) {
			// Run finalization logic
			if err := r.finalizeClusteredDB(ctx, log, db); err != nil {
				return ctrl.Result{}, err
			}

			// Remove finalizer. Once all finalizers are removed, the object will be deleted.
			controllerutil.RemoveFinalizer(db, clusteredDBFinalizer)
			err := r.Update(ctx, db)
			if err != nil {
				return ctrl.Result{}, err
			}
		}
		return ctrl.Result{}, nil
	}

	// Add finalizer for this CR if it doesn't exist
	if !controllerutil.ContainsFinalizer(db, clusteredDBFinalizer) {
		controllerutil.AddFinalizer(db, clusteredDBFinalizer)
		err = r.Update(ctx, db)
		if err != nil {
			return ctrl.Result{}, err
		}
	}

	// ... main reconciliation logic follows
    return r.reconcileCluster(ctx, log, db)
}

func (r *ClusteredDatabaseReconciler) finalizeClusteredDB(ctx context.Context, log logr.Logger, db *dbv1alpha1.ClusteredDatabase) error {
	log.Info("Finalizing ClusteredDatabase")
	// In a real operator, you might want to perform a graceful shutdown of the database,
	// backup data, or other cleanup tasks.
	// For this example, we'll just log.
	log.Info("ClusteredDatabase finalizer logic executed.")
	return nil
}

// controllers/clustereddatabase_controller.go

// A portion of the main reconcileCluster function
func (r *ClusteredDatabaseReconciler) reconcileCluster(ctx context.Context, log logr.Logger, db *dbv1alpha1.ClusteredDatabase) (ctrl.Result, error) {
	// Reconcile Headless Service
	headlessSvc := &corev1.Service{}
	err := r.Get(ctx, types.NamespacedName{Name: db.Name + "-headless", Namespace: db.Namespace}, headlessSvc)
	if err != nil && errors.IsNotFound(err) {
		// Define and create a new headless service
		hls := r.serviceForClusteredDB(db)
		log.Info("Creating a new Headless Service", "Service.Namespace", hls.Namespace, "Service.Name", hls.Name)
		if err := r.Create(ctx, hls); err != nil {
			log.Error(err, "Failed to create new Headless Service")
			return ctrl.Result{}, err
		}
	} else if err != nil {
		log.Error(err, "Failed to get Headless Service")
		return ctrl.Result{}, err
	}

	// Reconcile StatefulSet
	sts := &appsv1.StatefulSet{}
	err = r.Get(ctx, types.NamespacedName{Name: db.Name, Namespace: db.Namespace}, sts)
	if err != nil && errors.IsNotFound(err) {
		// Define and create a new StatefulSet
		newSts := r.statefulSetForClusteredDB(db)
		log.Info("Creating a new StatefulSet", "StatefulSet.Namespace", newSts.Namespace, "StatefulSet.Name", newSts.Name)
		if err = r.Create(ctx, newSts); err != nil {
			log.Error(err, "Failed to create new StatefulSet")
			return ctrl.Result{}, err
		}
		// StatefulSet created, requeue to check status after a short delay
		return ctrl.Result{RequeueAfter: time.Second * 5}, nil
	}

	// ... update logic follows

	return ctrl.Result{}, nil
}

// Helper function to define the StatefulSet
func (r *ClusteredDatabaseReconciler) statefulSetForClusteredDB(db *dbv1alpha1.ClusteredDatabase) *appsv1.StatefulSet {
	labels := map[string]string{"app": db.Name}
	replicas := db.Spec.Replicas

	// CRITICAL: Set the update strategy to OnDelete
	updateStrategy := appsv1.StatefulSetUpdateStrategy{
		Type: appsv1.OnDeleteStatefulSetStrategyType,
	}

	sts := &appsv1.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:      db.Name,
			Namespace: db.Namespace,
		},
		Spec: appsv1.StatefulSetSpec{
			ServiceName: db.Name + "-headless",
			Replicas:    replicas,
			Selector:    &metav1.LabelSelector{MatchLabels: labels},
			UpdateStrategy: updateStrategy,
			Template: corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{Labels: labels},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{{
						Image:     db.Spec.Image,
						Name:      "database",
						Ports:     []corev1.ContainerPort{{ContainerPort: 5432, Name: "db"}},
						Resources: db.Spec.Resources,
						// ... readiness/liveness probes are essential here
					}},
				},
			},
			VolumeClaimTemplates: []corev1.PersistentVolumeClaim{{
				ObjectMeta: metav1.ObjectMeta{Name: "data"},
				Spec: corev1.PersistentVolumeClaimSpec{
					AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce},
					StorageClassName: &db.Spec.Storage.StorageClassName,
					Resources: corev1.ResourceList{
						corev1.ResourceStorage: resource.MustParse(db.Spec.Storage.Size),
					},
				},
			}},
		},
	}

	controllerutil.SetControllerReference(db, sts, r.Scheme)
	return sts
}

// serviceForClusteredDB and other helpers would be defined similarly

// Inside reconcileCluster function

// ... after getting the StatefulSet

currentImage := sts.Spec.Template.Spec.Containers[0].Image
desiredImage := db.Spec.Image

if currentImage != desiredImage {
    log.Info("Image version mismatch, initiating zero-downtime update", "current", currentImage, "desired", desiredImage)
    return r.performUpdate(ctx, log, db, sts)
}

// ... logic to check other spec changes (replicas, resources) would go here

// If no updates are needed, we reconcile the current state (e.g., check primary health)
// and update the status.
return r.reconcileStatus(ctx, log, db, sts)

// controllers/clustereddatabase_controller.go

const (
	PrimaryAnnotation = "db.example.com/role"
)

func (r *ClusteredDatabaseReconciler) performUpdate(ctx context.Context, log logr.Logger, db *dbv1alpha1.ClusteredDatabase, sts *appsv1.StatefulSet) (ctrl.Result, error) {
	// First, update the StatefulSet's template. This won't trigger pod updates due to OnDelete strategy.
	// The StatefulSet controller will update its `updateRevision`.
	sts.Spec.Template.Spec.Containers[0].Image = db.Spec.Image
	if err := r.Update(ctx, sts); err != nil {
		log.Error(err, "Failed to update StatefulSet spec for image change")
		return ctrl.Result{}, err
	}

	// Get all pods for this StatefulSet
	podList := &corev1.PodList{}
	// ... code to list pods with the correct labels

	// Identify the current primary
	primaryPodName := findPrimary(podList) // Helper to find pod with PrimaryAnnotation
	if primaryPodName == "" {
		log.Info("No primary found during update, may be initial setup. Requeuing.")
		// This could happen if the primary pod is down. We need a robust recovery strategy here.
		// For now, we'll requeue.
		return ctrl.Result{RequeueAfter: time.Second * 10}, nil
	}

	// 1. Update replicas first, in reverse order
	for i := int(*db.Spec.Replicas - 1); i >= 0; i-- {
		podName := fmt.Sprintf("%s-%d", db.Name, i)
		if podName == primaryPodName {
			continue // Skip the primary for now
		}

		// Check if this replica pod needs an update
		pod := findPodByName(podList, podName)
		if pod == nil {
			log.Info("Replica pod not found, waiting for it to be created", "pod", podName)
			return ctrl.Result{RequeueAfter: time.Second * 10}, nil
		}

		if pod.Labels["controller-revision-hash"] != sts.Status.UpdateRevision {
			log.Info("Deleting outdated replica pod to trigger update", "pod", pod.Name)
			if err := r.Delete(ctx, pod); err != nil {
				log.Error(err, "Failed to delete replica pod for update")
				return ctrl.Result{}, err
			}
			// Wait for the pod to be recreated and become ready
			log.Info("Waiting for replica to be updated", "pod", pod.Name)
			return ctrl.Result{RequeueAfter: time.Second * 15}, nil
		}
	}

	// 2. All replicas are updated. Now, perform failover.
	log.Info("All replicas are up-to-date. Preparing for primary failover.")

	// Find a suitable candidate for promotion (must be ready and updated)
	newPrimaryCandidate := findPromotionCandidate(podList, primaryPodName, sts.Status.UpdateRevision)
	if newPrimaryCandidate == nil {
		log.Error(nil, "No healthy, updated replica found to promote. Aborting update.")
		// This is a critical state. We should set a 'Degraded' condition on the CRD.
		return ctrl.Result{RequeueAfter: time.Minute * 1}, nil
	}

	log.Info("Promoting new primary", "candidate", newPrimaryCandidate.Name)
	
	// In a real scenario, you would exec into the pods:
	// 1. exec into old primary: run pre-demotion script (e.g., `pg_ctl stop -m fast`)
	// 2. exec into new primary: run promotion script (e.g., `pg_ctl promote`)
	
	// We simulate this by updating annotations and the primary service
	primaryPod := findPodByName(podList, primaryPodName)
	
	// Update annotations
	primaryPod.Annotations[PrimaryAnnotation] = "replica"
	newPrimaryCandidate.Annotations[PrimaryAnnotation] = "primary"

	if err := r.Update(ctx, primaryPod); err != nil { /* ... handle error */ }
	if err := r.Update(ctx, newPrimaryCandidate); err != nil { /* ... handle error */ }

	// Update the primary service to point to the new primary
	if err := r.updatePrimaryService(ctx, db, newPrimaryCandidate.Name); err != nil {
		log.Error(err, "CRITICAL: Failed to update primary service. Manual intervention may be required.")
		// Revert annotations if this fails? Complex error recovery is key here.
		return ctrl.Result{}, err
	}

	// 3. Delete the old primary pod
	log.Info("Deleting old primary pod to complete the update", "pod", primaryPod.Name)
	if err := r.Delete(ctx, primaryPod); err != nil {
		log.Error(err, "Failed to delete old primary pod")
		return ctrl.Result{}, err
	}

	log.Info("Zero-downtime update process complete. Waiting for cluster to stabilize.")
	return ctrl.Result{RequeueAfter: time.Second * 10}, nil
}

# Example primary-service.yaml
apiVersion: v1
kind: Service
metadata:
  name: my-db-primary
spec:
  ports:
  - port: 5432
    protocol: TCP
    targetPort: 5432
  selector:
    app: my-db
    db.example.com/role: primary # This selector is key
  type: ClusterIP

// Inside reconcile loop

// If a pod is not ready after a timeout
meta.SetStatusCondition(&db.Status.Conditions, metav1.Condition{
    Type:    dbv1alpha1.DegradedCondition,
    Status:  metav1.ConditionTrue,
    Reason:  "UpdateFailed",
    Message: fmt.Sprintf("Pod %s failed to become ready after update", pod.Name),
})
// Update the status on the API server
if err := r.Status().Update(ctx, db); err != nil {
    // handle error
}
// Requeue with exponential backoff
return ctrl.Result{RequeueAfter: nextBackoffDuration}, nil

// main.go

func main() {
	// ... setup code

	ctrl.NewControllerManagedBy(mgr). 
		For(&dbv1alpha1.ClusteredDatabase{}).
		WithOptions(controller.Options{MaxConcurrentReconciles: 1}). // CRITICAL!
		Owns(&appsv1.StatefulSet{}).
		Owns(&corev1.Service{}).
		Complete(&r)

	// ... start manager
}