Coverage for apps / core / cache.py: 91%

1"""PostgreSQL-safe database cache backend. 

2 

3Django's built-in DatabaseCache has two concurrency issues on Postgres: 

4 

51. ``_base_set`` does SELECT-then-INSERT, which races into IntegrityError 

6 on concurrent inserts of the same key. 

72. ``incr`` / ``decr`` do SELECT-then-UPDATE on the raw row, which races 

8 into lost updates: N concurrent incrs on the same key can settle at 

9 fewer than N because each reads the same pre-value and each writes the 

10 same post-value. 

11 

12This backend wraps set in a savepoint + retry (fix #1) and wraps incr/decr 

13in a per-key advisory lock (fix #2) so quota counters and other 

14incr-based counters remain accurate under load. 

15""" 

16 

17import hashlib 

18 

19from django.core.cache.backends.db import DatabaseCache 

20from django.db import IntegrityError, connections, router, transaction 

21 

22 

23class PostgreSafeDatabaseCache(DatabaseCache): 

24 """DatabaseCache with Postgres-safe set/incr semantics.""" 

25 

26 def _base_set(self, mode, key, value, timeout=None): 

27 try: 

28 with transaction.atomic(): 

29 return super()._base_set(mode, key, value, timeout) 

30 except IntegrityError: 

31 # Another worker inserted the same key between our SELECT and INSERT. 

32 # The savepoint rolled back the failed INSERT, so the transaction is 

33 # clean. Retry — the second attempt will see the existing row and 

34 # UPDATE it. 

35 return super()._base_set(mode, key, value, timeout) 

36 

37 def _advisory_lock_id(self, full_key: str) -> int: 

38 """Stable non-negative 32-bit int from a cache key for pg_advisory_xact_lock. 

39 

40 blake2b with digest_size=4 gives exactly 4 bytes; non-cryptographic use 

41 (we just need a deterministic mapping from key-string to int32). 

42 """ 

43 digest = hashlib.blake2b(full_key.encode("utf-8"), digest_size=4).digest() 

44 return int.from_bytes(digest, "big", signed=False) & 0x7FFFFFFF 

45 

46 def incr(self, key, delta=1, version=None): 

47 """Atomic per-key increment using a transaction-scoped advisory lock. 

48 

49 Django's DatabaseCache.incr reads the row, computes the new value in 

50 Python, then UPDATEs — racey across connections. Advisory-lock 

51 serialisation makes concurrent incrs on the same key linearisable 

52 while allowing parallelism across different keys. 

53 """ 

54 full_key = self.make_key(key, version=version) 

55 lock_id = self._advisory_lock_id(full_key) 

56 db = router.db_for_write(self.cache_model_class) 

57 with transaction.atomic(using=db): 

58 with connections[db].cursor() as cursor: 

59 cursor.execute("SELECT pg_advisory_xact_lock(%s)", [lock_id]) 

60 return super().incr(key, delta=delta, version=version) 

61 

62 def decr(self, key, delta=1, version=None): 

63 """Atomic decrement — same story as incr.""" 

64 return self.incr(key, delta=-delta, version=version)