Adding support to bucket throttle param handling

Change-Id: I0a73b6b9074d1c9b149fa2fbb97fe625e23cfce5
Reviewed-on: https://review.couchbase.org/c/TAF/+/242445
Reviewed-by: <shiva.adky@couchbase.com>
Tested-by: Ashwin <ashwin.govindarajulu@couchbase.com>

Author: ashwin2002

couchbase_utils/bucket_utils/bucket_ready_functions.py

@@ -2665,7 +2665,10 @@ def create_default_bucket(
             enable_encryption_at_rest=False,
             encryption_at_rest_key_id=None,
             encryption_at_rest_dek_rotation_interval=None,
-            encryption_at_rest_dek_lifetime=None):
+            encryption_at_rest_dek_lifetime=None,
+            throttle_enabled=None,
+            throttle_hard_limit=None,
+            throttle_reserved=None):
 
         node_info = global_vars.cluster_util.get_nodes_self(cluster.master)
         if ram_quota:
@@ -2713,6 +2716,13 @@ def create_default_bucket(
         if fusion_log_store_uri is not None:
             bucket_obj.fusionLogstoreURI = fusion_log_store_uri
 
+        if throttle_enabled is not None:
+            bucket_obj.throttleEnabled = throttle_enabled
+        if throttle_hard_limit is not None:
+            bucket_obj.throttleHardLimit = throttle_hard_limit
+        if throttle_reserved is not None:
+            bucket_obj.throttleReserved = throttle_reserved
+
         if vbuckets is not None and storage == Bucket.StorageBackend.magma:
             bucket_obj.numVBuckets = vbuckets
             if vbuckets == 1024 and bucket_obj.ramQuotaMB < 1024:
@@ -3575,7 +3585,10 @@ def update_bucket_property(cluster_node, bucket, ram_quota_mb=None,
                                magma_key_tree_data_block_size=None,
                                magma_seq_tree_data_block_size=None,
                                durability_impossible_fallback=None,
-                               warmup_behavior=None):
+                               warmup_behavior=None,
+                               throttle_enabled=None,
+                               throttle_hard_limit=None,
+                               throttle_reserved=None):
         return BucketHelper(cluster_node).change_bucket_props(
             bucket, ramQuotaMB=ram_quota_mb, replicaNumber=replica_number,
             replicaIndex=replica_index, flushEnabled=flush_enabled,
@@ -3595,7 +3608,10 @@ def update_bucket_property(cluster_node, bucket, ram_quota_mb=None,
             magma_seq_tree_data_block_size=magma_seq_tree_data_block_size,
             storageBackend=storageBackend,
             durability_impossible_fallback=durability_impossible_fallback,
-            warmup_behavior=warmup_behavior)
+            warmup_behavior=warmup_behavior,
+            throttle_enabled=throttle_enabled,
+            throttle_hard_limit=throttle_hard_limit,
+            throttle_reserved=throttle_reserved)
 
     def update_all_bucket_maxTTL(self, cluster, maxttl=0):
         for bucket in cluster.buckets:

lib/BucketLib/BucketOperations_Rest.py

@@ -400,7 +400,10 @@ def create_bucket(self, bucket_params=dict()):
                 bucket_params.get(Bucket.warmupBehavior),
             Bucket.encryptionAtRestKeyId: bucket_params.get(Bucket.encryptionAtRestKeyId),
             Bucket.encryptionAtRestDekRotationInterval: bucket_params.get(Bucket.encryptionAtRestDekRotationInterval),
-            Bucket.encryptionAtRestDekLifetime: bucket_params.get(Bucket.encryptionAtRestDekLifetime)}
+            Bucket.encryptionAtRestDekLifetime: bucket_params.get(Bucket.encryptionAtRestDekLifetime),
+            Bucket.throttleReserved: bucket_params.get(Bucket.throttleReserved, None),
+            Bucket.throttleHardLimit: bucket_params.get(Bucket.throttleHardLimit, None),
+            Bucket.throttleEnabled: bucket_params.get(Bucket.throttleEnabled, False)}
 
         # Set Bucket's width/weight param only if serverless is enabled
         if bucket_params.get('serverless'):
