metered 1.0.0__py3-none-any.whl

metered/__init__.py

@@ -0,0 +1,17 @@
+from .core.limiter import Metered
+from .core.types import Limit, QuotaConfig, QuotaRule, LimitResult, Strategy, IdentifierType
+from .backends.memory import InMemoryBackend
+from .backends.redis import RedisBackend
+
+# Explicitly export the public API
+__all__ = [
+    "Metered",
+    "Limit",
+    "QuotaConfig",
+    "QuotaRule",
+    "LimitResult",
+    "Strategy",
+    "IdentifierType",
+    "InMemoryBackend",
+    "RedisBackend",
+]

metered/analytics/__init__.py

@@ -0,0 +1,4 @@
+from .models import AnalyticsConfig, UsageSummary
+from .tracker import UsageTracker
+
+__all__ = ["AnalyticsConfig", "UsageSummary", "UsageTracker"]

metered/analytics/models.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Optional
+
+@dataclass
+class AnalyticsConfig:
+    retention_days: int = 30
+    anonymize_ip: bool = True
+    granularity: str = "hourly" # or "daily"
+
+@dataclass
+class UsageSummary:
+    target: str
+    total_cost: int
+    period_start: float
+    period_end: float

metered/analytics/tracker.py

@@ -0,0 +1,30 @@
+from typing import Optional, Any
+import time
+from datetime import datetime
+from .models import AnalyticsConfig
+from ..core.backends import BaseBackend
+
+class UsageTracker:
+    def __init__(self, backend: BaseBackend, config: Optional[AnalyticsConfig] = None):
+        self.backend = backend
+        self.config = config or AnalyticsConfig()
+
+    async def track(self, target: str, cost: int, timestamp: Optional[float] = None) -> None:
+        if timestamp is None:
+            timestamp = time.time()
+            
+        dt = datetime.fromtimestamp(timestamp)
+        
+        if self.config.granularity == "hourly":
+            time_key = dt.strftime("%Y-%m-%d:%H")
+        else:
+            time_key = dt.strftime("%Y-%m-%d")
+            
+        key = f"metered:usage:{target}:{time_key}"
+        
+        if hasattr(self.backend, "increment_usage"):
+            await self.backend.increment_usage(key, cost, self.config.retention_days) # type: ignore
+
+    async def get_usage(self, target: str, start_ts: float, end_ts: float) -> int:
+        """Helper to query usage if backend supports it"""
+        return 0

metered/backends/memory.py

