Improvement: improved in-memory-bucket performance (#159)

* refactor algorithm for inmemoryBucket

* version bump

* update pre-commit config

* up

* fix

* fix

* revert master

.pre-commit-config.yaml

@@ -1,29 +1,25 @@
 repos:
--   repo: https://github.com/psf/black
-    rev: 22.3.0
-    hooks:
-    -   id: black
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.3.0
+    rev: v4.5.0
     hooks:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
 -   repo: https://github.com/PyCQA/flake8
-    rev: 5.0.4
+    rev: 7.0.0
     hooks:
     -   id: flake8
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.930
+    rev: v1.9.0
     hooks:
     -   id: mypy
         additional_dependencies: [types-filelock, types-redis]
 -   repo: https://github.com/pre-commit/pygrep-hooks
-    rev: v1.9.0
+    rev: v1.10.0
     hooks:
     -   id: python-no-eval
     -   id: python-use-type-annotations
 -   repo: https://github.com/asottile/reorder_python_imports
-    rev: v2.6.0
+    rev: v3.12.0
     hooks:
     -   id: reorder-python-imports

noxfile.py

@@ -4,7 +4,7 @@
 # Reuse virtualenv created by poetry instead of creating new ones
 nox.options.reuse_existing_virtualenvs = True
 
-PYTEST_ARGS = ["--verbose", "-s", "--full-trace", "--maxfail=1", "--numprocesses=auto", "tests/test_bucket_factory.py"]
+PYTEST_ARGS = ["--verbose", "-s", "--full-trace", "--maxfail=1", "--numprocesses=auto"]
 COVERAGE_ARGS = ["--cov", "--cov-report=term", "--cov-report=xml", "--cov-report=html"]
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyrate-limiter"
-version = "3.3.0"
+version = "3.4.0"
 description = "Python Rate-Limiter using Leaky-Bucket Algorithm"
 authors = ["vutr <[email protected]>"]
 license = "MIT"

pyrate_limiter/buckets/in_memory_bucket.py

@@ -26,7 +26,16 @@ def __init__(self, rates: List[Rate]):
         self.items = []
 
     def put(self, item: RateItem) -> bool:
+        if item.weight == 0:
+            return True
+
+        current_length = len(self.items)
+        after_length = item.weight + current_length
+
         for rate in self.rates:
+            if after_length < rate.limit:
+                break
+
             lower_bound_value = item.timestamp - rate.interval
             lower_bound_idx = binary_search(self.items, lower_bound_value)
 
@@ -41,7 +50,12 @@ def put(self, item: RateItem) -> bool:
                 return False
 
         self.failing_rate = None
-        self.items.extend(item.weight * [item])
+
+        if item.weight > 1:
+            self.items.extend([item for _ in range(item.weight)])
+        else:
+            self.items.append(item)
+
         return True
 
     def leak(self, current_timestamp: Optional[int] = None) -> int:

tests/conftest.py

@@ -94,19 +94,6 @@ async def create_async_redis_bucket(rates: List[Rate]):
     return bucket
 
 
-async def create_sentinel_redis_bucket(rates: List[Rate]):
-    from redis.sentinel import Sentinel
-
-    sentinel = Sentinel([('localhost', 26379), ('localhost', 26380)], socket_timeout=0.1)
-    master = sentinel.master_for('mymaster', socket_timeout=0.1)
-    bucket_key = f"test-bucket/{id_generator()}"
-    assert master.ping()
-    master.delete(bucket_key)
-    bucket = RedisBucket.init(rates, master, bucket_key)
-    assert bucket.count() == 0
-    return bucket
-
-
 async def create_sqlite_bucket(rates: List[Rate]):
     temp_dir = Path(gettempdir())
     default_db_path = temp_dir / f"pyrate_limiter_{id_generator(size=5)}.sqlite"
@@ -144,7 +131,6 @@ async def create_sqlite_bucket(rates: List[Rate]):
         create_redis_bucket,
         create_sqlite_bucket,
         create_async_redis_bucket,
-        create_sentinel_redis_bucket,
     ]
 )
 def create_bucket(request):