@@ -539,7 +542,10 @@ def change_bucket_props(self, bucket, ramQuotaMB=None,
                             encryptionAtRestKeyId=None,
                             encryptionAtRestDekLifetime=None,
                             encryptionAtRestDekRotationInterval=None,
-                            warmup_behavior=None):
+                            warmup_behavior=None,
+                            throttle_enabled=None,
+                            throttle_hard_limit=None,
+                            throttle_reserved=None):
         params_dict = {}
         if ramQuotaMB:
             params_dict["ramQuotaMB"] = ramQuotaMB
@@ -603,6 +609,12 @@ def change_bucket_props(self, bucket, ramQuotaMB=None,
             params_dict[Bucket.encryptionAtRestDekRotationInterval] = encryptionAtRestDekRotationInterval
         if encryptionAtRestDekLifetime is not None:
             params_dict[Bucket.encryptionAtRestDekLifetime] = encryptionAtRestDekLifetime
+        if throttle_enabled is not None:
+            params_dict[Bucket.throttleEnabled] = throttle_enabled
+        if throttle_hard_limit is not None:
+            params_dict[Bucket.throttleHardLimit] = throttle_hard_limit
+        if throttle_reserved is not None:
+            params_dict[Bucket.throttleReserved] = throttle_reserved
 
         self.log.info("Updating bucket properties for {}".format(bucket.name))
         status, content = self.edit_bucket(bucket.name, params_dict)

lib/BucketLib/bucket.py

@@ -298,9 +298,12 @@ def __init__(self, new_params=dict()):
             Bucket.encryptionAtRestDekRotationInterval, 2592000)
         self.encryptionAtRestDekLifetime = new_params.get(
             Bucket.encryptionAtRestDekLifetime, 31536000)
-        self.throttleReserved = 18446744073709551615
-        self.throttleHardLimit = 18446744073709551615
-        self.throttleEnabled = False
+        self.throttleReserved = new_params.get(
+            Bucket.throttleReserved, 18446744073709551615)
+        self.throttleHardLimit = new_params.get(
+            Bucket.throttleHardLimit, 18446744073709551615)
+        self.throttleEnabled = new_params.get(
+            Bucket.throttleEnabled, False)
         self.encryptionAtRestInfo = {
             'dataStatus': new_params.get('dataStatus', 'unencrypted'),
             'dekNumber': new_params.get('dekNumber', 0),

pytests/cb_basetest.py

@@ -103,6 +103,12 @@ def setUp(self):
             # Couchstore bucket + Eviction Policy not set explicitly
             if self.bucket_eviction_policy is None:
                 self.bucket_eviction_policy = Bucket.EvictionPolicy.VALUE_ONLY
+        self.bucket_throttle_enabled = \
+            self.input.param("bucket_throttle_enabled", None)
+        self.bucket_throttle_hard_limit = \
+            self.input.param("bucket_throttle_hard_limit", None)
+        self.bucket_throttle_reserved = \
+            self.input.param("bucket_throttle_reserved", None)
         # End of bucket params
 
         # Doc specific params

pytests/onPrem_basetestcase.py

@@ -92,8 +92,7 @@ def setUp(self):
             self.input.param("disable_file_based_rebalance", False)
         # End of cluster info parameters
 
-        self.bucket_replica_index = self.input.param("bucket_replica_index",
-                                                     1)
+        self.bucket_replica_index = self.input.param("bucket_replica_index", 1)
         self.bucket_num_vb = self.input.param("bucket_num_vb", None)
         # End of bucket parameters
 
@@ -1872,6 +1871,14 @@ def create_bucket(self, cluster, bucket_name="default"):
         if self.bucket_num_vb is not None:
             create_bucket_params["vbuckets"] = self.bucket_num_vb
 
+        # Rate limiting feature values
+        if self.bucket_throttle_enabled is not None:
+            create_bucket_params["throttle_enabled"] = self.bucket_throttle_enabled
+        if self.bucket_throttle_reserved is not None:
+            create_bucket_params["throttle_reserved"] = self.bucket_throttle_reserved
+        if self.bucket_throttle_hard_limit is not None:
+            create_bucket_params["throttle_hard_limit"] = self.bucket_throttle_hard_limit
+
         # This is needed because server will throw the error saying,
         # "Support for variable number of vbuckets is not enabled"
         if CbServer.cluster_profile == "serverless":