Add vhost limits. (#977)

* Add vhost limits.
* Remove outdated vhost limits.

Author: MirahImage

api/v1beta1/vhost_types.go

@@ -34,6 +34,10 @@ type VhostSpec struct {
 	// +kubebuilder:validation:Enum=delete;retain
 	// +kubebuilder:default:=delete
 	DeletionPolicy string `json:"deletionPolicy,omitempty"`
+	// Limits defines limits to be applied to the vhost.
+	// Supported limits include max-connections and max-queues.
+	// See https://www.rabbitmq.com/docs/vhosts#limits
+	VhostLimits *VhostLimits `json:"limits,omitempty"`
 }
 
 // VhostStatus defines the observed state of Vhost
@@ -67,6 +71,12 @@ type VhostList struct {
 	Items           []Vhost `json:"items"`
 }
 
+// VhostLimits defines limits to be applied to the vhost.
+type VhostLimits struct {
+	Connections *int32 `json:"connections,omitempty"`
+	Queues      *int32 `json:"queues,omitempty"`
+}
+
 func (v *Vhost) GroupResource() schema.GroupResource {
 	return schema.GroupResource{
 		Group:    v.GroupVersionKind().Group,

api/v1beta1/vhost_types_test.go

@@ -101,6 +101,109 @@ var _ = Describe("Vhost", func() {
 		}))
 	})
 