@@ -0,0 +1,174 @@
+import time
+import asyncio
+from typing import Dict, List, Tuple, Optional
+from collections import defaultdict
+from metered.core.backends import BaseBackend
+from metered.core.types import LimitResult, Strategy
+
+class InMemoryBackend(BaseBackend):
+    """In-memory backend with per-key locking and auto-cleanup."""
+    
+    def __init__(self, cleanup_interval: float = 300.0) -> None:
+        self._fixed_window: Dict[str, Tuple[int, float]] = {}
+        self._token_bucket: Dict[str, Tuple[float, float]] = {}
+        self._sliding_window: Dict[str, List[float]] = {}
+        self._leaky_bucket: Dict[str, Tuple[float, float]] = {}
+        self._locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
+        self._plans_cache: Dict[str, Dict[str, dict]] = defaultdict(dict)
+        self._cleanup_interval = cleanup_interval
+        self._last_cleanup = time.time()
+    
+    async def evaluate(
+        self,
+        key: str,
+        max_tokens: int,
+        period: float,
+        cost: int,
+        strategy: Strategy
+    ) -> LimitResult:
+        if cost <= 0:
+            raise ValueError(f"Cost must be positive, got {cost}")
+        
+        lock = self._locks[key]
+        async with lock:
+            try:
+                # Periodic cleanup
+                now = time.time()
+                if now - self._last_cleanup > self._cleanup_interval:
+                    self._cleanup_expired(now)
+                    self._last_cleanup = now
+                
+                if strategy == Strategy.FIXED_WINDOW:
+                    return self._eval_fixed(key, max_tokens, period, cost, now)
+                elif strategy == Strategy.TOKEN_BUCKET:
+                    return self._eval_token(key, max_tokens, period, cost, now)
+                elif strategy == Strategy.SLIDING_WINDOW:
+                    return self._eval_sliding(key, max_tokens, period, cost, now)
+                elif strategy == Strategy.LEAKY_BUCKET:
+                    return self._eval_leaky(key, max_tokens, period, cost, now)
+                raise ValueError(f"Unknown strategy: {strategy}")
+            finally:
+                # Prevent lock buildup
+                if len(self._locks) > 10000:
+                    self._locks = defaultdict(
+                        asyncio.Lock,
+                        {k: v for k, v in self._locks.items() 
+                         if k in self._fixed_window or k in self._token_bucket}
+                    )
+    
+    def _eval_fixed(
+        self, key: str, max_tokens: int, period: float, cost: int, now: float
+    ) -> LimitResult:
+        state = self._fixed_window.get(key)
+        
+        if state is None or now - state[1] >= period:
+            if cost > max_tokens:
+                return LimitResult(False, 0, now + period, period)
+            self._fixed_window[key] = (cost, now)
+            return LimitResult(True, max_tokens - cost, now + period, 0.0)
+            
+        used, window_start = state
+        reset_at = window_start + period
+        
+        if used + cost > max_tokens:
+            return LimitResult(False, max_tokens - used, reset_at, max(0, reset_at - now))
+            
+        self._fixed_window[key] = (used + cost, window_start)
+        return LimitResult(True, max_tokens - (used + cost), reset_at, 0.0)
+
+    def _eval_token(
+        self, key: str, max_tokens: int, period: float, cost: int, now: float
+    ) -> LimitResult:
+        state = self._token_bucket.get(key)
+        rate = max_tokens / period
+        
+        if state is None:
+            tokens, last_refill = float(max_tokens), now
+        else:
+            tokens, last_refill = state
+            elapsed = now - last_refill
+            if elapsed < 0:
+                elapsed = 0  # Clock skew protection
+            tokens = min(float(max_tokens), tokens + elapsed * rate)
+            
+        if tokens >= cost:
+            tokens -= cost
+            self._token_bucket[key] = (tokens, now)
+            return LimitResult(True, int(tokens), now + period, 0.0)
+            
+        retry_after = (cost - tokens) / rate
+        return LimitResult(False, int(tokens), now + period, retry_after)
+
+    def _eval_sliding(
+        self, key: str, max_tokens: int, period: float, cost: int, now: float
+    ) -> LimitResult:
+        window_start = now - period
+        history = [ts for ts in self._sliding_window.get(key, []) if ts > window_start]
+        
+        if len(history) + cost > max_tokens:
+            retry_after = max(0, (history[0] + period - now) if history else period)
+            return LimitResult(False, max_tokens - len(history), now + period, retry_after)
+            
+        history.extend([now] * cost)
+        self._sliding_window[key] = history
+        return LimitResult(True, max_tokens - len(history), now + period, 0.0)
+
+    def _eval_leaky(
+        self, key: str, max_tokens: int, period: float, cost: int, now: float
+    ) -> LimitResult:
+        state = self._leaky_bucket.get(key)
+        leak_rate = max_tokens / period
+        
+        if state is None:
+            water, last_leak = 0.0, now
+        else:
+            water, last_leak = state
+            elapsed = now - last_leak
+            if elapsed < 0:
+                elapsed = 0  # Clock skew protection
+            water = max(0.0, water - elapsed * leak_rate)
+            
+        if water + cost <= max_tokens:
+            water += cost
+            self._leaky_bucket[key] = (water, now)
+            return LimitResult(True, int(max_tokens - water), now + period, 0.0)
+            
+        retry_after = (water + cost - max_tokens) / leak_rate
+        return LimitResult(False, int(max_tokens - water), now + period, retry_after)
+    
+    def _cleanup_expired(self, now: float) -> None:
+        """Remove stale data to prevent memory leaks."""
+        max_period = 3600.0
+        
+        self._fixed_window = {
+            k: v for k, v in self._fixed_window.items()
+            if now - v[1] < max_period
+        }
+        self._token_bucket = {
+            k: v for k, v in self._token_bucket.items()
+            if now - v[1] < max_period
+        }
+        self._leaky_bucket = {
+            k: v for k, v in self._leaky_bucket.items()
+            if now - v[1] < max_period
+        }
+        
+        window_start = now - max_period
+        for key in list(self._sliding_window.keys()):
+            self._sliding_window[key] = [
+                ts for ts in self._sliding_window[key] if ts > window_start
+            ]
+            if not self._sliding_window[key]:
+                del self._sliding_window[key]
+
+    async def push_plan(self, target: str, plan_name: str, limit: int, reset_period: str, identifier_type: str, reset_day: int, timezone: str) -> None:
+        self._plans_cache[target][plan_name] = {
+            "limit": limit,
+            "reset_period": reset_period,
+            "identifier_type": identifier_type,
+            "reset_day": reset_day,
+            "timezone": timezone
+        }
+
+    async def pull_plan(self, target: str, plan_name: str) -> Optional[dict]:
+        return self._plans_cache.get(target, {}).get(plan_name)

