MB-70082 Add support for KV rate limiting to cluster-init

Change-Id: I9db2ddca2a82dc3fdaffe4e6355e5ce75f19e898
Reviewed-on: https://review.couchbase.org/c/couchbase-cli/+/242093
Tested-by: Build Bot <build@couchbase.com>
Reviewed-by: Matt Hall <matt.hall@couchbase.com>

Author: lubomarinski

cbmgr.py

@@ -1109,6 +1109,15 @@ def __init__(self):
                            help="Set the IP family for the cluster")
         group.add_argument("--node-to-node-encryption", dest="encryption", metavar="<on|off>", default="off",
                            choices=["on", "off"], help="Enable node to node encryption")
+        group.add_argument("--node-capacity", dest="node_capacity", metavar="<units>", type=(int),
+                           help="The number of units (bytes) per second. This is a global counter that any "
+                           "bucket can consume from")
+        group.add_argument("--throttle-enabled", dest="throttle_enabled", metavar="<1|0>", choices=["0", "1"],
+                           help="Enable/Disable throttling feature")
+        group.add_argument("--read-unit-size", dest="read_unit_size", metavar="<bytes>", type=(int),
+                           help="Size of read that will consume 1 unit")
+        group.add_argument("--write-unit-size", dest="write_unit_size", metavar="<bytes>", type=(int),
+                           help="Size of write that will consume 1 unit")
 
     @rest_initialiser(enterprise_check=False, enterprise_analytics_check=False)
     def execute(self, opts):
@@ -1182,6 +1191,13 @@ def execute(self, opts):
             send_stats=opts.notifications == "1")
         _exit_if_errors(errors)
 
+        # Set global throttling settings if any were provided
+        if (opts.node_capacity is not None or opts.throttle_enabled is not None or opts.read_unit_size is not None or
+                opts.write_unit_size is not None):
+            _, errors = self.rest.set_global_memcached_settings(
+                opts.node_capacity, opts.throttle_enabled, opts.read_unit_size, opts.write_unit_size)
+            _exit_if_errors(errors)
+
         _success("Cluster initialized")
 
     def setup_ip_family_and_encryption(self, opts):

cluster_manager.py

