[Feature][Add] Added PodDisruptionBudget support (#48)

* Added PodDisruptioBudget support

Signed-off-by: iamabhishek-dubey <abhishekbhardwaj510@gmail.com>

api/v1alpha1/mongodb_types.go

@@ -30,7 +30,6 @@ type MongoDBSpec struct {
 
 // MongoDBStatus defines the observed state of MongoDB
 type MongoDBStatus struct {
-	MongoDB MongoDBSpec `json:"mongodb,omitempty"`
 }
 
 //+kubebuilder:object:root=true

api/v1alpha1/mongodbcluster_types.go

@@ -22,17 +22,24 @@ import (
 
 // MongoDBClusterSpec defines the desired state of MongoDBCluster
 type MongoDBClusterSpec struct {
-	MongoDBClusterSize *int32             `json:"clusterSize"`
-	EnableArbiter      *bool              `json:"enableMongoArbiter,omitempty"`
-	KubernetesConfig   KubernetesConfig   `json:"kubernetesConfig"`
-	Storage            *Storage           `json:"storage,omitempty"`
-	MongoDBSecurity    *MongoDBSecurity   `json:"mongoDBSecurity"`
-	MongoDBMonitoring  *MongoDBMonitoring `json:"mongoDBMonitoring,omitempty"`
+	MongoDBClusterSize  *int32                      `json:"clusterSize"`
+	EnableArbiter       *bool                       `json:"enableMongoArbiter,omitempty"`
+	KubernetesConfig    KubernetesConfig            `json:"kubernetesConfig"`
+	Storage             *Storage                    `json:"storage,omitempty"`
+	MongoDBSecurity     *MongoDBSecurity            `json:"mongoDBSecurity"`
+	MongoDBMonitoring   *MongoDBMonitoring          `json:"mongoDBMonitoring,omitempty"`
+	PodDisruptionBudget *MongoDBPodDisruptionBudget `json:"podDisruptionBudget,omitempty"`
+}
+
+// MongoDBPodDisruptionBudget defines the struct for MongoDB cluster
+type MongoDBPodDisruptionBudget struct {
+	Enabled        bool   `json:"enabled,omitempty"`
+	MinAvailable   *int32 `json:"minAvailable,omitempty"`
+	MaxUnavailable *int32 `json:"maxUnavailable,omitempty"`
 }
 
 // MongoDBClusterStatus defines the observed state of MongoDBCluster
 type MongoDBClusterStatus struct {
-	MongoDBCluster MongoDBClusterSpec `json:"mongodbCluster,omitempty"`
 }
 
 //+kubebuilder:object:root=true

config/rbac/role.yaml

@@ -97,3 +97,15 @@ rules:
   - get
   - patch
   - update
+- apiGroups:
+  - policy
+  resources:
+  - poddisruptionbudgets
+  verbs:
+  - create
+  - delete
+  - get
+  - list
+  - patch
+  - update
+  - watch

controllers/mongodbcluster_controller.go

@@ -39,6 +39,7 @@ type MongoDBClusterReconciler struct {
 //+kubebuilder:rbac:groups=opstreelabs.in,resources=mongodbclusters,verbs=get;list;watch;create;update;patch;delete
 //+kubebuilder:rbac:groups=opstreelabs.in,resources=mongodbclusters/status,verbs=get;update;patch
 //+kubebuilder:rbac:groups=opstreelabs.in,resources=mongodbclusters/finalizers,verbs=update
+//+kubebuilder:rbac:groups="policy",resources=poddisruptionbudgets,verbs=get;list;watch;create;update;patch;delete
 
 // Reconcile is part of the main kubernetes reconciliation loop which aims to
 func (r *MongoDBClusterReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {

examples/disruption-budget/clusterd.yaml

@@ -0,0 +1,23 @@
+---
+apiVersion: opstreelabs.in/v1alpha1
+kind: MongoDBCluster
+metadata:
+  name: mongodb
+spec:
+  clusterSize: 3
+  kubernetesConfig:
+    image: quay.io/opstree/mongo:v5.0
+    imagePullPolicy: IfNotPresent
+  storage:
+    accessModes: ["ReadWriteOnce"]
+    storageSize: 1Gi
+    storageClass: gp2
+  mongoDBSecurity:
+    mongoDBAdminUser: admin
+    secretRef:
+      name: mongodb-secret
+      key: password
+  podDisruptionBudget:
+    enabled: true
+    maxUnavailable: 1
+#    minAvailable: 1

k8sgo/cluster.go

@@ -68,6 +68,13 @@ func CreateMongoClusterSetup(cr *opstreelabsinv1alpha1.MongoDBCluster) error {
 		logger.Error(err, "Cannot create cluster StatefulSet for MongoDB")
 		return err
 	}
+	if cr.Spec.PodDisruptionBudget != nil && cr.Spec.PodDisruptionBudget.Enabled {
+		err = CreateOrUpdatePodDisruption(getPodDisruptionParams(cr))
+		if err != nil {
+			logger.Error(err, "Cannot create PodDisruptionBudget for MongoDB")
+			return err
+		}
+	}
 	return nil
 }
 
@@ -167,3 +174,22 @@ func getMongoDBClusterParams(cr *opstreelabsinv1alpha1.MongoDBCluster) statefulS
 	}
 	return params
 }
+
+// getPodDisruptionParams is a method to create parameters for pod disruption budget
+func getPodDisruptionParams(cr *opstreelabsinv1alpha1.MongoDBCluster) PodDisruptionParameters {
+	appName := fmt.Sprintf("%s-%s", cr.ObjectMeta.Name, "cluster")
+	labels := map[string]string{
+		"app":           appName,
+		"mongodb_setup": "cluster",
+		"role":          "cluster",
+	}
+	params := PodDisruptionParameters{
+		PDBMeta:        generateObjectMetaInformation(appName, cr.Namespace, labels, generateAnnotations()),
+		OwnerDef:       mongoClusterAsOwner(cr),
+		Namespace:      cr.Namespace,
+		Labels:         labels,
+		MinAvailable:   cr.Spec.PodDisruptionBudget.MinAvailable,
+		MaxUnavailable: cr.Spec.PodDisruptionBudget.MaxUnavailable,
+	}
+	return params
+}

k8sgo/poddisruption.go

@@ -0,0 +1,145 @@
+package k8sgo
+
+import (
+	"context"
+	"github.com/banzaicloud/k8s-objectmatcher/patch"
+	policyv1 "k8s.io/api/policy/v1beta1"
+	"k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/intstr"
+)
+
+// PodDisruptionParameters is an input parameter structure for Pod disruption budget
+type PodDisruptionParameters struct {
+	PDBMeta        metav1.ObjectMeta
+	OwnerDef       metav1.OwnerReference
+	Labels         map[string]string
+	Namespace      string
+	MinAvailable   *int32
+	MaxUnavailable *int32
+}
+
+// CreateOrUpdatePodDisruption method will create or update MongoDB PodDisruptionBudgets
+func CreateOrUpdatePodDisruption(params PodDisruptionParameters) error {
+	logger := logGenerator(params.PDBMeta.Name, params.Namespace, "PodDisruptionBudget")
+	pdbDef := generatePodDisruption(params)
+	storedPDB, err := getPodDisruption(params.Namespace, params.PDBMeta.Name)
+	if err != nil {
+		if err := patch.DefaultAnnotator.SetLastAppliedAnnotation(pdbDef); err != nil {
+			logger.Error(err, "Unable to patch MongoDB PodDisruptionBudget with comparison object")
+			return err
+		}
+		if errors.IsNotFound(err) {
+			return createPodDisruption(params.Namespace, pdbDef)
+		}
+		return err
+	}
+	return patchPodDisruption(storedPDB, pdbDef, params.Namespace)
+}
+
+// patchPodDisruption will patch MongoDB Kubernetes PodDisruptionBudgets
+func patchPodDisruption(storedPdb *policyv1.PodDisruptionBudget, newPdb *policyv1.PodDisruptionBudget, namespace string) error {
+	logger := logGenerator(newPdb.Name, namespace, "PodDisruptionBudget")
+	newPdb.ResourceVersion = storedPdb.ResourceVersion
+	newPdb.CreationTimestamp = storedPdb.CreationTimestamp
+	newPdb.ManagedFields = storedPdb.ManagedFields
+
+	storedPdb.Kind = "PodDisruptionBudget"
+	storedPdb.APIVersion = "policy/v1"
+
+	patchResult, err := patch.DefaultPatchMaker.Calculate(storedPdb, newPdb,
+		patch.IgnorePDBSelector(),
+		patch.IgnoreStatusFields(),
+	)
+	if err != nil {
+		logger.Error(err, "Unable to patch MongoDB PodDisruption with comparison object")
+		return err
+	}
+	if !patchResult.IsEmpty() {
+		logger.Info("Changes in PodDisruptionBudget Detected, Updating...",
+			"patch", string(patchResult.Patch),
+			"Current", string(patchResult.Current),
+			"Original", string(patchResult.Original),
+			"Modified", string(patchResult.Modified),
+		)
+		for key, value := range storedPdb.Annotations {
+			if _, present := newPdb.Annotations[key]; !present {
+				newPdb.Annotations[key] = value
+			}
+		}
+		if err := patch.DefaultAnnotator.SetLastAppliedAnnotation(newPdb); err != nil {
+			logger.Error(err, "Unable to patch MongoDB PodDisruptionBudget with comparison object")
+			return err
+		}
+		return updatePodDisruption(namespace, newPdb)
+	}
+	logger.Info("PodDisruptionBudget is reconciled, nothing to change")
+	return nil
+}
+
+// updatePodDisruption is a method to create Pod disruption budget
+func updatePodDisruption(namespace string, pdb *policyv1.PodDisruptionBudget) error {
+	logger := logGenerator(pdb.Name, namespace, "PodDisruptionBudget")
+	_, err := generateK8sClient().PolicyV1beta1().PodDisruptionBudgets(namespace).Update(context.TODO(), pdb, metav1.UpdateOptions{})
+	if err != nil {
+		logger.Error(err, "MongoDB PodDisruptionBudget update failed")
+		return err
+	}
+	logger.Info("MongoDB PodDisruptionBudget update was successful", "PDB.Spec", pdb.Spec)
+	return nil
+}
+
+// createPodDisruption is a method to create Pod disruption budget
+func createPodDisruption(namespace string, pdb *policyv1.PodDisruptionBudget) error {
+	logger := logGenerator(pdb.Name, namespace, "PodDisruptionBudget")
+	_, err := generateK8sClient().PolicyV1beta1().PodDisruptionBudgets(namespace).Create(context.TODO(), pdb, metav1.CreateOptions{})
+	if err != nil {
+		logger.Error(err, "MongoDB PodDisruptionBudget creation failed")
+		return err
+	}
+	logger.Info("MongoDB PodDisruptionBudget creation was successful", "PDB.Spec", pdb.Spec)
+	return nil
+}
+
+// getPodDisruption is a method to get Pod disruption budget
+func getPodDisruption(namespace, name string) (*policyv1.PodDisruptionBudget, error) {
+	logger := logGenerator(name, namespace, "PodDisruptionBudget")
+	pdbInfo, err := generateK8sClient().PolicyV1beta1().PodDisruptionBudgets(namespace).Get(context.TODO(), name, metav1.GetOptions{})
+	if err != nil {
+		logger.Info("Unable to get pod disruption budget")
+		return nil, err
+	}
+	logger.Info("MongoDB PodDisruptionBudget get action was successful")
+	return pdbInfo, err
+}
+
+// generatePodDisruption is a method to generate Pod disruption budget definiton
+func generatePodDisruption(params PodDisruptionParameters) *policyv1.PodDisruptionBudget {
+	pdbTemplate := &policyv1.PodDisruptionBudget{
+		TypeMeta:   generateMetaInformation("PodDisruptionBudget", "policy/v1beta1"),
+		ObjectMeta: params.PDBMeta,
+		Spec: policyv1.PodDisruptionBudgetSpec{
+			Selector: &metav1.LabelSelector{
+				MatchLabels: params.Labels,
+			},
+			MaxUnavailable: &intstr.IntOrString{
+				Type:   intstr.Int,
+				IntVal: int32(*params.MaxUnavailable),
+			},
+		},
+	}
+	if params.MinAvailable != nil {
+		pdbTemplate.Spec.MinAvailable = &intstr.IntOrString{
+			Type:   intstr.Int,
+			IntVal: int32(*params.MinAvailable),
+		}
+	}
+	if params.MaxUnavailable != nil {
+		pdbTemplate.Spec.MaxUnavailable = &intstr.IntOrString{
+			Type:   intstr.Int,
+			IntVal: int32(*params.MaxUnavailable),
+		}
+	}
+	AddOwnerRefToObject(pdbTemplate, params.OwnerDef)
+	return pdbTemplate
+}