Coverage for ivatar/access_stats.py: 77%
124 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 11:51 +0000
1import time
2import threading
3import atexit
4import logging
5from typing import Set
6from django.core.cache import cache
7from django.db import models
8from django.db.models import F
9from django.conf import settings
10from django.apps import apps
12logger = logging.getLogger("ivatar.access_stats")
14# Maps parent model class name to (stat model class name, FK field name used in filter).
15# The stat model stores the access_count and has a OneToOneField (as PK) to the parent.
16MODEL_TO_STAT: dict[str, tuple[str, str]] = {
17 "Photo": ("PhotoAccessStat", "photo_id"),
18 "ConfirmedEmail": ("ConfirmedEmailAccessStat", "email_id"),
19 "ConfirmedOpenId": ("ConfirmedOpenIdAccessStat", "openid_id"),
20}
23class AccessStatsManager:
24 def __init__(self):
25 self._local_dirty_keys: Set[str] = set()
26 self._lock = threading.Lock()
27 atexit.register(self.flush_all_dirty_keys)
29 def _get_flush_timeout(self) -> int:
30 return getattr(settings, "STATS_FLUSH_TIMEOUT", 300)
32 def _get_batch_size(self) -> int:
33 return getattr(settings, "STATS_BATCH_SIZE", 100)
35 def _flush_key_sync(self, key: str) -> None:
36 """
37 Synchronously flush a specific cache key to the database.
38 Uses a distributed lock to prevent race conditions.
39 """
40 lock_key = f"{key}:lock"
41 # Acquire lock (expire in 60s to prevent deadlocks)
42 # cache.add returns True if the key was added (lock acquired), False otherwise
43 if not cache.add(lock_key, "locked", 60):
44 logger.debug("Could not acquire lock for %s, skipping flush", key)
45 return
47 try:
48 # key format: "stats:access_count:ModelName:pk"
49 parts = key.split(":")
50 if len(parts) != 4:
51 return
53 model_name = parts[2]
54 pk = parts[3]
56 # Get the cached value
57 try:
58 count = cache.get(key)
59 if not count:
60 return
61 count = int(count)
62 except (ValueError, TypeError):
63 return
65 if count == 0:
66 return
68 # Resolve the stat model class
69 stat_info = MODEL_TO_STAT.get(model_name)
70 if stat_info is None:
71 logger.warning("No stat model registered for %s", model_name)
72 return
74 stat_model_name, fk_field = stat_info
75 try:
76 stat_class = apps.get_model("ivataraccount", stat_model_name)
77 except LookupError:
78 logger.warning("Could not find stat model %s", stat_model_name)
79 return
81 logger.info(
82 "Flushing %d access counts for %s:%s to database", count, model_name, pk
83 )
85 # Update the stat table. Use get_or_create to handle the rare case where
86 # the stat row does not yet exist (e.g. objects created before the migration).
87 rows_updated = stat_class.objects.filter(**{fk_field: pk}).update(
88 access_count=F("access_count") + count
89 )
90 if rows_updated == 0:
91 # Stat row missing — create it with the accumulated count.
92 stat_class.objects.get_or_create(
93 **{fk_field: pk},
94 defaults={"access_count": count},
95 )
97 # Decrement cache
98 try:
99 cache.decr(key, count)
100 except ValueError:
101 cache.set(key, 0)
103 # Clear start time
104 start_time_key = f"stats:start_time:{model_name}:{pk}"
105 cache.delete(start_time_key)
107 except Exception as e:
108 logger.error("Error flushing key %s: %s", key, e)
109 finally:
110 # Release lock
111 cache.delete(lock_key)
113 def _flush_worker(self, key: str) -> None:
114 """
115 Worker function to flush key and remove from dirty set.
116 """
117 self._flush_key_sync(key)
118 with self._lock:
119 if key in self._local_dirty_keys:
120 self._local_dirty_keys.remove(key)
122 def flush_all_dirty_keys(self) -> None:
123 """
124 Flush all known dirty keys. Registered with atexit.
125 """
126 with self._lock:
127 keys_to_flush = list(self._local_dirty_keys)
128 self._local_dirty_keys.clear()
130 for key in keys_to_flush:
131 self._flush_key_sync(key)
133 def _update_stat_direct(self, obj: models.Model, delta: int = 1) -> None:
134 """
135 Update the stat table directly (used for sync mode and cache-failure fallback).
136 """
137 model_name = obj.__class__.__name__
138 stat_info = MODEL_TO_STAT.get(model_name)
139 if stat_info is None:
140 logger.warning("No stat model registered for %s", model_name)
141 return
142 stat_model_name, fk_field = stat_info
143 try:
144 stat_class = apps.get_model("ivataraccount", stat_model_name)
145 except LookupError:
146 logger.warning("Could not find stat model %s", stat_model_name)
147 return
148 rows_updated = stat_class.objects.filter(**{fk_field: obj.pk}).update(
149 access_count=F("access_count") + delta
150 )
151 if rows_updated == 0:
152 stat_class.objects.get_or_create(
153 **{fk_field: obj.pk},
154 defaults={"access_count": delta},
155 )
157 def update_access_count(self, obj: models.Model) -> None:
158 """
159 Update access count for an object (Photo, ConfirmedEmail, ConfirmedOpenId).
160 If ASYNC_ACCESS_COUNT is True, uses Memcached to batch updates.
161 Otherwise, updates the stat table directly.
162 """
163 if not getattr(settings, "ASYNC_ACCESS_COUNT", True):
164 # Fallback to synchronous/direct update on the stat table
165 self._update_stat_direct(obj)
166 return
168 # Key format: stats:access_count:<model_name>:<pk>
169 model_name = obj.__class__.__name__
170 pk = obj.pk
171 key = f"stats:access_count:{model_name}:{pk}"
172 start_time_key = f"stats:start_time:{model_name}:{pk}"
174 # Track as dirty
175 with self._lock:
176 self._local_dirty_keys.add(key)
178 try:
179 # Atomic increment.
180 try:
181 new_value = cache.incr(key)
182 except ValueError:
183 # Key didn't exist, set it to 1
184 cache.set(key, 1)
185 new_value = 1
186 # Set start time for timeout flush
187 cache.set(start_time_key, time.time(), timeout=None)
189 # Check triggers
190 batch_size = self._get_batch_size()
191 should_flush = False
193 if new_value >= batch_size:
194 should_flush = True
195 else:
196 # Check time-based trigger
197 start_time = cache.get(start_time_key)
198 if start_time:
199 try:
200 if time.time() - float(start_time) > self._get_flush_timeout():
201 should_flush = True
202 except (ValueError, TypeError):
203 pass
205 if should_flush:
206 # Run flush in a separate thread to avoid blocking the request
207 threading.Thread(target=self._flush_worker, args=(key,)).start()
209 except Exception as e:
210 # Fallback: update stat table directly if cache fails
211 logger.warning("Failed to update access count cache: %s", e)
212 self._update_stat_direct(obj)
215stats_manager = AccessStatsManager()