@@ -1533,6 +1533,26 @@ def cluster_init(self, services=None, username=None, password=None, port=None,
 
         return self._post_form_encoded(url, params)
 
+    def set_global_memcached_settings(
+            self,
+            node_capacity=None,
+            throttle_enabled=None,
+            read_unit_size=None,
+            write_unit_size=None):
+        url = f'{self.hostname}/pools/default/settings/memcached/global'
+        params = {}
+
+        if node_capacity is not None:
+            params["nodeCapacity"] = node_capacity
+        if throttle_enabled is not None:
+            params["throttleEnabled"] = "true" if throttle_enabled == "1" else "false"
+        if read_unit_size is not None:
+            params["readUnitSize"] = read_unit_size
+        if write_unit_size is not None:
+            params["writeUnitSize"] = write_unit_size
+
+        return self._post_form_encoded(url, params)
+
     def node_info(self):
         url = f'{self.hostname}/nodes/self'
         return self._get(url)

docs/modules/cli/pages/cbcli/couchbase-cli-cluster-init.adoc

@@ -22,6 +22,8 @@ _couchbase-cli cluster-init_ [--cluster <url>] [--cluster-username <username>] [
     [--index-storage-setting <setting>] [--services <services>]
     [--update-notifications <1|0>]
     [--ip-family <ipv4|ipv6|ipv4-only|ipv6-only>] [--node-to-node-encryption <on|off>]
+    [--node-capacity <capacity>] [--throttle-enabled <1|0>]
+    [--read-unit-size <size>] [--write-unit-size <size>]
 
 == DESCRIPTION
 
@@ -149,6 +151,21 @@ which will determine how the global secondary indexes (GSI) are stored.
   'off'. When set to 'on' all communications between nodes will be over a
   encrypted connection.
 
+--node-capacity::
+  The number of units (bytes) per second. This is a global counter that any
+  bucket can consume from. This sets the total amount of resources available on
+  the node.
+
+--throttle-enabled::
+  Enable or disable the throttling feature. Accepted values are "1" to enable
+  and "0" to disable.
+
+--read-unit-size::
+  The size in bytes of read that will consume 1 unit.
+
+--write-unit-size::
+  The size in bytes of write that will consume 1 unit.
+
 include::{partialsdir}/cbcli/part-host-formats.adoc[]
 
 include::{partialsdir}/cbcli/part-certificate-authentication.adoc[]

test/mock_server.py

@@ -570,6 +570,7 @@ def export_eventing_functions(rest_params=None, server_args=None, path="", endpo
     (r'/settings/retryRebalance$', {'GET': get_by_path, 'POST': do_nothing}),
     (r'/settings/replications/(\d|\w)+$', {'POST': do_nothing, 'GET': get_by_path}),
     (r'/settings/analytics', {'POST': do_nothing, 'GET': get_by_path}),
+    (r'/pools/default/settings/memcached/global$', {'POST': do_nothing}),
     (r'/nodes/self/controller/settings$', {'POST': do_nothing}),
     (r'/nodes/self$', {'GET': get_node_info}),
     (r'/nodeInit', {'POST': do_nothing}),

test/test_cli.py

@@ -409,6 +409,49 @@ def test_error_manager_only_service_enterprise_analytics(self):
                              ['--services', 'manager-only'], self.server_args)
         self.assertIn('--services cannot be specified on Enterprise Analytics', self.str_output)
 
+    def test_init_cluster_with_throttle_settings(self):
+        self.server_args['init'] = False
+        full_options = cluster_connect_args[:2] + self.command_args + [
+            '--cluster-ramsize', '512', '--services', 'data',
+            '--node-capacity', '10000', '--throttle-enabled', '1',
+            '--read-unit-size', '1024', '--write-unit-size', '2048'
+        ]
+
+        self.no_error_run(self.command + full_options, self.server_args)
+        self.assertIn('SUCCESS', self.str_output)
+        expected_params = ['memoryQuota=512', 'username=Administrator', 'password=asdasd', 'port=6789',
+                           'nodeCapacity=10000', 'throttleEnabled=true', 'readUnitSize=1024', 'writeUnitSize=2048']
+        self.rest_parameter_match(expected_params, False)
+        self.assertIn('POST:/pools/default/settings/memcached/global', self.server.trace)
+
+    def test_init_cluster_with_throttle_disabled(self):
+        self.server_args['init'] = False
+        full_options = cluster_connect_args[:2] + self.command_args + [
+            '--cluster-ramsize', '512', '--services', 'data',
+            '--node-capacity', '10000', '--throttle-enabled', '0'
+        ]
+
+        self.no_error_run(self.command + full_options, self.server_args)
+        self.assertIn('SUCCESS', self.str_output)
+        expected_params = ['memoryQuota=512', 'username=Administrator', 'password=asdasd', 'port=6789',
+                           'nodeCapacity=10000', 'throttleEnabled=false']
+        self.rest_parameter_match(expected_params, False)
+        self.assertIn('POST:/pools/default/settings/memcached/global', self.server.trace)
+
+    def test_init_cluster_with_only_node_capacity(self):
+        self.server_args['init'] = False
+        full_options = cluster_connect_args[:2] + self.command_args + [
+            '--cluster-ramsize', '512', '--services', 'data',
+            '--node-capacity', '50000'
+        ]
+
+        self.no_error_run(self.command + full_options, self.server_args)
+        self.assertIn('SUCCESS', self.str_output)
+        expected_params = ['memoryQuota=512', 'username=Administrator', 'password=asdasd', 'port=6789',
+                           'nodeCapacity=50000']
+        self.rest_parameter_match(expected_params, False)
+        self.assertIn('POST:/pools/default/settings/memcached/global', self.server.trace)
+
 
 class TestBucketCompact(CommandTest):
     def setUp(self):