Fix Postgres Distributed rates (#248) (#250)

paultiq · Copilot · web-flow · commit 1615e35502b8 · 2025-12-30T09:26:37.000+07:00
* fix: make put's atomic - lock, check, put... rather than check and put

* Update postgres.py

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;

* fix: Resolve comments - wrap _create_table in a pg_advisory_xact_lock and parameterize rate.interval

* fix: ruff format

* add log message for locknotavailable

* fix: Check sliding windows for both rates

---------

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/pyrate_limiter/buckets/postgres.py b/pyrate_limiter/buckets/postgres.py
@@ -26,6 +26,9 @@ class Queries:
     CREATE_INDEX_ON_TIMESTAMP = """
     CREATE INDEX IF NOT EXISTS {index} ON {table} (item_timestamp)
     """
+    LOCK_TABLE = """
+    LOCK TABLE {table} IN EXCLUSIVE MODE NOWAIT
+    """
     COUNT = """
     SELECT COUNT(*) FROM {table}
     """
@@ -73,6 +76,8 @@ def _get_conn(self):
 
     def _create_table(self):
         with self._get_conn() as conn:
+            lock_id = hash(self._full_tbl) & 0x7FFFFFFF
+            conn.execute("SELECT pg_advisory_xact_lock(%s)", (lock_id,))
             conn.execute(Queries.CREATE_BUCKET_TABLE.format(table=self._full_tbl))
             index_name = f"timestampIndex_{self.table}"
             conn.execute(Queries.CREATE_INDEX_ON_TIMESTAMP.format(table=self._full_tbl, index=index_name))
@@ -81,14 +86,38 @@ def put(self, item: RateItem) -> Union[bool, Awaitable[bool]]:
         """Put an item (typically the current time) in the bucket
         return true if successful, otherwise false
         """
+        from psycopg.errors import LockNotAvailable
+
         if item.weight == 0:
             return True
 
+        item_ts_seconds = item.timestamp / 1000
+
         with self._get_conn() as conn:
+            # Acquire an EXCLUSIVE MODE lock on the bucket table using NOWAIT.
+            # This ensures the "check current count" + "insert new items" sequence
+            # is atomic with respect to other writers, so rate limits cannot be
+            # exceeded due to concurrent requests interleaving.
+            #
+            # Because we use NOWAIT, if the table is already locked by another
+            # transaction, PostgreSQL raises LockNotAvailable and we immediately
+            # reject this request (return False) instead of blocking or retrying.
+            # This provides predictable, fail-fast behavior but may limit
+            # throughput under high contention since only one writer can perform
+            # the check-and-put at a time.
+            try:
+                conn.execute(Queries.LOCK_TABLE.format(table=self._full_tbl))
+            except LockNotAvailable:
+                logger.debug("LockNotAvailable")
+                self.failing_rate = self.rates[0]
+                return False
+
             for rate in self.rates:
-                bound = f"SELECT TO_TIMESTAMP({item.timestamp / 1000}) - INTERVAL '{rate.interval} milliseconds'"
-                query = f"SELECT COUNT(*) FROM {self._full_tbl} WHERE item_timestamp >= ({bound})"  # noqa: S608  # FIXME: SQL Parameterization and table name sanitization
-                cur = conn.execute(query)
+                cur = conn.execute(
+                    f"SELECT COUNT(*) FROM {self._full_tbl} "  # noqa: S608
+                    f"WHERE item_timestamp >= TO_TIMESTAMP(%s) - (%s * INTERVAL '1 milliseconds')",
+                    (item_ts_seconds, rate.interval),
+                )
                 count = int(cur.fetchone()[0])
                 cur.close()
 
@@ -97,13 +126,9 @@ def put(self, item: RateItem) -> Union[bool, Awaitable[bool]]:
                     return False
 
             self.failing_rate = None
-
             query = Queries.PUT.format(table=self._full_tbl)
-
-            # https://www.psycopg.org/docs/extras.html#fast-exec
-
             for _ in range(item.weight):
-                conn.execute(query, (item.name, item.weight, item.timestamp / 1000))
+                conn.execute(query, (item.name, item.weight, item_ts_seconds))
 
         return True
 
diff --git a/tests/test_postgres_concurrent.py b/tests/test_postgres_concurrent.py
@@ -0,0 +1,182 @@
+import threading
+
+import pytest
+
+from pyrate_limiter import Duration, PostgresBucket, Rate
+from pyrate_limiter.abstracts import RateItem
+
+
+@pytest.mark.postgres
+class TestPostgresConcurrent:
+
+    @pytest.fixture
+    def pg_pool(self):
+        from psycopg_pool import ConnectionPool
+
+        pool = ConnectionPool(
+            "postgresql://postgres:postgres@localhost:5432",
+            min_size=4,
+            max_size=10,
+            open=True,
+        )
+        yield pool
+        pool.close()
+
+    @pytest.fixture
+    def clean_table(self, pg_pool):
+        from pyrate_limiter import id_generator
+
+        table = f"test_concurrent_{id_generator()}"
+        yield table
+        with pg_pool.connection() as conn:
+            conn.execute(f"DROP TABLE IF EXISTS ratelimit___{table}")
+
+    def test_concurrent_put(self, pg_pool, clean_table):
+        rate_limit = 5
+        rates = [Rate(rate_limit, Duration.SECOND)]
+        num_threads = 8
+        attempts_per_thread = 10
+
+        results = []
+        results_lock = threading.Lock()
+
+        def worker(thread_id: int):
+            bucket = PostgresBucket(pg_pool, clean_table, rates)
+            thread_results = []
+
+            for _ in range(attempts_per_thread):
+                timestamp = bucket.now()
+                item = RateItem(f"thread_{thread_id}", timestamp, weight=1)
+                success = bucket.put(item)
+                thread_results.append((timestamp, success))
+
+            with results_lock:
+                results.extend(thread_results)
+
+        threads = [
+            threading.Thread(target=worker, args=(i,)) for i in range(num_threads)
+        ]
+
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # Verify db state
+        full_table = f"ratelimit___{clean_table}"
+        with pg_pool.connection() as conn:
+            cur = conn.execute(f"SELECT COUNT(*) FROM {full_table}")  # noqa: S608
+            total_in_db = cur.fetchone()[0]
+            cur.close()
+
+            # Check in sliding windows to make sure rate didn't exceed
+            cur = conn.execute(
+                f"SELECT EXTRACT(EPOCH FROM item_timestamp)::bigint as ts FROM {full_table}"  # noqa: S608
+            )
+            db_timestamps = [row[0] for row in cur.fetchall()]
+            cur.close()
+
+        for ts in db_timestamps:
+            window_start = ts - 1  # 1 second window
+            count_in_window = sum(1 for t in db_timestamps if t > window_start and t <= ts)
+            assert count_in_window <= rate_limit, (
+                f"Rate limit exceeded in DB: {count_in_window} items in 1-second window ending at {ts}"
+            )
+
+        # Verify anything worked
+        total_success = sum(1 for _, success in results if success)
+        assert total_success > 0, "No successful acquisitions"
+        assert total_success == total_in_db, (
+            f"Mismatch: {total_success} reported successes but {total_in_db} items in DB"
+        )
+
+        # Verify some rejections
+        total_rejected = sum(1 for _, success in results if not success)
+        assert total_rejected > 0, (
+            "No rejections occurred - rate limiting may not be working"
+        )
+
+    def test_concurrent_put_multiple_rates(self, pg_pool, clean_table):
+        rates = [
+            Rate(3, 500),   # 3 per 500ms
+            Rate(5, 1000),  # 5 per second
+        ]
+        num_threads = 4
+        attempts_per_thread = 5
+
+        results = []
+        results_lock = threading.Lock()
+
+        def worker(thread_id: int):
+            bucket = PostgresBucket(pg_pool, clean_table, rates)
+            thread_results = []
+
+            for _ in range(attempts_per_thread):
+                timestamp = bucket.now()
+                item = RateItem(f"thread_{thread_id}", timestamp, weight=1)
+                success = bucket.put(item)
+                thread_results.append((timestamp, success))
+
+            with results_lock:
+                results.extend(thread_results)
+
+        threads = [
+            threading.Thread(target=worker, args=(i,)) for i in range(num_threads)
+        ]
+
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        successful_timestamps = sorted([ts for ts, success in results if success])
+
+        # Check sliding windows for both rates
+        for ts in successful_timestamps:
+            # 1-second sliding window
+            count_1s = sum(1 for t in successful_timestamps if ts - 1000 <= t <= ts)
+            assert count_1s <= 5, f"1-second rate exceeded: {count_1s} items in window ending at {ts}"
+
+            # 500ms sliding window
+            count_500ms = sum(1 for t in successful_timestamps if ts - 500 <= t <= ts)
+            assert count_500ms <= 3, f"500ms rate exceeded: {count_500ms} items in window ending at {ts}"
+
+    def test_concurrent_put_weighted(self, pg_pool, clean_table):
+        rate_limit = 10
+        rates = [Rate(rate_limit, Duration.SECOND)]
+        num_threads = 4
+        weight = 3
+
+        results = []
+        results_lock = threading.Lock()
+
+        def worker(thread_id: int):
+            bucket = PostgresBucket(pg_pool, clean_table, rates)
+            thread_results = []
+
+            for _ in range(5):
+                timestamp = bucket.now()
+                item = RateItem(f"thread_{thread_id}", timestamp, weight=weight)
+                success = bucket.put(item)
+                thread_results.append((timestamp, success, weight))
+
+            with results_lock:
+                results.extend(thread_results)
+
+        threads = [
+            threading.Thread(target=worker, args=(i,)) for i in range(num_threads)
+        ]
+
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        successful_results = sorted([(ts, w) for ts, success, w in results if success])
+
+        # Check sliding windows for weighted items
+        for ts, _ in successful_results:
+            weight_in_window = sum(w for t, w in successful_results if ts - 1000 <= t <= ts)
+            assert weight_in_window <= rate_limit, (
+                f"Rate limit exceeded: weight {weight_in_window} in 1-second window ending at {ts}"
+            )
