diff --git a/packages/smooth-api-py/smooth_api/state.py b/packages/smooth-api-py/smooth_api/state.py index d1026d9..d2e0c46 100644 --- a/packages/smooth-api-py/smooth_api/state.py +++ b/packages/smooth-api-py/smooth_api/state.py @@ -1,5 +1,6 @@ from __future__ import annotations +import abc import asyncio import inspect import sys @@ -15,19 +16,185 @@ class CircuitEntry: state: CircuitState failure_count: int - last_failure_time: float #epoch seconds (not ms) + last_failure_time: float # epoch seconds (not ms) + + +class CircuitBreakerStorage(abc.ABC): + @abc.abstractmethod + async def get_state(self, domain: str) -> CircuitEntry: + """Retrieve current circuit entry state.""" + ... + + @abc.abstractmethod + async def record_failure( + self, domain: str, threshold: int, cooldown_ms: int + ) -> CircuitState: + """Record failure and return new circuit state.""" + ... + + @abc.abstractmethod + async def record_success(self, domain: str) -> CircuitState: + """Record success and return new circuit state.""" + ... + + +class InMemoryStorage(CircuitBreakerStorage): + def __init__(self) -> None: + self._map: dict[str, CircuitEntry] = {} + self._lock = threading.Lock() + + def _cleanup(self) -> None: + if len(self._map) <= 1000: + return + keys_to_delete = [ + k for k, v in self._map.items() + if v.state == "CLOSED" and v.failure_count == 0 + ] + for k in keys_to_delete: + del self._map[k] + + def _get_or_create(self, domain: str) -> CircuitEntry: + if domain not in self._map: + self._cleanup() + self._map[domain] = CircuitEntry( + state="CLOSED", + failure_count=0, + last_failure_time=0.0, + ) + return self._map[domain] + + async def get_state(self, domain: str) -> CircuitEntry: + with self._lock: + entry = self._get_or_create(domain) + return CircuitEntry( + state=entry.state, + failure_count=entry.failure_count, + last_failure_time=entry.last_failure_time, + ) + + async def record_failure( + self, domain: str, threshold: int, cooldown_ms: int + ) -> CircuitState: + with self._lock: + entry = self._get_or_create(domain) + entry.failure_count += 1 + if entry.state == "HALF_OPEN" or entry.failure_count >= threshold: + entry.state = "OPEN" + entry.last_failure_time = time.time() + return entry.state + + async def record_success(self, domain: str) -> CircuitState: + with self._lock: + entry = self._get_or_create(domain) + entry.failure_count = 0 + entry.state = "CLOSED" + return "CLOSED" + + +class RedisCircuitBreakerStorage(CircuitBreakerStorage): + def __init__(self, redis_client: Any, key_prefix: str = "smoothapi:cb:") -> None: + self._redis = redis_client + self._key_prefix = key_prefix + self._fallback = InMemoryStorage() + + def _get_key(self, domain: str) -> str: + return f"{self._key_prefix}{domain}" + + async def get_state(self, domain: str) -> CircuitEntry: + key = self._get_key(domain) + lua = """ + local state = redis.call('HGET', KEYS[1], 'state') + local failureCount = redis.call('HGET', KEYS[1], 'failure_count') + local lastFailureTime = redis.call('HGET', KEYS[1], 'last_failure_time') + return { + state or 'CLOSED', + failureCount or '0', + lastFailureTime or '0' + } + """ + try: + if inspect.iscoroutinefunction(getattr(self._redis, "eval", None)): + res = await self._redis.eval(lua, 1, key) + else: + res = self._redis.eval(lua, 1, key) + state = str(res[0]) if res and res[0] else "CLOSED" + failure_count = int(res[1]) if res and len(res) > 1 else 0 + last_failure_time = float(res[2]) if res and len(res) > 2 else 0.0 + return CircuitEntry( + state=state, # type: ignore[arg-type] + failure_count=failure_count, + last_failure_time=last_failure_time, + ) + except Exception: + return await self._fallback.get_state(domain) + + async def record_failure( + self, domain: str, threshold: int, cooldown_ms: int + ) -> CircuitState: + key = self._get_key(domain) + now = time.time() + lua = """ + local state = redis.call('HGET', KEYS[1], 'state') or 'CLOSED' + local count = tonumber(redis.call('HINCRBY', KEYS[1], 'failure_count', 1)) + local now = tonumber(ARGV[1]) + local threshold = tonumber(ARGV[2]) + local cooldownMs = tonumber(ARGV[3]) + + local newState = state + if state == 'HALF_OPEN' or count >= threshold then + newState = 'OPEN' + redis.call('HSET', KEYS[1], 'state', 'OPEN') + redis.call('HSET', KEYS[1], 'last_failure_time', tostring(now)) + local ttlSec = math.ceil((cooldownMs * 2) / 1000) + if ttlSec > 0 then + redis.call('EXPIRE', KEYS[1], ttlSec) + end + end + return newState + """ + try: + if inspect.iscoroutinefunction(getattr(self._redis, "eval", None)): + result = await self._redis.eval(lua, 1, key, now, threshold, cooldown_ms) + else: + result = self._redis.eval(lua, 1, key, now, threshold, cooldown_ms) + return str(result) # type: ignore[return-value] + except Exception: + return await self._fallback.record_failure(domain, threshold, cooldown_ms) + + async def record_success(self, domain: str) -> CircuitState: + key = self._get_key(domain) + lua = """ + redis.call('HSET', KEYS[1], 'state', 'CLOSED') + redis.call('HSET', KEYS[1], 'failure_count', '0') + return 'CLOSED' + """ + try: + if inspect.iscoroutinefunction(getattr(self._redis, "eval", None)): + result = await self._redis.eval(lua, 1, key) + else: + result = self._redis.eval(lua, 1, key) + return str(result) # type: ignore[return-value] + except Exception: + return await self._fallback.record_success(domain) + class CircuitBreakerState: def __init__( self, config: CircuitBreakerConfig | None = None, on_state_change: Callable[[CircuitStateChangeEvent], Any] | None = None, + storage: CircuitBreakerStorage | None = None, ): self._config = config or CircuitBreakerConfig() self._on_state_change = on_state_change + self._storage = storage or InMemoryStorage() self._map: dict[str, CircuitEntry] = {} self._lock = threading.Lock() + @property + def storage(self) -> CircuitBreakerStorage: + return self._storage + def _transition(self, domain: str, entry: CircuitEntry, to: CircuitState) -> CircuitStateChangeEvent | None: from_state = entry.state if from_state != to: @@ -58,7 +225,6 @@ def _handle_task_done(t: asyncio.Task) -> None: except Exception as ex: sys.stderr.write(f"[smoothAPI] Error in onCircuitStateChange hook: {ex}\n") - # No lock needed here def _get_or_create(self, domain: str) -> CircuitEntry: if domain not in self._map: self._cleanup() @@ -70,7 +236,6 @@ def _get_or_create(self, domain: str) -> CircuitEntry: return self._map[domain] def _cleanup(self) -> None: - # Enforce a strict domain limit to bound memory usage if len(self._map) <= 1000: return keys_to_delete = [ @@ -80,13 +245,43 @@ def _cleanup(self) -> None: for k in keys_to_delete: del self._map[k] + def _schedule_async(self, coro: Any) -> None: + try: + loop = asyncio.get_running_loop() + loop.create_task(coro) + except RuntimeError: + try: + asyncio.run(coro) + except Exception: + pass + def can_request(self, domain: str) -> bool: event = None with self._lock: entry = self._get_or_create(domain) if entry.state == 'CLOSED' or entry.state == 'HALF_OPEN': return True - # cooldown_ms is in ms; time.time() is in seconds + elapsed = time.time() - entry.last_failure_time + if elapsed > self._config.cooldown_ms / 1000: + event = self._transition(domain, entry, 'HALF_OPEN') + allowed = True + else: + allowed = False + if event: + self._notify_state_change(event) + return allowed + + async def can_request_async(self, domain: str) -> bool: + remote_entry = await self._storage.get_state(domain) + event = None + with self._lock: + entry = self._get_or_create(domain) + entry.state = remote_entry.state + entry.failure_count = remote_entry.failure_count + entry.last_failure_time = remote_entry.last_failure_time + + if entry.state == 'CLOSED' or entry.state == 'HALF_OPEN': + return True elapsed = time.time() - entry.last_failure_time if elapsed > self._config.cooldown_ms / 1000: event = self._transition(domain, entry, 'HALF_OPEN') @@ -105,6 +300,17 @@ def record_success(self, domain: str) -> None: event = self._transition(domain, entry, 'CLOSED') if event: self._notify_state_change(event) + self._schedule_async(self._storage.record_success(domain)) + + async def record_success_async(self, domain: str) -> None: + new_state = await self._storage.record_success(domain) + event = None + with self._lock: + entry = self._get_or_create(domain) + entry.failure_count = 0 + event = self._transition(domain, entry, new_state) + if event: + self._notify_state_change(event) def record_failure(self, domain: str) -> None: event = None @@ -119,7 +325,34 @@ def record_failure(self, domain: str) -> None: entry.last_failure_time = time.time() if event: self._notify_state_change(event) + self._schedule_async( + self._storage.record_failure( + domain, self._config.failure_threshold, self._config.cooldown_ms + ) + ) + + async def record_failure_async(self, domain: str) -> None: + new_state = await self._storage.record_failure( + domain, self._config.failure_threshold, self._config.cooldown_ms + ) + event = None + with self._lock: + entry = self._get_or_create(domain) + entry.failure_count += 1 + entry.last_failure_time = time.time() + event = self._transition(domain, entry, new_state) + if event: + self._notify_state_change(event) def get_state(self, domain: str) -> CircuitState: with self._lock: - return self._get_or_create(domain).state \ No newline at end of file + return self._get_or_create(domain).state + + async def get_state_async(self, domain: str) -> CircuitState: + remote_entry = await self._storage.get_state(domain) + with self._lock: + entry = self._get_or_create(domain) + entry.state = remote_entry.state + entry.failure_count = remote_entry.failure_count + entry.last_failure_time = remote_entry.last_failure_time + return entry.state diff --git a/packages/smooth-api-ts/src/state.ts b/packages/smooth-api-ts/src/state.ts index 6755445..819a416 100644 --- a/packages/smooth-api-ts/src/state.ts +++ b/packages/smooth-api-ts/src/state.ts @@ -1,20 +1,176 @@ -import type { CircuitState, CircuitBreakerConfig, CircuitEntry, CircuitStateChangeEvent } from "./types.js"; +import type { CircuitState, CircuitBreakerConfig, CircuitEntry, CircuitStateChangeEvent } from "./types.js"; const DEFAULTS: Required = { failureThreshold: 3, cooldownMs: 10_000, +}; + +export interface CircuitBreakerStorage { + getState(domain: string): Promise; + recordFailure(domain: string, threshold: number, cooldownMs: number): Promise; + recordSuccess(domain: string): Promise; +} + +export class InMemoryStorage implements CircuitBreakerStorage { + private readonly map = new Map(); + + private cleanup(): void { + if (this.map.size <= 1000) return; + for (const [key, entry] of this.map.entries()) { + if (entry.state === "CLOSED" && entry.failureCount === 0) { + this.map.delete(key); + } + } + } + + private getOrCreate(domain: string): CircuitEntry { + if (!this.map.has(domain)) { + this.cleanup(); + this.map.set(domain, { + state: "CLOSED", + failureCount: 0, + lastFailureTime: 0, + }); + } + return this.map.get(domain)!; + } + + async getState(domain: string): Promise { + const entry = this.getOrCreate(domain); + return { ...entry }; + } + + async recordFailure(domain: string, threshold: number, cooldownMs: number): Promise { + const entry = this.getOrCreate(domain); + entry.failureCount++; + const now = Date.now(); + if (entry.state === "HALF_OPEN" || entry.failureCount >= threshold) { + entry.state = "OPEN"; + entry.lastFailureTime = now; + } + return entry.state; + } + + async recordSuccess(domain: string): Promise { + const entry = this.getOrCreate(domain); + entry.failureCount = 0; + entry.state = "CLOSED"; + return "CLOSED"; + } + + // Direct synchronous access for local fast-path + getEntrySync(domain: string): CircuitEntry { + return this.getOrCreate(domain); + } +} + +export interface RedisClientLike { + eval(script: string, numkeys: number, ...args: (string | number)[]): Promise; + get?(key: string): Promise; +} + +export class RedisCircuitBreakerStorage implements CircuitBreakerStorage { + private readonly fallback: InMemoryStorage = new InMemoryStorage(); + private readonly keyPrefix: string; + + constructor( + private readonly redis: RedisClientLike, + keyPrefix = "smoothapi:cb:" + ) { + this.keyPrefix = keyPrefix; + } + + private getKey(domain: string): string { + return `${this.keyPrefix}${domain}`; + } + + async getState(domain: string): Promise { + const key = this.getKey(domain); + const script = ` + local state = redis.call('HGET', KEYS[1], 'state') + local failureCount = redis.call('HGET', KEYS[1], 'failure_count') + local lastFailureTime = redis.call('HGET', KEYS[1], 'last_failure_time') + return { + state or 'CLOSED', + failureCount or '0', + lastFailureTime or '0' + } + `; + try { + const res = (await this.redis.eval(script, 1, key)) as [string, string, string]; + return { + state: (res[0] as CircuitState) || "CLOSED", + failureCount: Number.parseInt(res[1] || "0", 10), + lastFailureTime: Number.parseFloat(res[2] || "0"), + }; + } catch { + return this.fallback.getState(domain); + } + } + + async recordFailure(domain: string, threshold: number, cooldownMs: number): Promise { + const key = this.getKey(domain); + const now = Date.now(); + const script = ` + local state = redis.call('HGET', KEYS[1], 'state') or 'CLOSED' + local count = tonumber(redis.call('HINCRBY', KEYS[1], 'failure_count', 1)) + local now = tonumber(ARGV[1]) + local threshold = tonumber(ARGV[2]) + local cooldownMs = tonumber(ARGV[3]) + + local newState = state + if state == 'HALF_OPEN' or count >= threshold then + newState = 'OPEN' + redis.call('HSET', KEYS[1], 'state', 'OPEN') + redis.call('HSET', KEYS[1], 'last_failure_time', tostring(now)) + local ttlSec = math.ceil((cooldownMs * 2) / 1000) + if ttlSec > 0 then + redis.call('EXPIRE', KEYS[1], ttlSec) + end + end + return newState + `; + try { + const result = (await this.redis.eval(script, 1, key, now, threshold, cooldownMs)) as string; + return result as CircuitState; + } catch { + return this.fallback.recordFailure(domain, threshold, cooldownMs); + } + } + + async recordSuccess(domain: string): Promise { + const key = this.getKey(domain); + const script = ` + redis.call('HSET', KEYS[1], 'state', 'CLOSED') + redis.call('HSET', KEYS[1], 'failure_count', '0') + return 'CLOSED' + `; + try { + const result = (await this.redis.eval(script, 1, key)) as string; + return result as CircuitState; + } catch { + return this.fallback.recordSuccess(domain); + } + } } export class CircuitBreakerState { private readonly config: Required; private readonly map: Map; + private readonly storage: CircuitBreakerStorage; constructor( config: Partial = {}, - private readonly onStateChange?: (event: CircuitStateChangeEvent) => void + private readonly onStateChange?: (event: CircuitStateChangeEvent) => void, + storage?: CircuitBreakerStorage ) { - this.config = { ...DEFAULTS, ...config}; + this.config = { ...DEFAULTS, ...config }; this.map = new Map(); + this.storage = storage ?? new InMemoryStorage(); + } + + public getStorage(): CircuitBreakerStorage { + return this.storage; } private transition(domain: string, entry: CircuitEntry, to: CircuitState): void { @@ -33,57 +189,101 @@ export class CircuitBreakerState { } private getOrCreate(domain: string): CircuitEntry { - if (!this.map.has(domain)) { - this.cleanup(); - this.map.set(domain, { - state: 'CLOSED', - failureCount: 0, - lastFailureTime: 0, - }); - } - return this.map.get(domain)!; + if (!this.map.has(domain)) { + this.cleanup(); + this.map.set(domain, { + state: "CLOSED", + failureCount: 0, + lastFailureTime: 0, + }); + } + return this.map.get(domain)!; } canRequest(domain: string): boolean { const entry = this.getOrCreate(domain); - if(entry.state === "CLOSED" || entry.state === "HALF_OPEN") + if (entry.state === "CLOSED" || entry.state === "HALF_OPEN") { return true; + } const elapsedTime = Date.now() - entry.lastFailureTime; - if(elapsedTime > this.config.cooldownMs){ + if (elapsedTime > this.config.cooldownMs) { this.transition(domain, entry, "HALF_OPEN"); return true; } - else{ - return false; - } + return false; + } + + async canRequestAsync(domain: string): Promise { + const remoteEntry = await this.storage.getState(domain); + const entry = this.getOrCreate(domain); + entry.state = remoteEntry.state; + entry.failureCount = remoteEntry.failureCount; + entry.lastFailureTime = remoteEntry.lastFailureTime; + + if (entry.state === "CLOSED" || entry.state === "HALF_OPEN") { + return true; + } + const elapsedTime = Date.now() - entry.lastFailureTime; + if (elapsedTime > this.config.cooldownMs) { + this.transition(domain, entry, "HALF_OPEN"); + return true; + } + return false; } recordSuccess(domain: string): void { const entry = this.getOrCreate(domain); entry.failureCount = 0; this.transition(domain, entry, "CLOSED"); + this.storage.recordSuccess(domain).catch(() => {}); + } + + async recordSuccessAsync(domain: string): Promise { + const entry = this.getOrCreate(domain); + const newState = await this.storage.recordSuccess(domain); + entry.failureCount = 0; + this.transition(domain, entry, newState); } recordFailure(domain: string): void { const entry = this.getOrCreate(domain); entry.failureCount++; - if(entry.state === "HALF_OPEN"){ + if (entry.state === "HALF_OPEN") { this.transition(domain, entry, "OPEN"); entry.lastFailureTime = Date.now(); - return - } - else if( entry.failureCount >= this.config.failureThreshold){ + } else if (entry.failureCount >= this.config.failureThreshold) { this.transition(domain, entry, "OPEN"); entry.lastFailureTime = Date.now(); } + this.storage.recordFailure(domain, this.config.failureThreshold, this.config.cooldownMs).catch(() => {}); + } + + async recordFailureAsync(domain: string): Promise { + const entry = this.getOrCreate(domain); + const newState = await this.storage.recordFailure( + domain, + this.config.failureThreshold, + this.config.cooldownMs + ); + entry.failureCount++; + entry.lastFailureTime = Date.now(); + this.transition(domain, entry, newState); } getState(domain: string): CircuitState { return this.getOrCreate(domain).state; } - + + async getStateAsync(domain: string): Promise { + const entry = await this.storage.getState(domain); + const local = this.getOrCreate(domain); + local.state = entry.state; + local.failureCount = entry.failureCount; + local.lastFailureTime = entry.lastFailureTime; + return entry.state; + } + private cleanup(): void { - // Enforce a strict domain limit to bound memory usage if (this.map.size <= 1000) return; for (const [key, entry] of this.map.entries()) { if (entry.state === "CLOSED" && entry.failureCount === 0) { @@ -92,4 +292,3 @@ export class CircuitBreakerState { } } } -