metered/backends/redis.py

@@ -0,0 +1,249 @@
+from typing import Any, Optional
+import time
+import math
+import uuid
+from metered.core.backends import BaseBackend
+from metered.core.types import LimitResult, Strategy
+
+FIXED_WINDOW_LUA = """
+local current = redis.call("GET", KEYS[1])
+if current then
+    if tonumber(current) + tonumber(ARGV[1]) > tonumber(ARGV[2]) then
+        return {0, tonumber(ARGV[2]) - tonumber(current), redis.call("PTTL", KEYS[1])}
+    end
+    redis.call("INCRBY", KEYS[1], tonumber(ARGV[1]))
+    return {1, tonumber(ARGV[2]) - tonumber(current) - tonumber(ARGV[1]), redis.call("PTTL", KEYS[1])}
+else
+    if tonumber(ARGV[1]) > tonumber(ARGV[2]) then
+        return {0, tonumber(ARGV[2]), math.floor(tonumber(ARGV[3]) * 1000)}
+    end
+    redis.call("SET", KEYS[1], tonumber(ARGV[1]), "PX", math.floor(tonumber(ARGV[3]) * 1000))
+    return {1, tonumber(ARGV[2]) - tonumber(ARGV[1]), math.floor(tonumber(ARGV[3]) * 1000)}
+end
+"""
+
+TOKEN_BUCKET_LUA = """
+local key = KEYS[1]
+local cost = tonumber(ARGV[1])
+local max_tokens = tonumber(ARGV[2])
+local period_sec = tonumber(ARGV[3])
+local now = tonumber(ARGV[4])
+
+local tokens = redis.call('HGET', key, 'tokens')
+local last_refill = redis.call('HGET', key, 'last_refill')
+
+if not tokens then
+    tokens = max_tokens
+    last_refill = now
+else
+    tokens = tonumber(tokens)
+    last_refill = tonumber(last_refill)
+end
+
+if last_refill > now then
+    last_refill = now
+end
+
+local elapsed = now - last_refill
+local rate = max_tokens / period_sec
+local new_tokens = math.min(max_tokens, tokens + (elapsed * rate))
+
+if new_tokens >= cost then
+    new_tokens = new_tokens - cost
+    redis.call('HSET', key, 'tokens', tostring(new_tokens), 'last_refill', tostring(now))
+    redis.call('EXPIRE', key, math.ceil(period_sec))
+    return {1, math.floor(new_tokens), 0}
+else
+    local wait = (cost - new_tokens) / rate
+    return {0, math.floor(new_tokens), math.ceil(wait * 1000)}
+end
+"""
+
+SLIDING_WINDOW_LUA = """
+local key = KEYS[1]
+local cost = tonumber(ARGV[1])
+local max_tokens = tonumber(ARGV[2])
+local period = tonumber(ARGV[3])
+local now = tonumber(ARGV[4])
+local random_id = ARGV[5]
+
+local window_start = now - period
+local dropped = redis.call('ZRANGEBYSCORE', key, '-inf', window_start)
+local dropped_cost = 0
+
+for i, member in ipairs(dropped) do
+    local c = string.match(member, ":(%d+):")
+    if c then
+        dropped_cost = dropped_cost + tonumber(c)
+    end
+end
+
+local sum_exists = redis.call('EXISTS', key .. ':sum')
+if sum_exists == 1 and dropped_cost > 0 then
+    local current_sum_val = tonumber(redis.call('GET', key .. ':sum') or "0")
+    if current_sum_val - dropped_cost < 0 then
+        redis.call('SET', key .. ':sum', 0)
+    else
+        redis.call('DECRBY', key .. ':sum', dropped_cost)
+    end
+end
+redis.call('ZREMRANGEBYSCORE', key, '-inf', window_start)
+
+local current_sum = tonumber(redis.call('GET', key .. ':sum') or "0")
+if current_sum + cost <= max_tokens then
+    redis.call('ZADD', key, now, now .. ":" .. cost .. ":" .. random_id)
+    redis.call('INCRBY', key .. ':sum', cost)
+    redis.call('EXPIRE', key, math.ceil(period) + 1)
+    redis.call('EXPIRE', key .. ':sum', math.ceil(period) + 1)
+    return {1, max_tokens - (current_sum + cost), 0}
+else
+    return {0, max_tokens - current_sum, math.ceil(period * 1000)}
+end
+"""
+
+LEAKY_BUCKET_LUA = """
+local key = KEYS[1]
+local cost = tonumber(ARGV[1])
+local capacity = tonumber(ARGV[2])
+local period_sec = tonumber(ARGV[3])
+local now = tonumber(ARGV[4])
+
+local water = redis.call('HGET', key, 'water')
+local last_leak = redis.call('HGET', key, 'last_leak')
+
+if not water then
+    water = 0
+    last_leak = now
+else
+    water = tonumber(water)
+    last_leak = tonumber(last_leak)
+end
+
+if last_leak > now then
+    last_leak = now
+end
+
+local elapsed = now - last_leak
+local leak_rate = capacity / period_sec
+local leaked = elapsed * leak_rate
+water = math.max(0, water - leaked)
+
+if water + cost <= capacity then
+    water = water + cost
+    redis.call('HSET', key, 'water', tostring(water), 'last_leak', tostring(now))
+    redis.call('EXPIRE', key, math.ceil(period_sec))
+    return {1, math.floor(capacity - water), 0}
+else
+    local wait = (water + cost - capacity) / leak_rate
+    return {0, math.floor(capacity - water), math.ceil(wait * 1000)}
+end
+"""
+
+class RedisBackend(BaseBackend):
+    def __init__(self, redis_client: Any) -> None:
+        self.client = redis_client
+        self._scripts = {
+            Strategy.FIXED_WINDOW: self.client.register_script(FIXED_WINDOW_LUA),
+            Strategy.TOKEN_BUCKET: self.client.register_script(TOKEN_BUCKET_LUA),
+            Strategy.SLIDING_WINDOW: self.client.register_script(SLIDING_WINDOW_LUA),
+            Strategy.LEAKY_BUCKET: self.client.register_script(LEAKY_BUCKET_LUA),
+        }
+
+    async def evaluate(self, key: str, max_tokens: int, period: float, cost: int, strategy: Strategy) -> LimitResult:
+        if cost <= 0:
+            raise ValueError(f"Cost must be positive, got {cost}")
+            
+        now = time.time()
+        
+        if strategy not in self._scripts:
+            raise NotImplementedError(f"Strategy {strategy} not supported")
+            
+        script = self._scripts[strategy]
+        
+        try:
+            if strategy == Strategy.FIXED_WINDOW:
+                allowed, remaining, pttl = await script(keys=[key], args=[cost, max_tokens, period])
+                reset_at = now + (pttl / 1000.0)
+                retry_after = (pttl / 1000.0) if not allowed else 0.0
+            
+            elif strategy in (Strategy.TOKEN_BUCKET, Strategy.LEAKY_BUCKET):
+                allowed, remaining, pttl = await script(keys=[key], args=[cost, max_tokens, period, now])
+                reset_at = now + period
+                retry_after = (pttl / 1000.0) if not allowed else 0.0
+                
+            elif strategy == Strategy.SLIDING_WINDOW:
+                random_id = f"{uuid.uuid4()}:{time.time_ns()}"
+                allowed, remaining, pttl = await script(keys=[key], args=[cost, max_tokens, period, now, random_id])
+                reset_at = now + period
+                retry_after = (pttl / 1000.0) if not allowed else 0.0
+                
+            return LimitResult(bool(allowed), int(remaining), reset_at, retry_after)
+            
+        except Exception as e:
+            import logging
+            logging.getLogger("metered").error(f"Redis backend evaluation failed: {e}")
+            # Fail-safe mode: reject request on Redis failure to prevent abuse,
+            # or allow it if you prefer high-availability. Here we follow fail-safe:
+            return LimitResult(False, 0, now + period, period)
+        
+    async def push_event(self, queue: str, event_name: str, payload: dict[str, Any]) -> None:
+        import json
+        data = {
+            "event": event_name,
+            "payload": payload,
+            "timestamp": time.time()
+        }
+        await self.client.lpush(queue, json.dumps(data))
+        
+    async def pop_event(self, queue: str, timeout: int = 0) -> Any:
+        # BRPOP returns a tuple (queue_name, data)
+        result = await self.client.brpop(queue, timeout=timeout)
+        if result:
+            data = result[1]
+            unique_id = str(uuid.uuid4())
+            await self.client.setex(f"event:{unique_id}", 30, data)
+            return (unique_id, data)
+        return None
+
+    async def ack_event(self, unique_id: str) -> None:
+        await self.client.delete(f"event:{unique_id}")
+
+    async def nack_event(self, unique_id: str, requeue: bool = True) -> None:
+        data = await self.client.get(f"event:{unique_id}")
+        await self.client.delete(f"event:{unique_id}")
+        if requeue and data:
+            await self.client.lpush("metered_events", data)
+    async def increment_usage(self, key: str, cost: int, retention_days: int) -> None:
+        await self.client.incrby(key, cost)
+        # Only set expire if it's a new key, or we can just always set it
+        await self.client.expire(key, retention_days * 86400)
+
+    async def push_plan(self, target: str, plan_name: str, limit: int, reset_period: str, identifier_type: str, reset_day: int, timezone: str) -> None:
+        key = f"metered:plan:{target}:{plan_name}"
+        data = {
+            "limit": str(limit),
+            "reset_period": reset_period,
+            "identifier_type": identifier_type,
+            "reset_day": str(reset_day),
+            "timezone": timezone
+        }
+        await self.client.hset(key, mapping=data)
+
+    async def pull_plan(self, target: str, plan_name: str) -> Optional[dict]:
+        key = f"metered:plan:{target}:{plan_name}"
+        data = await self.client.hgetall(key)
+        if not data:
+            return None
+        
+        # Redis hgetall returns bytes if not decode_responses=True, but we assume it might be strings or bytes
+        # Let's decode if needed
+        def _decode(val):
+            return val.decode('utf-8') if isinstance(val, bytes) else val
+            
+        return {
+            "limit": int(_decode(data.get(b"limit", data.get("limit", 0)))),
+            "reset_period": _decode(data.get(b"reset_period", data.get("reset_period", ""))),
+            "identifier_type": _decode(data.get(b"identifier_type", data.get("identifier_type", ""))),
+            "reset_day": int(_decode(data.get(b"reset_day", data.get("reset_day", 1)))),
+            "timezone": _decode(data.get(b"timezone", data.get("timezone", "UTC")))
+        }

metered/core/backends.py

@@ -0,0 +1,19 @@
+from typing import Protocol, Optional
+from .types import LimitResult, Strategy
+
+class BaseBackend(Protocol):
+    async def evaluate(
+        self, 
+        key: str, 
+        max_tokens: int, 
+        period: float, 
+        cost: int, 
+        strategy: Strategy
+    ) -> LimitResult:
+        ...
+
+    async def push_plan(self, target: str, plan_name: str, limit: int, reset_period: str, identifier_type: str, reset_day: int, timezone: str) -> None:
+        ...
+        
+    async def pull_plan(self, target: str, plan_name: str) -> Optional[dict]:
+        ...