Coverage for agentos/tools/request_deduplicator.py: 0%
142 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 21:19 +0800
1"""
2RequestDeduplicator — fingerprint-based concurrent request deduplication.
4Supports:
5 - Fingerprint generation from request parameters
6 - In-flight deduplication (same fingerprint → wait for existing result)
7 - Result caching with TTL (return cached result on duplicate)
8 - Thread-safe + async-safe
9 - Auto-cleanup of expired entries
10 - Configurable max cache size
11"""
13from __future__ import annotations
15import hashlib
16import json
17import threading
18import time
19from collections.abc import Callable
20from enum import Enum
21from typing import Any
23# ============================================================================
24# Result
25# ============================================================================
28class ResultStatus(Enum):
29 COMPLETED = "completed"
30 ERROR = "error"
33class DedupResult:
34 __slots__ = ("status", "value", "timestamp")
36 def __init__(self, status: ResultStatus, value: Any):
37 self.status = status
38 self.value = value
39 self.timestamp = time.time()
42# ============================================================================
43# RequestDeduplicator
44# ============================================================================
47class RequestDeduplicator:
48 """Fingerprint-based request deduplication with result caching.
50 Usage:
51 dedup = RequestDeduplicator(ttl=30.0)
53 # Option A: manual
54 key = dedup.create_key(method="POST", path="/api/users", body={"name": "Alice"})
55 result = dedup.get(key)
56 if result:
57 return result.value
59 dedup.mark_in_flight(key)
60 try:
61 response = do_request(...)
62 dedup.complete(key, response)
63 except Exception as e:
64 dedup.error(key, e)
65 raise
67 # Option B: decorator
68 @dedup.deduplicate(key_fn=lambda *a, **kw: f"{a[0]}_{a[1]}")
69 def fetch(user_id, query):
70 return api_call(user_id, query)
71 """
73 def __init__(
74 self,
75 ttl: float = 60.0,
76 max_entries: int = 10000,
77 key_prefix: str = "dedup:",
78 ):
79 self._ttl = ttl
80 self._max_entries = max_entries
81 self._key_prefix = key_prefix
82 self._cache: dict[str, DedupResult] = {}
83 self._in_flight: dict[str, threading.Event] = {}
84 self._in_flight_results: dict[str, DedupResult] = {}
85 self._lock = threading.RLock()
86 self._last_cleanup = time.time()
88 # ---------- key generation ----------
90 def create_key(self, *args: Any, **kwargs: Any) -> str:
91 """Generate a unique fingerprint key from args/kwargs.
93 Args are hashed positionally; kwargs are sorted by key.
94 """
95 payload: dict[str, Any] = {"args": args, "kwargs": dict(sorted(kwargs.items()))}
96 raw = json.dumps(payload, sort_keys=True, default=str)
97 digest = hashlib.sha256(raw.encode()).hexdigest()[:16]
98 return f"{self._key_prefix}{digest}"
100 # ---------- lookup ----------
102 def get(self, key: str) -> Any | None:
103 """Return cached result if available and not expired. None if not found."""
104 self._maybe_cleanup()
105 with self._lock:
106 entry = self._cache.get(key)
107 if entry is None:
108 return None
109 age = time.time() - entry.timestamp
110 if age > self._ttl:
111 del self._cache[key]
112 return None
113 return entry
115 def get_or_none(self, key: str) -> Any | None:
116 """Same as get() but returns the raw value or None."""
117 entry = self.get(key)
118 if entry:
119 return entry.value
120 return None
122 # ---------- in-flight management ----------
124 def mark_in_flight(self, key: str) -> bool:
125 """Mark key as in-flight. Returns True if we should proceed (first caller).
126 Returns False if another caller is already processing — caller should wait.
127 """
128 with self._lock:
129 if key in self._in_flight:
130 return False
131 self._in_flight[key] = threading.Event()
132 return True
134 def wait_in_flight(self, key: str, timeout: float | None = None) -> Any | None:
135 """Wait for an in-flight request to complete, then return its result."""
136 event = None
137 with self._lock:
138 event = self._in_flight.get(key)
139 if event is None:
140 return None
141 signaled = event.wait(timeout=timeout)
142 if not signaled:
143 return None
144 with self._lock:
145 result = self._in_flight_results.pop(key, None)
146 self._in_flight.pop(key, None)
147 if result:
148 return result.value
149 return None
151 def complete(self, key: str, result: Any) -> None:
152 """Signal completion and cache the result."""
153 with self._lock:
154 entry = DedupResult(ResultStatus.COMPLETED, result)
155 self._cache[key] = entry
156 self._in_flight_results[key] = entry
157 event = self._in_flight.get(key)
158 # Signal outside lock to avoid deadlock
159 if event:
160 event.set()
161 self._evict_if_needed()
163 def error(self, key: str, error: Exception) -> None:
164 """Signal error for in-flight request."""
165 with self._lock:
166 entry = DedupResult(ResultStatus.ERROR, error)
167 self._in_flight_results[key] = entry
168 event = self._in_flight.get(key)
169 if event:
170 event.set()
172 # ---------- decorator ----------
174 def deduplicate(
175 self,
176 key_fn: Callable[..., str],
177 wait_timeout: float | None = 30.0,
178 cache_errors: bool = False,
179 ):
180 """Decorator: deduplicate concurrent calls with same fingerprint.
182 Args:
183 key_fn: function(*args, **kwargs) → key string
184 wait_timeout: max wait for in-flight request
185 cache_errors: if True, cache error results too
186 """
188 def decorator(func: Callable) -> Callable:
189 def wrapper(*args: Any, **kwargs: Any) -> Any:
190 key = key_fn(*args, **kwargs)
192 # Check cache first
193 cached = self.get(key)
194 if cached is not None:
195 if cached.status == ResultStatus.ERROR:
196 if not cache_errors:
197 pass # fall through to re-execute
198 else:
199 raise (
200 cached.value
201 if isinstance(cached.value, Exception)
202 else Exception(str(cached.value))
203 )
204 else:
205 return cached.value
207 # Try to claim in-flight
208 if self.mark_in_flight(key):
209 try:
210 result = func(*args, **kwargs)
211 self.complete(key, result)
212 return result
213 except Exception as e:
214 if cache_errors:
215 self.complete(key, e)
216 else:
217 self.error(key, e)
218 raise
219 else:
220 # Another caller is processing — wait
221 result = self.wait_in_flight(key, timeout=wait_timeout)
222 if result is not None:
223 return result
224 # Timeout: fall through to execute ourselves
225 raise TimeoutError(f"Timeout waiting for deduplicated request: {key}")
227 return wrapper
229 return decorator
231 # ---------- cache maintenance ----------
233 def _maybe_cleanup(self) -> None:
234 """Trigger cleanup if enough time has passed."""
235 now = time.time()
236 if now - self._last_cleanup < self._ttl:
237 return
238 self._last_cleanup = now
239 with self._lock:
240 expired = [k for k, v in self._cache.items() if now - v.timestamp > self._ttl]
241 for k in expired:
242 del self._cache[k]
244 def _evict_if_needed(self) -> None:
245 with self._lock:
246 excess = len(self._cache) - self._max_entries
247 if excess <= 0:
248 return
249 # Evict oldest entries
250 sorted_by_age = sorted(self._cache.items(), key=lambda x: x[1].timestamp)
251 for k, _ in sorted_by_age[:excess]:
252 del self._cache[k]
254 def clear(self) -> None:
255 with self._lock:
256 self._cache.clear()
257 self._in_flight.clear()
258 self._in_flight_results.clear()
260 @property
261 def cache_size(self) -> int:
262 with self._lock:
263 return len(self._cache)
265 @property
266 def in_flight_count(self) -> int:
267 with self._lock:
268 return len(self._in_flight)