+	Context("Vhost Limits", func() {
+		var connections, queues int32
+
+		When("vhost limits are configured", func() {
+			It("creates a vhost with limits", func() {
+				connections = 1000
+				queues = 500
+				vhost := Vhost{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "vhost-with-limits",
+						Namespace: namespace,
+					},
+					Spec: VhostSpec{
+						Name: "vhost-with-limits",
+						VhostLimits: &VhostLimits{
+							Connections: &connections,
+							Queues:      &queues,
+						},
+						RabbitmqClusterReference: RabbitmqClusterReference{
+							Name: "random-cluster",
+						},
+					},
+				}
+				Expect(k8sClient.Create(ctx, &vhost)).To(Succeed())
+				fetched := &Vhost{}
+				Expect(k8sClient.Get(ctx, types.NamespacedName{
+					Name:      vhost.Name,
+					Namespace: vhost.Namespace,
+				}, fetched)).To(Succeed())
+
+				Expect(*fetched.Spec.VhostLimits.Connections).To(Equal(connections))
+				Expect(*fetched.Spec.VhostLimits.Queues).To(Equal(queues))
+				Expect(fetched.Spec.Name).To(Equal("vhost-with-limits"))
+				Expect(fetched.Spec.RabbitmqClusterReference).To(Equal(RabbitmqClusterReference{
+					Name: "random-cluster",
+				}))
+			})
+		})
+
+		When("No vhost limits are provided", func() {
+			It("Does not set VhostLimits", func() {
+				vhost := Vhost{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "vhost-with-no-provided-limits",
+						Namespace: namespace,
+					},
+					Spec: VhostSpec{
+						Name: "vhost-with-no-provided-limits",
+						RabbitmqClusterReference: RabbitmqClusterReference{
+							Name: "random-cluster",
+						},
+					},
+				}
+				Expect(k8sClient.Create(ctx, &vhost)).To(Succeed())
+				fetched := &Vhost{}
+				Expect(k8sClient.Get(ctx, types.NamespacedName{
+					Name:      vhost.Name,
+					Namespace: vhost.Namespace,
+				}, fetched)).To(Succeed())
+
+				Expect(fetched.Spec.VhostLimits).To(BeNil())
+				Expect(fetched.Spec.Name).To(Equal("vhost-with-no-provided-limits"))
+				Expect(fetched.Spec.RabbitmqClusterReference).To(Equal(RabbitmqClusterReference{
+					Name: "random-cluster",
+				}))
+			})
+		})
+
+		When("Only some vhost limits are provided", func() {
+			It("Configures those limits and lifts other limits", func() {
+				queues = 800
+				vhost := Vhost{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "vhost-some-limits",
+						Namespace: namespace,
+					},
+					Spec: VhostSpec{
+						Name: "vhost-some-limits",
+						VhostLimits: &VhostLimits{
+							Queues: &queues,
+						},
+						RabbitmqClusterReference: RabbitmqClusterReference{
+							Name: "random-cluster",
+						},
+					},
+				}
+				Expect(k8sClient.Create(ctx, &vhost)).To(Succeed())
+				fetched := &Vhost{}
+				Expect(k8sClient.Get(ctx, types.NamespacedName{
+					Name:      vhost.Name,
+					Namespace: vhost.Namespace,
+				}, fetched)).To(Succeed())
+
+				Expect(fetched.Spec.VhostLimits.Connections).To(BeNil())
+				Expect(*fetched.Spec.VhostLimits.Queues).To(Equal(queues))
+				Expect(fetched.Spec.Name).To(Equal("vhost-some-limits"))
+				Expect(fetched.Spec.RabbitmqClusterReference).To(Equal(RabbitmqClusterReference{
+					Name: "random-cluster",
+				}))
+			})
+		})
+	})
+
 	Context("Default queue types", func() {
 		var qTypeVhost = &Vhost{
 			ObjectMeta: metav1.ObjectMeta{

api/v1beta1/zz_generated.deepcopy.go

@@ -59,6 +59,19 @@ spec:
                 - delete
                 - retain
                 type: string
+              limits:
+                description: |-
+                  Limits defines limits to be applied to the vhost.
+                  Supported limits include max-connections and max-queues.
+                  See https://www.rabbitmq.com/docs/vhosts#limits
+                properties:
+                  connections:
+                    format: int32
+                    type: integer
+                  queues:
+                    format: int32
+                    type: integer
+                type: object
               name:
                 description: Name of the vhost; see https://www.rabbitmq.com/vhosts.html.
                 type: string

config/crd/bases/rabbitmq.com_vhosts.yaml

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 
+	rabbithole "github.com/michaelklishin/rabbit-hole/v3"
 	"github.com/rabbitmq/messaging-topology-operator/internal"
 	"github.com/rabbitmq/messaging-topology-operator/rabbitmqclient"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -20,9 +21,35 @@ type VhostReconciler struct {
 	client.Client
 }
 
-func (r *VhostReconciler) DeclareFunc(_ context.Context, client rabbitmqclient.Client, obj topology.TopologyResource) error {
+func (r *VhostReconciler) DeclareFunc(ctx context.Context, client rabbitmqclient.Client, obj topology.TopologyResource) error {
+	logger := ctrl.LoggerFrom(ctx)
 	vhost := obj.(*topology.Vhost)
-	return validateResponse(client.PutVhost(vhost.Spec.Name, *internal.GenerateVhostSettings(vhost)))
+	settings := internal.GenerateVhostSettings(vhost)
+	logger.Info("generated vhost settings", "vhost", vhost.Spec.Name, "settings", settings)
+	err := validateResponse(client.PutVhost(vhost.Spec.Name, *settings))
+	if err != nil {
+		return err
+	}
+
+	newVhostLimits := internal.GenerateVhostLimits(vhost.Spec.VhostLimits)
+	logger.Info("getting existing vhost limits", vhost, vhost.Spec.Name)
+	existingVhostLimits, err := r.getVhostLimits(client, vhost.Spec.Name)
+	if err != nil {
+		return err
+	}
+	limitsToDelete := r.vhostLimitsToDelete(existingVhostLimits, newVhostLimits)
+	if len(limitsToDelete) > 0 {
+		logger.Info("Deleting outdated vhost limits", "vhost", vhost.Spec.Name, "limits", limitsToDelete)
+		err = validateResponseForDeletion(client.DeleteVhostLimits(vhost.Spec.Name, limitsToDelete))
+		if err != nil && !errors.Is(err, NotFound) {
+			return err
+		}
+	}
+	if len(newVhostLimits) > 0 {
+		logger.Info("creating new vhost limits", "vhost", vhost.Spec.Name, "limits", newVhostLimits)
+		return validateResponse(client.PutVhostLimits(vhost.Spec.Name, newVhostLimits))
+	}
+	return nil
 }
 
 // DeleteFunc deletes vhost from server
@@ -41,3 +68,28 @@ func (r *VhostReconciler) DeleteFunc(ctx context.Context, client rabbitmqclient.
 	}
 	return nil
 }
+
+func (r *VhostReconciler) vhostLimitsToDelete(existingVhostLimits, newVhostLimits rabbithole.VhostLimitsValues) (limitsToDelete rabbithole.VhostLimits) {
+	vhostLimitKeys := []string{"max-connections", "max-queues"}
+	for _, limit := range vhostLimitKeys {
+		_, oldExists := existingVhostLimits[limit]
+		_, newExists := newVhostLimits[limit]
+		if oldExists && !newExists {
+			limitsToDelete = append(limitsToDelete, limit)
+		}
+	}
+	return limitsToDelete
+}
+
+func (r *VhostReconciler) getVhostLimits(client rabbitmqclient.Client, vhost string) (rabbithole.VhostLimitsValues, error) {
+	vhostLimitsInfo, err := client.GetVhostLimits(vhost)
+	if errors.Is(err, error(rabbithole404)) {
+		return rabbithole.VhostLimitsValues{}, nil
+	} else if err != nil {
+		return rabbithole.VhostLimitsValues{}, err
+	}
+	if len(vhostLimitsInfo) == 0 {
+		return rabbithole.VhostLimitsValues{}, nil
+	}
+	return vhostLimitsInfo[0].Value, nil
+}