Write rate limit batch update per point count (#6152)

agourlay · web-flow · commit 1a0862ebd3df · 2025-03-11T16:10:07.000+01:00
diff --git a/lib/collection/src/shards/replica_set/update.rs b/lib/collection/src/shards/replica_set/update.rs
@@ -10,7 +10,7 @@ use super::{ReplicaSetState, ReplicaState, ShardReplicaSet, clock_set};
 use crate::operations::point_ops::WriteOrdering;
 use crate::operations::types::{CollectionError, CollectionResult, UpdateResult, UpdateStatus};
 use crate::operations::{ClockTag, CollectionUpdateOperations, OperationWithClockTag};
-use crate::shards::shard::PeerId;
+use crate::shards::shard::{PeerId, Shard};
 use crate::shards::shard_trait::ShardOperation as _;
 
 /// Maximum number of attempts for applying an update with a new clock.
@@ -54,8 +54,7 @@ impl ShardReplicaSet {
         let result = match state {
             ReplicaState::Active => {
                 // Rate limit update operations on Active replica
-                // TODO(ratelimits) determine cost of update based on operation
-                self.check_write_rate_limiter(1, &hw_measurement)?;
+                self.check_operation_write_rate_limiter(&hw_measurement, local, &operation)?;
                 local.get().update(operation, wait, hw_measurement).await
             }
 
@@ -296,19 +295,11 @@ impl ShardReplicaSet {
 
                 if self.peer_is_active(this_peer_id) {
                     // Check write rate limiter before proceeding if replica active
-                    // TODO(ratelimits) determine cost of update based on operation
-
-                    self.check_write_rate_limiter_lazy(&hw_measurement_acc, || {
-                        let mut ratelimiter_cost = 1;
-
-                        // Estimate the cost based on affected points if filter is available.
-                        match local.estimate_request_cardinality(&operation.operation) {
-                            Ok(est) => ratelimiter_cost = 1.max(est.exp),
-                            Err(err) => log::error!("Estimating cardinality: {err:?}"),
-                        }
-
-                        ratelimiter_cost
-                    })?;
+                    self.check_operation_write_rate_limiter(
+                        &hw_measurement_acc,
+                        local,
+                        &operation,
+                    )?;
                 }
 
                 let operation = operation.clone();
@@ -541,6 +532,29 @@ impl ShardReplicaSet {
         Ok(Some(res))
     }
 
+    /// Check write rate limiter for the operation
+    ///
+    /// Lazily compute the cost of the operation and check against the write rate limiter
+    fn check_operation_write_rate_limiter(
+        &self,
+        hw_measurement: &HwMeasurementAcc,
+        local: &Shard,
+        operation: &OperationWithClockTag,
+    ) -> CollectionResult<()> {
+        self.check_write_rate_limiter_lazy(hw_measurement, || {
+            let mut ratelimiter_cost = 1;
+
+            // Estimate the cost based on affected points if filter is available.
+            match local.estimate_request_cardinality(&operation.operation) {
+                Ok(est) => ratelimiter_cost = 1.max(est.exp),
+                Err(err) => log::error!("Estimating cardinality: {err:?}"),
+            }
+
+            ratelimiter_cost
+        })?;
+        Ok(())
+    }
+
     /// Whether to send updates to the given peer
     ///
     /// A peer in dead state, or a locally disabled peer, will not accept updates.
diff --git a/tests/openapi/test_strictmode.py b/tests/openapi/test_strictmode.py
@@ -1091,7 +1091,7 @@ def test_strict_mode_write_rate_limiting(collection_name):
         assert response.ok, "Rate limiting should be disabled now"
 
 
-def test_strict_mode_write_rate_limiting_update_op(collection_name):
+def test_strict_mode_write_rate_limiting_filtered_update_op(collection_name):
     set_strict_mode(collection_name, {
         "enabled": True,
         "write_rate_limit": 7,
@@ -1131,6 +1131,47 @@ def test_strict_mode_write_rate_limiting_update_op(collection_name):
     assert response.status_code == 429
     assert "Rate limiting exceeded: Write rate limit exceeded: Operation requires 5 tokens but only" in response.json()['status']['error']
 
+def test_strict_mode_write_rate_limiting_batch_update_op(collection_name):
+    def upsert_points(ids: list[int]):
+        length = len(ids)
+        payloads = [{} for _ in range(length)]
+        vectors = [[1, 2, 3, 5] for _ in range(length)]
+        return request_with_validation(
+            api='/collections/{collection_name}/points/batch',
+            method="POST",
+            path_params={'collection_name': collection_name},
+            body={
+                "operations": [
+                    {
+                        "upsert": {
+                            "batch": {
+                                "ids": ids,
+                                "payloads": payloads,
+                                "vectors": vectors
+                            }
+                        }
+                    }
+                ]
+            }
+        )
+
+    set_strict_mode(collection_name, {
+        "enabled": True,
+        "write_rate_limit": 10,
+    })
+
+    # validate that updates with 11 points will never be allowed
+    response = upsert_points([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+    assert response.status_code == 429
+    assert "Rate limiting exceeded: Write rate limit exceeded, request larger than than rate limiter capacity, please try to split your request" in response.json()['status']['error']
+
+    # validate that updates with 10 points is allowed because there are enough tokens for each point
+    upsert_points([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).raise_for_status()
+
+    # doing it again fails because we already consumed 10 tokens
+    response = upsert_points([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+    assert response.status_code == 429
+    assert "Rate limiting exceeded: Write rate limit exceeded: Operation requires 10 tokens but only 0.0 were available. Retry after 60s" in response.json()['status']['error']
 
 def test_filter_many_conditions(collection_name):
     def search_request(condition_count: int):
