Stateful Workload Orchestration with Custom Kubernetes Operators

apiVersion: apiextensions.k8s.io/v1
kind: CustomResourceDefinition
metadata:
  name: replicatedkvstores.database.mycorp.com
spec:
  group: database.mycorp.com
  names:
    kind: ReplicatedKVStore
    listKind: ReplicatedKVStoreList
    plural: replicatedkvstores
    singular: replicatedkvstore
  scope: Namespaced
  versions:
    - name: v1alpha1
      served: true
      storage: true
      schema:
        openAPIV3Schema:
          type: object
          properties:
            spec:
              type: object
              properties:
                replicas:
                  type: integer
                  minimum: 1
                  description: The desired number of replicas in the key-value store cluster.
                version:
                  type: string
                  description: The container image version to deploy.
                storageClassName:
                  type: string
                  description: The storage class for persistent volumes.
                volumeSize:
                  type: string
                  pattern: '^[0-9]+(Gi|Mi|Ki)$'
                  description: The size of the persistent volume for each replica (e.g., 10Gi).
              required:
                - replicas
                - version
                - storageClassName
                - volumeSize
            status:
              type: object
              properties:
                phase:
                  type: string
                  description: The current phase of the cluster (e.g., Creating, Ready, Failed).
                readyReplicas:
                  type: integer
                  description: The number of replicas that are fully initialized and ready.
                currentVersion:
                  type: string
                  description: The currently deployed version of the cluster.
                conditions:
                  type: array
                  items:
                    type: object
                    properties:
                      type:
                        type: string
                      status:
                        type: string
                      lastTransitionTime:
                        type: string
                        format: date-time
                      message:
                        type: string

package controller

import (
	"context"

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

	databasev1alpha1 "my-operator/api/v1alpha1"
)

// ReplicatedKVStoreReconciler reconciles a ReplicatedKVStore object
type ReplicatedKVStoreReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

func (r *ReplicatedKVStoreReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	// 1. Fetch the ReplicatedKVStore instance
	instance := &databasev1alpha1.ReplicatedKVStore{}
	err := r.Get(ctx, req.NamespacedName, instance)
	if err != nil {
		if errors.IsNotFound(err) {
			// Object was deleted. The reconciliation is complete.
			logger.Info("ReplicatedKVStore resource not found. Ignoring since object must be deleted")
			return ctrl.Result{}, nil
		}
		// Error reading the object - requeue the request.
		logger.Error(err, "Failed to get ReplicatedKVStore")
		return ctrl.Result{}, err
	}

	// *** CORE RECONCILIATION LOGIC GOES HERE ***

	// 2. Check for underlying StatefulSet
	foundSts := &appsv1.StatefulSet{}
	err = r.Get(ctx, types.NamespacedName{Name: instance.Name, Namespace: instance.Namespace}, foundSts)
	if err != nil && errors.IsNotFound(err) {
		// Define and create a new StatefulSet
		sts := r.statefulSetForKVStore(instance)
		logger.Info("Creating a new StatefulSet", "StatefulSet.Namespace", sts.Namespace, "StatefulSet.Name", sts.Name)
		err = r.Create(ctx, sts)
		if err != nil {
			logger.Error(err, "Failed to create new StatefulSet", "StatefulSet.Namespace", sts.Namespace, "StatefulSet.Name", sts.Name)
			return ctrl.Result{}, err
		}
		// StatefulSet created successfully - return and requeue
		return ctrl.Result{Requeue: true}, nil
	} else if err != nil {
		logger.Error(err, "Failed to get StatefulSet")
		return ctrl.Result{}, err
	}

	// 3. Ensure StatefulSet spec matches the CR spec
	// Example: Check replica count
	desiredReplicas := instance.Spec.Replicas
	if *foundSts.Spec.Replicas != desiredReplicas {
		foundSts.Spec.Replicas = &desiredReplicas
		err = r.Update(ctx, foundSts)
		if err != nil {
			logger.Error(err, "Failed to update StatefulSet replicas")
			return ctrl.Result{}, err
		}
		// Spec updated, requeue to check status after update
		return ctrl.Result{Requeue: true}, nil
	}

	// 4. Update the status of the ReplicatedKVStore resource
	// ... logic to check pod readiness and update instance.Status ...

	return ctrl.Result{}, nil
}

// statefulSetForKVStore returns a ReplicatedKVStore StatefulSet object
func (r *ReplicatedKVStoreReconciler) statefulSetForKVStore(cr *databasev1alpha1.ReplicatedKVStore) *appsv1.StatefulSet {
	// ... Logic to construct the StatefulSet YAML programmatically ...
	// This is where you define the container spec, volumes, probes, etc.
	// It's CRITICAL to set the OwnerReference here!
	sts := &appsv1.StatefulSet{ /* ... */ }
	ctrl.SetControllerReference(cr, sts, r.Scheme)
	return sts
}

// SetupWithManager sets up the controller with the Manager.
func (r *ReplicatedKVStoreReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&databasev1alpha1.ReplicatedKVStore{}).
		Owns(&appsv1.StatefulSet{}). // This makes the controller watch StatefulSets it creates
		Complete(r)
}

func (r *ReplicatedKVStoreReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	instance := &databasev1alpha1.ReplicatedKVStore{}
	// ... fetch instance ...

	// The finalizer name
	myFinalizerName := "database.mycorp.com/finalizer"

	// Check if the object is being deleted
	if instance.ObjectMeta.DeletionTimestamp.IsZero() {
		// The object is not being deleted, so if it does not have our finalizer,
		// then lets add the finalizer and update the object.
		if !controllerutil.ContainsFinalizer(instance, myFinalizerName) {
			controllerutil.AddFinalizer(instance, myFinalizerName)
			if err := r.Update(ctx, instance); err != nil {
				return ctrl.Result{}, err
			}
		}
	} else {
		// The object is being deleted
		if controllerutil.ContainsFinalizer(instance, myFinalizerName) {
			// Our finalizer is present, so lets handle any external dependency
			if err := r.cleanupExternalResources(ctx, instance); err != nil {
				// if fail to delete the external dependency here, return with error
				// so that it can be retried
				logger.Error(err, "Failed to run cleanup logic")
				return ctrl.Result{}, err
			}

			// Once cleanup is successful, remove the finalizer from the list and update it.
			controllerutil.RemoveFinalizer(instance, myFinalizerName)
			if err := r.Update(ctx, instance); err != nil {
				return ctrl.Result{}, err
			}
		}

		// Stop reconciliation as the item is being deleted
		return ctrl.Result{}, nil
	}

	// ... rest of the reconciliation logic for create/update ...

	return ctrl.Result{}, nil
}

func (r *ReplicatedKVStoreReconciler) cleanupExternalResources(ctx context.Context, cr *databasev1alpha1.ReplicatedKVStore) error {
	logger := log.FromContext(ctx)
	logger.Info("Executing finalizer logic for ReplicatedKVStore")

	// This is where you implement the application-specific cleanup.
	// For example, connect to the cluster and gracefully remove nodes.
	// This logic must be idempotent.
	
	// Example: Call a 'decommission' endpoint on each pod before it's terminated.
	// This would require a more complex orchestration than shown here, possibly
	// by scaling down the statefulset one by one and waiting for confirmation.

	logger.Info("Finalizer logic completed successfully")
	return nil
}