Coverage for ivatar/access_stats.py: 77%

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 

11 

12logger = logging.getLogger("ivatar.access_stats") 

13 

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} 

21 

22 

23class AccessStatsManager: 

24 def __init__(self): 

25 self._local_dirty_keys: Set[str] = set() 

26 # Worker-local start time per stat key. Avoids a synchronous 

27 # cache.get(start_time_key) on every single hit. 

28 self._local_start_times: dict = {} 

29 self._lock = threading.Lock() 

30 atexit.register(self.flush_all_dirty_keys) 

31 

32 def _get_flush_timeout(self) -> int: 

33 return getattr(settings, "STATS_FLUSH_TIMEOUT", 300) 

34 

35 def _get_batch_size(self) -> int: 

36 return getattr(settings, "STATS_BATCH_SIZE", 100) 

37 

38 def _flush_key_sync(self, key: str) -> None: 

39 """ 

40 Synchronously flush a specific cache key to the database. 

41 Uses a distributed lock to prevent race conditions. 

42 """ 

43 lock_key = f"{key}:lock" 

44 # Acquire lock (expire in 60s to prevent deadlocks) 

45 # cache.add returns True if the key was added (lock acquired), False otherwise 

46 if not cache.add(lock_key, "locked", 60): 

47 logger.debug("Could not acquire lock for %s, skipping flush", key) 

48 return 

49 

50 try: 

51 # key format: "stats:access_count:ModelName:pk" 

52 parts = key.split(":") 

53 if len(parts) != 4: 

54 return 

55 

56 model_name = parts[2] 

57 pk = parts[3] 

58 

59 # Get the cached value 

60 try: 

61 count = cache.get(key) 

62 if not count: 

63 return 

64 count = int(count) 

65 except (ValueError, TypeError): 

66 return 

67 

68 if count == 0: 

69 return 

70 

71 # Resolve the stat model class 

72 stat_info = MODEL_TO_STAT.get(model_name) 

73 if stat_info is None: 

74 logger.warning("No stat model registered for %s", model_name) 

75 return 

76 

77 stat_model_name, fk_field = stat_info 

78 try: 

79 stat_class = apps.get_model("ivataraccount", stat_model_name) 

80 except LookupError: 

81 logger.warning("Could not find stat model %s", stat_model_name) 

82 return 

83 

84 logger.info( 

85 "Flushing %d access counts for %s:%s to database", count, model_name, pk 

86 ) 

87 

88 # Update the stat table. Use get_or_create to handle the rare case where 

89 # the stat row does not yet exist (e.g. objects created before the migration). 

90 rows_updated = stat_class.objects.filter(**{fk_field: pk}).update( 

91 access_count=F("access_count") + count 

92 ) 

93 if rows_updated == 0: 

94 # Stat row missing — create it with the accumulated count. 

95 stat_class.objects.get_or_create( 

96 **{fk_field: pk}, 

97 defaults={"access_count": count}, 

98 ) 

99 

100 # Decrement cache 

101 try: 

102 cache.decr(key, count) 

103 except ValueError: 

104 cache.set(key, 0) 

105 

106 except Exception as e: 

107 logger.error("Error flushing key %s: %s", key, e) 

108 finally: 

109 # Release lock 

110 cache.delete(lock_key) 

111 

112 def _flush_worker(self, key: str) -> None: 

113 """ 

114 Worker function to flush key and remove from dirty set. 

115 """ 

116 self._flush_key_sync(key) 

117 with self._lock: 

118 if key in self._local_dirty_keys: 

119 self._local_dirty_keys.remove(key) 

120 self._local_start_times.pop(key, None) 

121 

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() 

129 self._local_start_times.clear() 

130 

131 for key in keys_to_flush: 

132 self._flush_key_sync(key) 

133 

134 def _update_stat_direct(self, obj: models.Model, delta: int = 1) -> None: 

135 """ 

136 Update the stat table directly (used for sync mode and cache-failure fallback). 

137 """ 

138 model_name = obj.__class__.__name__ 

139 stat_info = MODEL_TO_STAT.get(model_name) 

140 if stat_info is None: 

141 logger.warning("No stat model registered for %s", model_name) 

142 return 

143 stat_model_name, fk_field = stat_info 

144 try: 

145 stat_class = apps.get_model("ivataraccount", stat_model_name) 

146 except LookupError: 

147 logger.warning("Could not find stat model %s", stat_model_name) 

148 return 

149 rows_updated = stat_class.objects.filter(**{fk_field: obj.pk}).update( 

150 access_count=F("access_count") + delta 

151 ) 

152 if rows_updated == 0: 

153 stat_class.objects.get_or_create( 

154 **{fk_field: obj.pk}, 

155 defaults={"access_count": delta}, 

156 ) 

157 

158 def update_access_count(self, obj: models.Model) -> None: 

159 """ 

160 Update access count for an object (Photo, ConfirmedEmail, ConfirmedOpenId). 

161 If ASYNC_ACCESS_COUNT is True, uses Memcached to batch updates. 

162 Otherwise, updates the stat table directly. 

163 """ 

164 if not getattr(settings, "ASYNC_ACCESS_COUNT", True): 

165 # Fallback to synchronous/direct update on the stat table 

166 self._update_stat_direct(obj) 

167 return 

168 

169 # Key format: stats:access_count:<model_name>:<pk> 

170 model_name = obj.__class__.__name__ 

171 pk = obj.pk 

172 key = f"stats:access_count:{model_name}:{pk}" 

173 

174 # Track as dirty and record the (worker-local) first-seen time. 

175 # Bounded so high-cardinality traffic cannot grow memory without limit. 

176 max_dirty_keys = getattr(settings, "STATS_MAX_DIRTY_KEYS", 10000) 

177 with self._lock: 

178 if len(self._local_dirty_keys) >= max_dirty_keys: 

179 for old_key in list(self._local_dirty_keys)[:1000]: 

180 self._local_dirty_keys.discard(old_key) 

181 self._local_start_times.pop(old_key, None) 

182 self._local_dirty_keys.add(key) 

183 if key not in self._local_start_times: 

184 self._local_start_times[key] = time.time() 

185 

186 try: 

187 # Atomic increment. 

188 try: 

189 new_value = cache.incr(key) 

190 except ValueError: 

191 # Key didn't exist, set it to 1 

192 cache.set(key, 1) 

193 new_value = 1 

194 

195 # Check triggers 

196 batch_size = self._get_batch_size() 

197 should_flush = False 

198 

199 if new_value >= batch_size: 

200 should_flush = True 

201 else: 

202 # Time-based trigger, evaluated from worker-local state 

203 start_time = self._local_start_times.get(key) 

204 if start_time and ( 

205 time.time() - start_time > self._get_flush_timeout() 

206 ): 

207 should_flush = True 

208 

209 if should_flush: 

210 # Run flush in a separate thread to avoid blocking the request 

211 threading.Thread(target=self._flush_worker, args=(key,)).start() 

212 

213 except Exception as e: 

214 # Fallback: update stat table directly if cache fails 

215 logger.warning("Failed to update access count cache: %s", e) 

216 self._update_stat_direct(obj) 

217 

218 

219stats_manager = AccessStatsManager()