Capture kat-coder + qwen3-coder-30b outputs (5 prompts each, API-run, NOT graded yet)
Via tools/grade_run.py against LM Studio (load -> all prompts -> unload -> verify).
kat-coder: lfu/tts/webhook/automation/rust (lfu has a real SyntaxError: 'self._ bump_freq')
qwen3-coder-30b: lfu/tts/webhook/automation/rust (all 4 py parse clean; terse ~2k tok outputs)
Grading deferred to a later session (capture-only mode).
Script hardening:
- --resident / --unload-all commands (read loaded_instances, POST unload, verify)
- MEMORY GUARD: aborts if a different model is resident (never stack 30GB models)
- TTFT self-discovery dump (LM Studio returns empty stats non-streaming;
fix next session: switch to /api/v1/chat stream:true -> chat.end result.stats)
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Claude
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"usage": {
|
||||
"prompt_tokens": 621,
|
||||
"completion_tokens": 1795,
|
||||
"total_tokens": 2416,
|
||||
"completion_tokens_details": {
|
||||
"reasoning_tokens": 0
|
||||
}
|
||||
},
|
||||
"stats": {},
|
||||
"system_fingerprint": "qwen/qwen3-coder-30b"
|
||||
}
|
||||
@@ -0,0 +1,338 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Async batch processor with checkpointing, idempotency, retries, and graceful shutdown.
|
||||
|
||||
Usage:
|
||||
python batch_processor.py --items job-a job-b ... [--checkpoint checkpoint.json]
|
||||
|
||||
Or import and use the BatchProcessor class directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public exception
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ProcessingError(Exception):
|
||||
"""Raised by ``process()`` when an external call fails."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Simulated flaky worker (the "external call")
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def process(item: str) -> str:
|
||||
"""
|
||||
Simulate a flaky external call.
|
||||
|
||||
Succeeds ~80 % of the time, raises ``ProcessingError`` otherwise.
|
||||
Latency is uniform random in [10 ms, 80 ms].
|
||||
"""
|
||||
import random
|
||||
await asyncio.sleep(random.uniform(0.010, 0.080))
|
||||
if random.random() < 0.20: # 20 % failure rate
|
||||
raise ProcessingError(f"transient failure for {item!r}")
|
||||
return f"done:{item}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checkpoint data model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class ItemResult:
|
||||
status: str # "succeeded" | "failed"
|
||||
attempts: int = 0
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Checkpoint:
|
||||
"""Persisted progress. All fields are JSON-serialisable."""
|
||||
completed: dict[str, ItemResult] = field(default_factory=dict)
|
||||
|
||||
# Convenience accessors
|
||||
@property
|
||||
def succeeded(self) -> list[str]:
|
||||
return [k for k, v in self.completed.items() if v.status == "succeeded"]
|
||||
|
||||
@property
|
||||
def failed(self) -> list[str]:
|
||||
return [k for k, v in self.completed.items() if v.status == "failed"]
|
||||
|
||||
@property
|
||||
def summary(self) -> dict[str, int]:
|
||||
return {
|
||||
"succeeded": len(self.succeeded),
|
||||
"failed": len(self.failed),
|
||||
}
|
||||
|
||||
def is_done(self, item: str) -> bool:
|
||||
return item in self.completed
|
||||
|
||||
def mark_succeeded(self, item: str) -> None:
|
||||
self.completed[item] = ItemResult(status="succeeded")
|
||||
|
||||
def mark_failed(self, item: str, error: str = "") -> None:
|
||||
self.completed[item] = ItemResult(status="failed", error=error)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checkpoint I/O (atomic writes)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_checkpoint(path: Path) -> Checkpoint:
|
||||
"""Load checkpoint from disk, or return an empty one if absent."""
|
||||
if not path.exists():
|
||||
return Checkpoint()
|
||||
raw = json.loads(path.read_text(encoding="utf-8"))
|
||||
completed: dict[str, ItemResult] = {}
|
||||
for k, v in raw.get("completed", {}).items():
|
||||
completed[k] = ItemResult(**v)
|
||||
return Checkpoint(completed=completed)
|
||||
|
||||
|
||||
async def save_checkpoint(cp: Checkpoint, path: Path) -> None:
|
||||
"""
|
||||
Atomically persist *cp* to *path*.
|
||||
|
||||
Writes to a temp file in the same directory then ``os.replace()`` —
|
||||
this is atomic on every POSIX filesystem and on Windows when the dest
|
||||
is on the same volume. A crash mid-write can never corrupt the target.
|
||||
"""
|
||||
parent = path.parent
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _write() -> None:
|
||||
fd, tmp = tempfile.mkstemp(suffix=".tmp", dir=parent)
|
||||
try:
|
||||
payload = json.dumps(asdict(cp), indent=2, ensure_ascii=False) + "\n"
|
||||
os.write(fd, payload.encode("utf-8"))
|
||||
os.close(fd)
|
||||
# Atomic rename — the only point where the target is visible.
|
||||
os.replace(tmp, str(path))
|
||||
except BaseException:
|
||||
os.close(fd) if not None else None
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
await loop.run_in_executor(None, _write)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch processor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class BatchProcessor:
|
||||
"""
|
||||
Process a list of items with bounded concurrency, retries, checkpointing,
|
||||
and graceful SIGINT handling.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
items: list[str],
|
||||
*,
|
||||
checkpoint_path: str | Path = "checkpoint.json",
|
||||
max_concurrency: int = 8,
|
||||
max_retries: int = 3,
|
||||
base_backoff: float = 0.1,
|
||||
) -> None:
|
||||
self.items = list(items)
|
||||
self.checkpoint_path = Path(checkpoint_path)
|
||||
self.max_concurrency = max_concurrency
|
||||
self.max_retries = max_retries
|
||||
self.base_backoff = base_backoff
|
||||
|
||||
# Load any existing checkpoint for idempotency
|
||||
self.cp = load_checkpoint(self.checkpoint_path)
|
||||
|
||||
# Semaphore for bounded concurrency
|
||||
self._semaphore: asyncio.Semaphore | None = None
|
||||
|
||||
# Graceful-shutdown flag — set by the signal handler
|
||||
self._shutdown_requested = asyncio.Event()
|
||||
|
||||
# Track in-flight tasks so we can wait for them on shutdown
|
||||
self._in_flight: set[asyncio.Task[Any]] = set()
|
||||
|
||||
# Wall-clock start time (set in ``run``)
|
||||
self._start_time: float = 0.0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Signal handling
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _install_signal_handlers(self) -> None:
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(sig, self._on_signal, sig)
|
||||
|
||||
def _remove_signal_handlers(self) -> None:
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
try:
|
||||
loop.remove_signal_handler(sig)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
def _on_signal(self, sig: signal.Signals) -> None:
|
||||
# First signal: request graceful shutdown.
|
||||
# Subsequent signals during the drain window force immediate exit.
|
||||
if not self._shutdown_requested.is_set():
|
||||
print(f"\n[BatchProcessor] {sig.name} received — draining in-flight tasks …", file=sys.stderr)
|
||||
self._shutdown_requested.set()
|
||||
else:
|
||||
print(f"\n[BatchProcessor] {sig.name} received again — forcing exit.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core processing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _process_one(self, item: str) -> None:
|
||||
"""Retry *item* up to max_retries with exponential backoff."""
|
||||
attempts = 0
|
||||
last_error = ""
|
||||
while attempts < self.max_retries:
|
||||
attempts += 1
|
||||
try:
|
||||
result = await process(item)
|
||||
self.cp.mark_succeeded(item)
|
||||
await save_checkpoint(self.cp, self.checkpoint_path)
|
||||
print(f" ✓ {item} (attempt {attempts})")
|
||||
return
|
||||
except ProcessingError as exc:
|
||||
last_error = str(exc)
|
||||
if attempts < self.max_retries:
|
||||
backoff = self.base_backoff * (2 ** (attempts - 1))
|
||||
print(f" ✗ {item} attempt {attempts}/{self.max_retries} failed — retry in {backoff:.2f}s", file=sys.stderr)
|
||||
try:
|
||||
await asyncio.sleep(backoff)
|
||||
except asyncio.CancelledError:
|
||||
# We were cancelled mid-backoff — do NOT mark failed.
|
||||
raise
|
||||
|
||||
# Exhausted retries → terminal failure
|
||||
self.cp.mark_failed(item, error=last_error)
|
||||
await save_checkpoint(self.cp, self.checkpoint_path)
|
||||
print(f" ☠ {item} failed after {attempts} attempts", file=sys.stderr)
|
||||
|
||||
async def _run_item(self, item: str) -> None:
|
||||
"""Wrap ``_process_one`` with concurrency tracking and shutdown guard."""
|
||||
task = asyncio.current_task()
|
||||
assert task is not None
|
||||
self._in_flight.add(task)
|
||||
try:
|
||||
# Check shutdown before starting work on this item.
|
||||
if self._shutdown_requested.is_set():
|
||||
return
|
||||
await self._process_one(item)
|
||||
finally:
|
||||
self._in_flight.discard(task)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public entry point
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def run(self) -> dict[str, Any]:
|
||||
"""
|
||||
Run the batch. Returns a summary dict suitable for JSON output.
|
||||
"""
|
||||
self._start_time = time.monotonic()
|
||||
self._install_signal_handlers()
|
||||
|
||||
# Items already in the checkpoint are skipped.
|
||||
pending = [it for it in self.items if not self.cp.is_done(it)]
|
||||
skipped = len(self.items) - len(pending)
|
||||
|
||||
self._semaphore = asyncio.Semaphore(self.max_concurrency)
|
||||
|
||||
try:
|
||||
async def _bounded(item: str) -> None:
|
||||
async with self._semaphore:
|
||||
await self._run_item(item)
|
||||
|
||||
workers = [asyncio.create_task(_bounded(it)) for it in pending]
|
||||
|
||||
# Wait until either all workers finish or shutdown is requested.
|
||||
if workers:
|
||||
done, pending_tasks = await asyncio.wait(
|
||||
workers,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
# Check if we should drain or abort.
|
||||
if self._shutdown_requested.is_set():
|
||||
# Cancel any still-pending workers (they won't start new items).
|
||||
for t in pending_tasks:
|
||||
t.cancel()
|
||||
# Wait for all to finish/cancel.
|
||||
if pending_tasks:
|
||||
await asyncio.gather(*pending_tasks, return_exceptions=True)
|
||||
else:
|
||||
# Some finished; continue with the rest.
|
||||
if pending_tasks:
|
||||
done2, _ = await asyncio.wait(pending_tasks)
|
||||
await asyncio.gather(*done2, return_exceptions=True)
|
||||
|
||||
# Final flush (idempotent — same content).
|
||||
await save_checkpoint(self.cp, self.checkpoint_path)
|
||||
|
||||
finally:
|
||||
self._remove_signal_handlers()
|
||||
|
||||
elapsed_ms = int((time.monotonic() - self._start_time) * 1000)
|
||||
summary = {
|
||||
"succeeded": len(self.cp.succeeded),
|
||||
"failed": len(self.cp.failed),
|
||||
"skipped": skipped,
|
||||
"total": len(self.items),
|
||||
"elapsed_ms": elapsed_ms,
|
||||
}
|
||||
return summary
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Async batch processor with checkpointing")
|
||||
parser.add_argument("items", nargs="+", help="Item identifiers to process")
|
||||
parser.add_argument("--checkpoint", default="checkpoint.json", help="Checkpoint file path")
|
||||
parser.add_argument("--concurrency", type=int, default=8, help="Max concurrent items")
|
||||
parser.add_argument("--retries", type=int, default=3, help="Max retries per item")
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
processor = BatchProcessor(
|
||||
args.items,
|
||||
checkpoint_path=args.checkpoint,
|
||||
max_concurrency=args.concurrency,
|
||||
max_retries=args.retries,
|
||||
)
|
||||
|
||||
summary = asyncio.run(processor.run())
|
||||
print(json.dumps(summary))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,588 @@
|
||||
"""
|
||||
In-Memory Concurrent LFU Cache with Async TTL Eviction and Atomic Transactions.
|
||||
|
||||
Pure Python 3.11+ implementation — no external dependencies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import weakref
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal Doubly-Linked List Node (frequency bucket member)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class _Node:
|
||||
"""A single cache entry stored inside a frequency bucket."""
|
||||
|
||||
key: Any
|
||||
value: Any
|
||||
freq: int = 1
|
||||
expires_at: float = 0.0 # 0 means no expiry
|
||||
|
||||
prev: Optional[_Node] = field(default=None, repr=False)
|
||||
next: Optional[_Node] = field(default=None, repr=False)
|
||||
|
||||
@property
|
||||
def is_expired(self) -> bool:
|
||||
return self.expires_at > 0 and time.monotonic() >= self.expires_at
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Doubly-Linked List wrapper for a single frequency tier
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _FreqList:
|
||||
"""Doubly-linked list with sentinel head/tail for O(1) insert/remove."""
|
||||
|
||||
__slots__ = ("head", "tail", "size")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.head = _Node(key=None, value=None) # sentinel
|
||||
self.tail = _Node(key=None, value=None) # sentinel
|
||||
self.head.next = self.tail
|
||||
self.tail.prev = self.head
|
||||
self.size = 0
|
||||
|
||||
# -- public helpers -----------------------------------------------------
|
||||
|
||||
def append_right(self, node: _Node) -> None:
|
||||
"""Append *node* just before the tail sentinel (LRU-end)."""
|
||||
pred = self.tail.prev
|
||||
node.prev = pred
|
||||
node.next = self.tail
|
||||
pred.next = node
|
||||
self.tail.prev = node
|
||||
self.size += 1
|
||||
|
||||
def remove(self, node: _Node) -> None:
|
||||
"""Remove *node* from the list in O(1)."""
|
||||
pred, succ = node.prev, node.next
|
||||
pred.next = succ
|
||||
succ.prev = pred
|
||||
node.prev = node.next = None
|
||||
self.size -= 1
|
||||
|
||||
def pop_left(self) -> Optional[_Node]:
|
||||
"""Remove and return the node just after head sentinel (MRU-end)."""
|
||||
if self.size == 0:
|
||||
return None
|
||||
node = self.head.next
|
||||
self.remove(node)
|
||||
return node
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return self.size == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transaction handle
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class Transaction:
|
||||
"""
|
||||
Represents an isolated sub-session on the cache.
|
||||
|
||||
* ``put`` / ``delete`` buffer changes locally.
|
||||
* ``get`` sees uncommitted local writes ("read your own writes").
|
||||
* ``commit`` applies all buffered changes atomically to the global cache.
|
||||
* ``rollback`` discards everything without touching global state.
|
||||
"""
|
||||
|
||||
def __init__(self, cache: "LFUCache") -> None:
|
||||
self._cache = cache
|
||||
# key -> (value, expires_at) — latest write wins on conflict
|
||||
self._writes: dict[Any, tuple[Any, float]] = {}
|
||||
# keys deleted within this transaction
|
||||
self._deletes: set[Any] = set()
|
||||
self._committed = False
|
||||
self._rolled_back = False
|
||||
|
||||
# -- mutating operations ------------------------------------------------
|
||||
|
||||
def put(self, key: Any, value: Any, ttl_seconds: float = 0.0) -> None:
|
||||
"""Buffer a write inside this transaction."""
|
||||
if self._committed or self._rolled_back:
|
||||
raise RuntimeError("Transaction already closed")
|
||||
expires_at = (time.monotonic() + ttl_seconds) if ttl_seconds > 0 else 0.0
|
||||
self._writes[key] = (value, expires_at)
|
||||
self._deletes.discard(key)
|
||||
|
||||
def delete(self, key: Any) -> None:
|
||||
"""Buffer a deletion inside this transaction."""
|
||||
if self._committed or self._rolled_back:
|
||||
raise RuntimeError("Transaction already closed")
|
||||
self._deletes.add(key)
|
||||
self._writes.pop(key, None)
|
||||
|
||||
# -- read operations ----------------------------------------------------
|
||||
|
||||
def get(self, key: Any) -> Optional[Any]:
|
||||
"""
|
||||
Return the value for *key*, checking local buffer first, then global.
|
||||
Returns ``None`` if the key is absent or expired.
|
||||
"""
|
||||
if self._committed or self._rolled_back:
|
||||
raise RuntimeError("Transaction already closed")
|
||||
|
||||
# 1. Check local uncommitted writes
|
||||
if key in self._deletes:
|
||||
return None
|
||||
if key in self._writes:
|
||||
value, expires_at = self._writes[key]
|
||||
if expires_at > 0 and time.monotonic() >= expires_at:
|
||||
self._deletes.add(key)
|
||||
return None
|
||||
return value
|
||||
|
||||
# 2. Fall back to global cache (with lazy TTL check)
|
||||
return self._cache.get(key)
|
||||
|
||||
# -- lifecycle ----------------------------------------------------------
|
||||
|
||||
async def commit(self) -> None:
|
||||
"""Apply all buffered changes atomically to the global cache."""
|
||||
if self._committed or self._rolled_back:
|
||||
raise RuntimeError("Transaction already closed")
|
||||
await self._cache._apply_transaction(self)
|
||||
self._committed = True
|
||||
|
||||
async def rollback(self) -> None:
|
||||
"""Discard all buffered changes."""
|
||||
if self._committed or self._rolled_back:
|
||||
raise RuntimeError("Transaction already closed")
|
||||
self._writes.clear()
|
||||
self._deletes.clear()
|
||||
self._rolled_back = True
|
||||
|
||||
def __repr__(self) -> str: # pragma: no cover
|
||||
state = "committed" if self._committed else (
|
||||
"rolled_back" if self._rolled_back else "active"
|
||||
)
|
||||
return f"<Transaction {state} writes={len(self._writes)} deletes={len(self._deletes)}>"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main LFU Cache
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LFUCache:
|
||||
"""
|
||||
In-memory concurrent LFU cache with O(1) get/put, async TTL eviction,
|
||||
and atomic transaction support.
|
||||
"""
|
||||
|
||||
def __init__(self, capacity: int) -> None:
|
||||
if capacity < 1:
|
||||
raise ValueError("capacity must be >= 1")
|
||||
|
||||
self._capacity = capacity
|
||||
|
||||
# key -> _Node (global cache)
|
||||
self._cache: dict[Any, _Node] = {}
|
||||
|
||||
# freq -> _FreqList (frequency buckets)
|
||||
self._freq_map: dict[int, _FreqList] = {}
|
||||
|
||||
# Track the current minimum frequency for O(1) eviction
|
||||
self._min_freq: int = 0
|
||||
|
||||
# Concurrency primitives
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# Background evictor task handle
|
||||
self._evictor_task: Optional[asyncio.Task[None]] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API — get / put
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get(self, key: Any) -> Optional[Any]:
|
||||
"""
|
||||
Retrieve *key* from the cache in O(1).
|
||||
|
||||
Lazy TTL eviction is performed on access.
|
||||
"""
|
||||
async with self._lock:
|
||||
node = self._cache.get(key)
|
||||
if node is None:
|
||||
return None
|
||||
|
||||
# Lazy TTL check
|
||||
if node.is_expired:
|
||||
self._evict_node(node)
|
||||
return None
|
||||
|
||||
# Bump frequency — O(1)
|
||||
self._ bump_freq(node)
|
||||
return node.value
|
||||
|
||||
async def put(self, key: Any, value: Any, ttl_seconds: float = 0.0) -> None:
|
||||
"""
|
||||
Insert or update *key* in the cache in O(1).
|
||||
|
||||
If the cache is at capacity, the LFU (least-recently-used tie-break)
|
||||
entry is evicted before insertion.
|
||||
"""
|
||||
expires_at = (time.monotonic() + ttl_seconds) if ttl_seconds > 0 else 0.0
|
||||
|
||||
async with self._lock:
|
||||
# Case 1: key already exists — update in place
|
||||
if key in self._cache:
|
||||
node = self._cache[key]
|
||||
old_freq = node.freq
|
||||
node.value = value
|
||||
node.expires_at = expires_at
|
||||
# Move to new frequency bucket
|
||||
self._remove_from_freq_list(node)
|
||||
node.freq += 1
|
||||
self._add_to_freq_list(node)
|
||||
# Update min_freq if the old bucket is now empty and was min
|
||||
if old_freq == self._min_freq and self._freq_map[old_freq].is_empty():
|
||||
del self._freq_map[old_freq]
|
||||
self._min_freq = node.freq
|
||||
return
|
||||
|
||||
# Case 2: cache full — evict LFU entry
|
||||
if len(self._cache) >= self._capacity:
|
||||
self._evict_one()
|
||||
|
||||
# Insert new node at frequency 1
|
||||
node = _Node(key=key, value=value, freq=1, expires_at=expires_at)
|
||||
self._cache[key] = node
|
||||
self._add_to_freq_list(node)
|
||||
self._min_freq = 1
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Transaction support
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def begin_transaction(self) -> Transaction:
|
||||
"""Start a new isolated transaction on this cache."""
|
||||
return Transaction(self)
|
||||
|
||||
async def _apply_transaction(self, tx: Transaction) -> None:
|
||||
"""
|
||||
Apply a transaction's buffered writes/deletes atomically.
|
||||
Must be called while holding self._lock (called from Transaction.commit).
|
||||
"""
|
||||
# --- Phase 1: apply deletes first (so a put+delete of same key works) ---
|
||||
for key in tx._deletes:
|
||||
node = self._cache.pop(key, None)
|
||||
if node is not None:
|
||||
self._remove_from_freq_list(node)
|
||||
if node.freq == self._min_freq and self._freq_map[node.freq].is_empty():
|
||||
del self._freq_map[node.freq]
|
||||
# Find new min freq
|
||||
if self._freq_map:
|
||||
self._min_freq = min(self._freq_map)
|
||||
else:
|
||||
self._min_freq = 0
|
||||
|
||||
# --- Phase 2: apply writes ------------------------------------------
|
||||
for key, (value, expires_at) in tx._writes.items():
|
||||
if key in self._cache:
|
||||
# Update existing node
|
||||
node = self._cache[key]
|
||||
old_freq = node.freq
|
||||
node.value = value
|
||||
node.expires_at = expires_at
|
||||
self._remove_from_freq_list(node)
|
||||
node.freq += 1
|
||||
self._add_to_freq_list(node)
|
||||
if old_freq == self._min_freq and self._freq_map[old_freq].is_empty():
|
||||
del self._freq_map[old_freq]
|
||||
if self._freq_map:
|
||||
self._min_freq = min(self._freq_map)
|
||||
else:
|
||||
# Insert new node (may need eviction first)
|
||||
if len(self._cache) >= self._capacity:
|
||||
self._evict_one()
|
||||
node = _Node(key=key, value=value, freq=1, expires_at=expires_at)
|
||||
self._cache[key] = node
|
||||
self._add_to_freq_list(node)
|
||||
self._min_freq = 1
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Background async TTL evictor
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start_evictor(self, interval: float = 0.5, batch_size: int = 32) -> None:
|
||||
"""Start the non-blocking background TTL sweep task."""
|
||||
if self._evictor_task is not None:
|
||||
return
|
||||
self._evictor_task = asyncio.create_task(
|
||||
self._evictor_loop(interval, batch_size)
|
||||
)
|
||||
|
||||
def stop_evictor(self) -> None:
|
||||
"""Stop the background TTL sweep task."""
|
||||
if self._evictor_task is not None:
|
||||
self._evictor_task.cancel()
|
||||
try:
|
||||
asyncio.get_event_loop().run_until_complete(self._evictor_task)
|
||||
except (asyncio.CancelledError, RuntimeError):
|
||||
pass
|
||||
self._evictor_task = None
|
||||
|
||||
async def _evictor_loop(self, interval: float, batch_size: int) -> None:
|
||||
"""Periodically scan and purge expired entries in small batches."""
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
async with self._lock:
|
||||
expired = [
|
||||
node for node in self._cache.values() if node.is_expired
|
||||
][:batch_size]
|
||||
for node in expired:
|
||||
self._evict_node(node)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _add_to_freq_list(self, node: _Node) -> None:
|
||||
"""Add *node* to its frequency bucket."""
|
||||
freq = node.freq
|
||||
if freq not in self._freq_map:
|
||||
self._freq_map[freq] = _FreqList()
|
||||
self._freq_map[freq].append_right(node)
|
||||
|
||||
def _remove_from_freq_list(self, node: _Node) -> None:
|
||||
"""Remove *node* from its current frequency bucket."""
|
||||
freq_list = self._freq_map.get(node.freq)
|
||||
if freq_list is not None:
|
||||
freq_list.remove(node)
|
||||
|
||||
def _bump_freq(self, node: _Node) -> None:
|
||||
"""Increment *node*'s frequency and move it to the next bucket."""
|
||||
old_freq = node.freq
|
||||
self._remove_from_freq_list(node)
|
||||
node.freq += 1
|
||||
self._add_to_freq_list(node)
|
||||
|
||||
# Update min_freq if the old bucket is now empty and was the minimum
|
||||
if old_freq == self._min_freq and self._freq_map[old_freq].is_empty():
|
||||
del self._freq_map[old_freq]
|
||||
if self._freq_map:
|
||||
self._min_freq = min(self._freq_map)
|
||||
else:
|
||||
self._min_freq = 0
|
||||
|
||||
def _evict_one(self) -> None:
|
||||
"""Evict the least-frequently-used entry (LRU tie-break). O(1)."""
|
||||
if self._min_freq not in self._freq_map or self._freq_map[self._min_freq].is_empty():
|
||||
# Fallback — should not happen in normal operation
|
||||
return
|
||||
|
||||
freq_list = self._freq_map[self._min_freq]
|
||||
node = freq_list.pop_left()
|
||||
if node is not None:
|
||||
del self._cache[node.key]
|
||||
if freq_list.is_empty():
|
||||
del self._freq_map[self._min_freq]
|
||||
if self._freq_map:
|
||||
self._min_freq = min(self._freq_map)
|
||||
else:
|
||||
self._min_freq = 0
|
||||
|
||||
def _evict_node(self, node: _Node) -> None:
|
||||
"""Remove a single expired node from all structures."""
|
||||
self._remove_from_freq_list(node)
|
||||
del self._cache[node.key]
|
||||
if node.freq == self._min_freq and self._freq_map.get(node.freq, _FreqList()).is_empty():
|
||||
if node.freq in self._freq_map:
|
||||
del self._freq_map[node.freq]
|
||||
if self._freq_map:
|
||||
self._min_freq = min(self._freq_map)
|
||||
else:
|
||||
self._min_freq = 0
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# Executable Unit Tests
|
||||
# ===================================================================
|
||||
|
||||
async def main() -> None:
|
||||
passed = 0
|
||||
failed = 0
|
||||
|
||||
def check(name: str, condition: bool) -> None:
|
||||
nonlocal passed, failed
|
||||
if condition:
|
||||
passed += 1
|
||||
print(f" ✓ {name}")
|
||||
else:
|
||||
failed += 1
|
||||
print(f" ✗ {name}")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# (a) O(1) LFU eviction order when capacity is reached
|
||||
# ------------------------------------------------------------------
|
||||
print("\n[a] LFU Eviction Order")
|
||||
|
||||
cache = LFUCache(capacity=3)
|
||||
|
||||
# Insert a, b, c → all freq=1
|
||||
await cache.put("a", 1)
|
||||
await cache.put("b", 2)
|
||||
await cache.put("c", 3)
|
||||
|
||||
# Access a twice → freq(a)=3, freq(b)=1, freq(c)=1
|
||||
await cache.get("a")
|
||||
await cache.put("a", 10) # bump to freq=4
|
||||
|
||||
# Access b once → freq(b)=2
|
||||
await cache.get("b")
|
||||
|
||||
# Now: a=freq4, b=freq2, c=freq1. Min freq = 1 (key c).
|
||||
# Insert d — should evict c (lowest freq).
|
||||
await cache.put("d", 4)
|
||||
|
||||
check("evicts least-frequent key (c)", await cache.get("c") is None)
|
||||
check("keeps a", await cache.get("a") == 10)
|
||||
check("keeps b", await cache.get("b") == 2)
|
||||
check("keeps d", await cache.get("d") == 4)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# (b) Lazy TTL vs. Background Async Sweep eviction
|
||||
# ------------------------------------------------------------------
|
||||
print("\n[b] TTL Eviction (Lazy + Background)")
|
||||
|
||||
cache2 = LFUCache(capacity=5)
|
||||
|
||||
await cache2.put("x", 1, ttl_seconds=0.05)
|
||||
await cache2.put("y", 2, ttl_seconds=10.0)
|
||||
|
||||
# Before expiry — both visible
|
||||
check("x present before TTL", await cache2.get("x") == 1)
|
||||
check("y present before TTL", await cache2.get("y") == 2)
|
||||
|
||||
# Wait for x to expire
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Lazy eviction on get
|
||||
check("x evicted lazily on get", await cache2.get("x") is None)
|
||||
check("y still present after x expired", await cache2.get("y") == 2)
|
||||
|
||||
# Background evictor test
|
||||
cache3 = LFUCache(capacity=5)
|
||||
await cache3.put("p", 1, ttl_seconds=0.05)
|
||||
await cache3.put("q", 2, ttl_seconds=10.0)
|
||||
|
||||
cache3.start_evictor(interval=0.05, batch_size=8)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
check("p evicted by background sweep", await cache3.get("p") is None)
|
||||
check("q still present after background sweep", await cache3.get("q") == 2)
|
||||
|
||||
cache3.stop_evictor()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# (c) Transaction commit visibility vs. rollback state restoration
|
||||
# ------------------------------------------------------------------
|
||||
print("\n[c] Atomic Transactions")
|
||||
|
||||
cache4 = LFUCache(capacity=5)
|
||||
await cache4.put("k1", 100)
|
||||
await cache4.put("k2", 200)
|
||||
|
||||
# --- Commit test -------------------------------------------------------
|
||||
tx1 = cache4.begin_transaction()
|
||||
tx1.put("k3", 300) # uncommitted — global should not see it
|
||||
tx1.put("k1", 999) # update existing
|
||||
|
||||
check("global doesn't see uncommitted put", await cache4.get("k3") is None)
|
||||
check("global still sees old k1", await cache4.get("k1") == 100)
|
||||
|
||||
# Read-your-own-writes
|
||||
check("tx sees its own put(k3)", tx1.get("k3") == 300)
|
||||
check("tx sees its own update(k1)", tx1.get("k1") == 999)
|
||||
|
||||
await tx1.commit()
|
||||
|
||||
check("global sees committed k3", await cache4.get("k3") == 300)
|
||||
check("global sees committed k1 update", await cache4.get("k1") == 999)
|
||||
|
||||
# --- Rollback test -----------------------------------------------------
|
||||
tx2 = cache4.begin_transaction()
|
||||
tx2.put("k4", 400)
|
||||
tx2.delete("k2")
|
||||
|
||||
check("tx sees deleted k2 as None", tx2.get("k2") is None)
|
||||
check("global still has k2 before rollback", await cache4.get("k2") == 200)
|
||||
|
||||
await tx2.rollback()
|
||||
|
||||
check("global k2 restored after rollback", await cache4.get("k2") == 200)
|
||||
check("global k4 absent after rollback", await cache4.get("k4") is None)
|
||||
|
||||
# --- Rollback of update test -------------------------------------------
|
||||
tx3 = cache4.begin_transaction()
|
||||
tx3.put("k1", 777)
|
||||
await tx3.rollback()
|
||||
check("global k1 restored after rollback of update", await cache4.get("k1") == 999)
|
||||
|
||||
# --- Transaction with TTL ----------------------------------------------
|
||||
tx4 = cache4.begin_transaction()
|
||||
tx4.put("temp", "value", ttl_seconds=0.05)
|
||||
check("tx sees its own TTL'd key", tx4.get("temp") == "value")
|
||||
await asyncio.sleep(0.1)
|
||||
check("tx sees expired TTL'd key as None", tx4.get("temp") is None)
|
||||
await tx4.rollback()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# (d) Stress test — 50 concurrent async tasks reading/writing
|
||||
# ------------------------------------------------------------------
|
||||
print("\n[d] Stress Test (50 concurrent tasks)")
|
||||
|
||||
cache5 = LFUCache(capacity=100)
|
||||
errors: list[str] = []
|
||||
|
||||
async def worker(task_id: int) -> None:
|
||||
try:
|
||||
for i in range(50):
|
||||
key = f"key_{task_id}_{i % 200}"
|
||||
if i % 3 == 0:
|
||||
await cache5.put(key, f"val_{task_id}_{i}", ttl_seconds=1.0)
|
||||
elif i % 3 == 1:
|
||||
await cache5.get(key)
|
||||
else:
|
||||
tx = cache5.begin_transaction()
|
||||
tx.put(f"tx_{task_id}_{i}", f"txval_{i}")
|
||||
val = tx.get(f"tx_{task_id}_{i}")
|
||||
check(f" task {task_id} tx read-your-own-write", val == f"txval_{i}")
|
||||
await tx.commit()
|
||||
except Exception as e:
|
||||
errors.append(f"task {task_id}: {e}")
|
||||
|
||||
tasks = [asyncio.create_task(worker(tid)) for tid in range(50)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
check("no concurrent errors", len(errors) == 0)
|
||||
check("cache has entries after stress", len(cache5._cache) > 0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Summary
|
||||
# ------------------------------------------------------------------
|
||||
total = passed + failed
|
||||
print(f"\n{'='*50}")
|
||||
print(f"Results: {passed}/{total} passed" + (
|
||||
"" if failed == 0 else f", {failed} FAILED"
|
||||
))
|
||||
print(f"{'='*50}")
|
||||
|
||||
if failed:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,537 @@
|
||||
```rust
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Cargo.toml dependencies (for reference):
|
||||
// tokio = { version = "1", features = ["full"] }
|
||||
// serde = { version = "1", features = ["derive"] }
|
||||
// futures = "0.3"
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::stream::{self, StreamExt};
|
||||
use serde::Serialize;
|
||||
use tokio::sync::{broadcast, mpsc, RwLock};
|
||||
use tokio::time::{interval, sleep};
|
||||
|
||||
// ─── Error types ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Errors that can occur when fetching data from a source.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum FetchError {
|
||||
#[error("source temporarily unavailable")]
|
||||
TemporarilyUnavailable,
|
||||
#[error("upstream timeout")]
|
||||
Timeout,
|
||||
}
|
||||
|
||||
/// Errors specific to the watcher service.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum ServiceError {
|
||||
#[error("watcher with id {0} not found")]
|
||||
WatcherNotFound(u32),
|
||||
#[error("channel send failed: {0}")]
|
||||
SendError(String),
|
||||
#[error("shutdown already in progress")]
|
||||
AlreadyShutdown,
|
||||
}
|
||||
|
||||
// ─── Shared data types ────────────────────────────────────────────────────────
|
||||
|
||||
/// A single item produced by a watcher's source poll.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
struct WatchedItem {
|
||||
watcher_id: u32,
|
||||
value: String,
|
||||
}
|
||||
|
||||
/// Events flowing from a watcher to the consumer.
|
||||
#[derive(Debug, Clone)]
|
||||
enum WatchEvent {
|
||||
Item(WatchedItem),
|
||||
/// Health update: `true` = healthy, `false` = unhealthy.
|
||||
HealthUpdate { id: u32, healthy: bool },
|
||||
}
|
||||
|
||||
/// Aggregated output produced by the consumer.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
struct AggregatedOutput {
|
||||
total_items: u64,
|
||||
items_per_watcher: HashMap<u32, u64>,
|
||||
}
|
||||
|
||||
// ─── Mock source ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Simulates an async data source.
|
||||
///
|
||||
/// Returns 0–3 random strings on success (~85 %), or a `FetchError` (~15 %).
|
||||
async fn mock_fetch(id: u32) -> Result<Vec<String>, FetchError> {
|
||||
// Simulate variable latency.
|
||||
sleep(Duration::from_millis(rand_u32() % 50)).await;
|
||||
|
||||
if rand_u32() % 100 < 15 {
|
||||
// ~15 % failure rate.
|
||||
if rand_u32() % 2 == 0 {
|
||||
return Err(FetchError::TemporarilyUnavailable);
|
||||
}
|
||||
return Err(FetchError::Timeout);
|
||||
}
|
||||
|
||||
let count = (rand_u32() % 4) as usize; // 0..=3
|
||||
Ok((0..count)
|
||||
.map(|i| format!("item-{}-{}", id, i))
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Tiny deterministic PRNG helper so we don't pull in a crate.
|
||||
fn rand_u32() -> u32 {
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
// In a real service you'd use a proper RNG; this is fine for tests.
|
||||
static SEED: AtomicU32 = AtomicU32::new(0xDEADBEEF);
|
||||
let s = SEED.fetch_add(6364136223846793005, Ordering::Relaxed);
|
||||
s.wrapping_mul(6364136223846793005).wrapping_add(1)
|
||||
}
|
||||
|
||||
// ─── Watcher task ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// A single watcher polls `mock_fetch` on a schedule and forwards results.
|
||||
///
|
||||
/// * Bounded channel capacity: **32** events. When full, `send()` awaits
|
||||
/// until the consumer drains space — this is the backpressure mechanism.
|
||||
/// * Repeated failures (>5 consecutive) mark the watcher unhealthy and it
|
||||
/// stops polling, sending a `HealthUpdate` before exiting.
|
||||
async fn watcher_task(
|
||||
id: u32,
|
||||
tx: mpsc::Sender<WatchEvent>,
|
||||
mut shutdown_rx: broadcast::Receiver<()>,
|
||||
) {
|
||||
let mut interval = interval(Duration::from_millis(80));
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
let mut consecutive_failures: u32 = 0;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
tracing::info!(watcher_id = id, "shutdown signal received");
|
||||
break;
|
||||
}
|
||||
_ = interval.tick() => {
|
||||
match mock_fetch(id).await {
|
||||
Ok(values) => {
|
||||
consecutive_failures = 0;
|
||||
for value in values {
|
||||
let item = WatchedItem { watcher_id: id, value };
|
||||
if tx.send(WatchEvent::Item(item)).await.is_err() {
|
||||
tracing::warn!(watcher_id = id, "channel closed, stopping");
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
consecutive_failures += 1;
|
||||
tracing::warn!(
|
||||
watcher_id = id,
|
||||
failures = consecutive_failures,
|
||||
error = %e,
|
||||
"poll failed"
|
||||
);
|
||||
if consecutive_failures > 5 {
|
||||
tracing::error!(watcher_id = id, "watcher marked unhealthy after {} consecutive failures", consecutive_failures);
|
||||
let _ = tx.send(WatchEvent::HealthUpdate { id, healthy: false }).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(watcher_id = id, "watcher task exited cleanly");
|
||||
}
|
||||
|
||||
// ─── Consumer task ────────────────────────────────────────────────────────────
|
||||
|
||||
/// Receives from all watcher channels (merged via `select_all`) and aggregates.
|
||||
///
|
||||
/// The merged stream automatically removes closed receivers, so when a watcher
|
||||
/// drops its sender the consumer adapts without explicit coordination.
|
||||
async fn consumer_task(
|
||||
mut receivers: Vec<mpsc::Receiver<WatchEvent>>,
|
||||
output_tx: mpsc::Sender<AggregatedOutput>,
|
||||
mut shutdown_rx: broadcast::Receiver<()>,
|
||||
) -> AggregatedOutput {
|
||||
let mut output = AggregatedOutput {
|
||||
total_items: 0,
|
||||
items_per_watcher: HashMap::new(),
|
||||
};
|
||||
|
||||
let mut stream = stream::select_all(
|
||||
receivers
|
||||
.drain(..)
|
||||
.map(|rx| rx.map(Ok::<_, Infallible>))
|
||||
.collect::<Vec<_>>(),
|
||||
);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
tracing::info!("consumer: shutdown signal received");
|
||||
break;
|
||||
}
|
||||
next = stream.next() => {
|
||||
match next {
|
||||
Some(Ok(WatchEvent::Item(item))) => {
|
||||
output.total_items += 1;
|
||||
*output.items_per_watcher.entry(item.watcher_id).or_insert(0) += 1;
|
||||
tracing::trace!(
|
||||
watcher_id = item.watcher_id,
|
||||
value = %item.value,
|
||||
"consumed item"
|
||||
);
|
||||
}
|
||||
Some(Ok(WatchEvent::HealthUpdate { id, healthy })) => {
|
||||
tracing::info!(watcher_id = id, healthy, "health update received");
|
||||
}
|
||||
Some(Err(_)) => unreachable!("Infallible"),
|
||||
None => {
|
||||
// All senders dropped — this happens on shutdown.
|
||||
tracing::info!("consumer: all watcher channels closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush final output before exiting.
|
||||
if output.total_items > 0 {
|
||||
let _ = output_tx.send(output.clone()).await;
|
||||
}
|
||||
|
||||
tracing::info!(total_items = output.total_items, "consumer exited");
|
||||
output
|
||||
}
|
||||
|
||||
// ─── WatcherManager ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Manages a dynamic set of watcher tasks with add/remove and clean shutdown.
|
||||
///
|
||||
/// Shared state is protected by `Arc<RwLock<WatcherSet>>` so that add/remove
|
||||
/// operations are race-free and concurrent with running watchers.
|
||||
#[derive(Debug)]
|
||||
struct WatcherSet {
|
||||
/// Live watcher entries keyed by their ID.
|
||||
watchers: HashMap<u32, WatcherEntry>,
|
||||
/// Current health status per watcher (true = healthy).
|
||||
health_status: HashMap<u32, bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct WatcherEntry {
|
||||
id: u32,
|
||||
/// Sender half of the watcher→consumer channel. Dropping this closes the
|
||||
/// receiver, causing the consumer's `select_all` stream to drop it.
|
||||
tx: mpsc::Sender<WatchEvent>,
|
||||
/// Join handle for the watcher task.
|
||||
handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct WatcherManager {
|
||||
inner: Arc<RwLock<WatcherSet>>,
|
||||
/// Broadcast channel used to signal all tasks to shut down.
|
||||
shutdown_tx: broadcast::Sender<()>,
|
||||
/// Sender for the final aggregated output.
|
||||
output_tx: mpsc::Sender<AggregatedOutput>,
|
||||
}
|
||||
|
||||
impl WatcherManager {
|
||||
/// Creates a new manager and spawns the consumer task.
|
||||
pub fn new() -> (Self, mpsc::Receiver<AggregatedOutput>) {
|
||||
let (output_tx, output_rx) = mpsc::channel(1);
|
||||
let (shutdown_tx, _) = broadcast::channel(1);
|
||||
|
||||
let manager = Self {
|
||||
inner: Arc::new(RwLock::new(WatcherSet {
|
||||
watchers: HashMap::new(),
|
||||
health_status: HashMap::new(),
|
||||
})),
|
||||
shutdown_tx: shutdown_tx.clone(),
|
||||
output_tx,
|
||||
};
|
||||
|
||||
// Spawn the consumer; it will receive a clone of `shutdown_tx` internally.
|
||||
tokio::spawn(manager.clone_consumer(shutdown_tx));
|
||||
|
||||
(manager, output_rx)
|
||||
}
|
||||
|
||||
/// Returns the number of currently active watchers.
|
||||
pub async fn watcher_count(&self) -> usize {
|
||||
self.inner.read().await.watchers.len()
|
||||
}
|
||||
|
||||
/// Returns the health status of a watcher, if it exists.
|
||||
pub async fn is_healthy(&self, id: u32) -> Option<bool> {
|
||||
self.inner.read().await.health_status.get(&id).copied()
|
||||
}
|
||||
|
||||
/// Adds a new watcher with the given ID. Returns `Ok(())` or an error if
|
||||
/// a watcher with that ID already exists.
|
||||
pub async fn add_watcher(&self, id: u32) -> Result<(), ServiceError> {
|
||||
let mut set = self.inner.write().await;
|
||||
|
||||
if set.watchers.contains_key(&id) {
|
||||
return Err(ServiceError::WatcherNotFound(id));
|
||||
}
|
||||
|
||||
// Bounded channel: capacity 32. When full, `send()` awaits until the
|
||||
// consumer drains space — this is our backpressure mechanism. If the
|
||||
// consumer is permanently slow, watchers will block on send rather than
|
||||
// buffering unboundedly in memory.
|
||||
let (tx, rx) = mpsc::channel(32);
|
||||
|
||||
let mut shutdown_rx = self.shutdown_tx.subscribe();
|
||||
let handle = tokio::spawn(watcher_task(id, tx.clone(), shutdown_rx));
|
||||
|
||||
set.watchers.insert(
|
||||
id,
|
||||
WatcherEntry { id, tx, handle },
|
||||
);
|
||||
set.health_status.insert(id, true);
|
||||
|
||||
// We need to give the consumer its receiver. Since the consumer was
|
||||
// spawned before any watchers existed, we store receivers inside the
|
||||
// shared state and the consumer reads them dynamically.
|
||||
//
|
||||
// Actually, the consumer uses `select_all` on a Vec it owns. We need
|
||||
// to push the new receiver into that Vec. We'll store receivers in the
|
||||
// shared state and have the consumer re-build its stream periodically.
|
||||
//
|
||||
// Simpler approach: store receivers in the shared state and have a
|
||||
// separate "receiver manager" that the consumer polls.
|
||||
drop(set);
|
||||
|
||||
// Push receiver into the consumer's collection via a side channel.
|
||||
// (In production you'd use a dedicated coordination mechanism.)
|
||||
// For this demo, we'll store receivers in the shared state and have
|
||||
// the consumer rebuild its select_all stream when notified.
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Removes a watcher by ID, dropping its sender so the task exits.
|
||||
pub async fn remove_watcher(&self, id: u32) -> Result<(), ServiceError> {
|
||||
let mut set = self.inner.write().await;
|
||||
|
||||
if let Some(entry) = set.watchers.remove(&id) {
|
||||
// Dropping `tx` closes the channel; the watcher task will see
|
||||
// `send()` fail and exit. We await the handle to ensure clean join.
|
||||
let id = entry.id;
|
||||
tokio::spawn(async move {
|
||||
if entry.handle.await.is_err() {
|
||||
tracing::warn!(watcher_id = id, "watcher task panicked");
|
||||
}
|
||||
});
|
||||
set.health_status.remove(&id);
|
||||
} else {
|
||||
return Err(ServiceError::WatcherNotFound(id));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Signals all tasks to shut down and waits for them to finish.
|
||||
pub async fn shutdown(&self) -> Result<(), ServiceError> {
|
||||
// Broadcast the shutdown signal. All watchers and the consumer listen
|
||||
// to this channel and will exit their select! loops.
|
||||
let _ = self.shutdown_tx.send(());
|
||||
|
||||
// Drop all senders to ensure receivers close even if a task misses the
|
||||
// broadcast (e.g., it's blocked on `send()`).
|
||||
let set = self.inner.read().await;
|
||||
for entry in set.watchers.values() {
|
||||
// tx is dropped here when `set` is dropped at end of scope.
|
||||
}
|
||||
drop(set);
|
||||
|
||||
// Give tasks a moment to notice the signal and exit.
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
|
||||
// Join all handles.
|
||||
let set = self.inner.read().await;
|
||||
for entry in set.watchers.values() {
|
||||
let handle = entry.handle.clone();
|
||||
let id = entry.id;
|
||||
tokio::spawn(async move {
|
||||
if handle.await.is_err() {
|
||||
tracing::warn!(watcher_id = id, "watcher task panicked during shutdown");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clone for sharing between manager and consumer.
|
||||
fn clone_consumer(&self, shutdown_tx: broadcast::Sender<()>) -> Self {
|
||||
Self {
|
||||
inner: self.inner.clone(),
|
||||
shutdown_tx,
|
||||
output_tx: self.output_tx.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Internal consumer loop that rebuilds its receiver stream when watchers
|
||||
/// are added or removed.
|
||||
async fn run_consumer(self) {
|
||||
// We store receivers in the shared state alongside entries.
|
||||
// Re-implementing with a cleaner design: store (tx, rx) pairs in the set.
|
||||
// For this demo, we'll use a simpler approach with a dedicated receiver vec.
|
||||
|
||||
// Actually, let me restructure: the consumer will poll the shared state
|
||||
// to collect current receivers and rebuild select_all each iteration.
|
||||
// This is slightly less efficient but correct for a demo.
|
||||
|
||||
let mut receivers: Vec<mpsc::Receiver<WatchEvent>> = Vec::new();
|
||||
let mut last_count = 0;
|
||||
|
||||
loop {
|
||||
// Collect current receivers from shared state.
|
||||
let set = self.inner.read().await;
|
||||
let new_receivers: Vec<_> = set
|
||||
.watchers
|
||||
.values()
|
||||
.map(|entry| {
|
||||
let (tx, rx) = mpsc::channel(32);
|
||||
// We can't actually create new receivers for existing senders.
|
||||
// This approach won't work.
|
||||
drop(tx);
|
||||
rx
|
||||
})
|
||||
.collect();
|
||||
drop(set);
|
||||
|
||||
// This approach is flawed. Let me use a different design.
|
||||
break;
|
||||
}
|
||||
|
||||
// ── Revised consumer using a receiver coordination channel ─────────
|
||||
// Each watcher stores its receiver in the shared state. The consumer
|
||||
// reads receivers directly from there using select_all rebuilt on
|
||||
// changes. We use a oneshot-per-change notification.
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Revised architecture (cleaner) ──────────────────────────────────────────
|
||||
//
|
||||
// To avoid the dynamic-receiver problem, we use a single shared mpsc channel
|
||||
// from all watchers to the consumer. Each watcher clones the sender before
|
||||
// spawning. The consumer reads from one receiver. This is the standard
|
||||
// tokio pattern and avoids compile-time select! limitations entirely.
|
||||
//
|
||||
// Per-watcher health is tracked via an Arc<AtomicBool> stored in the shared
|
||||
// state, updated by the watcher task itself.
|
||||
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
/// Revised WatcherSet that stores health atomics for per-watcher status.
|
||||
#[derive(Debug)]
|
||||
struct WatcherSetV2 {
|
||||
watchers: HashMap<u32, WatcherEntry>,
|
||||
health_flags: HashMap<u32, Arc<AtomicBool>>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct WatcherEntryV2 {
|
||||
id: u32,
|
||||
handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
/// Revised manager using a single shared consumer channel.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WatcherManagerV2 {
|
||||
inner: Arc<RwLock<WatcherSetV2>>,
|
||||
/// Single shared sender — cloned for each watcher. Bounded at 32; when
|
||||
/// full, `send()` blocks (backpressure). Dropping all senders closes the
|
||||
/// channel, causing the consumer to exit.
|
||||
shared_tx: mpsc::Sender<WatchedItem>,
|
||||
shutdown_tx: broadcast::Sender<()>,
|
||||
output_tx: mpsc::Sender<AggregatedOutput>,
|
||||
}
|
||||
|
||||
impl WatcherManagerV2 {
|
||||
/// Creates a new manager and spawns the consumer task.
|
||||
pub fn new() -> (Self, mpsc::Receiver<AggregatedOutput>) {
|
||||
let (shared_tx, shared_rx) = mpsc::channel(32); // backpressure bound: 32
|
||||
let (output_tx, output_rx) = mpsc::channel(1);
|
||||
let (shutdown_tx, _) = broadcast::channel(1);
|
||||
|
||||
let manager = Self {
|
||||
inner: Arc::new(RwLock::new(WatcherSetV2 {
|
||||
watchers: HashMap::new(),
|
||||
health_flags: HashMap::new(),
|
||||
})),
|
||||
shared_tx: shared_tx.clone(),
|
||||
shutdown_tx: shutdown_tx.clone(),
|
||||
output_tx,
|
||||
};
|
||||
|
||||
// Spawn the consumer.
|
||||
tokio::spawn(manager.clone().run_consumer(shared_rx, shutdown_tx));
|
||||
|
||||
(manager, output_rx)
|
||||
}
|
||||
|
||||
/// Returns the number of currently active watchers.
|
||||
pub async fn watcher_count(&self) -> usize {
|
||||
self.inner.read().await.watchers.len()
|
||||
}
|
||||
|
||||
/// Returns the health status of a watcher, if it exists.
|
||||
pub async fn is_healthy(&self, id: u32) -> Option<bool> {
|
||||
self.inner
|
||||
.read()
|
||||
.await
|
||||
.health_flags
|
||||
.get(&id)
|
||||
.map(|flag| flag.load(Ordering::Relaxed))
|
||||
}
|
||||
|
||||
/// Adds a new watcher with the given ID. Returns `Ok(())` or an error if
|
||||
/// a watcher with that ID already exists.
|
||||
pub async fn add_watcher(&self, id: u32) -> Result<(), ServiceError> {
|
||||
let mut set = self.inner.write().await;
|
||||
|
||||
if set.watchers.contains_key(&id) {
|
||||
return Err(ServiceError::WatcherNotFound(id));
|
||||
}
|
||||
|
||||
let health_flag = Arc::new(AtomicBool::new(true));
|
||||
set.health_flags.insert(id, health_flag.clone());
|
||||
|
||||
let mut shutdown_rx = self.shutdown_tx.subscribe();
|
||||
let tx = self.shared_tx.clone(); // clone sender for this watcher
|
||||
|
||||
let handle = tokio::spawn(watcher_task_v2(
|
||||
id,
|
||||
tx,
|
||||
health_flag,
|
||||
shutdown_rx,
|
||||
));
|
||||
|
||||
set.watchers.insert(id, WatcherEntryV2 { id, handle });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Removes a watcher by ID, dropping its sender so the task exits.
|
||||
pub async fn remove_watcher(&self, id: u32) -> Result<(), ServiceError> {
|
||||
let mut set = self.inner.write().await;
|
||||
|
||||
if !set.watchers.contains_key(&id) {
|
||||
return Err(ServiceError::WatcherNotFound(id));
|
||||
}
|
||||
|
||||
let entry = set.watchers.remove(&id).unwrap();
|
||||
set.health_flags.remove(&id);
|
||||
@@ -0,0 +1,416 @@
|
||||
"""
|
||||
Async TTS Job Pipeline — stdlib-only, Python 3.11+
|
||||
|
||||
A bounded-concurrency job pipeline with retry, cancellation, backpressure,
|
||||
and lifecycle callbacks. Includes a self-testing ``main()``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import random
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Awaitable, Callable
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# Mock synthesizer
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
async def mock_synthesize(text: str) -> bytes:
|
||||
"""Simulate TTS synthesis.
|
||||
|
||||
* Sleeps a random 50–300 ms.
|
||||
* Returns ``b"\\x00" * len(text)`` on success.
|
||||
* Raises ``RuntimeError`` ~10 % of the time to exercise retry logic.
|
||||
"""
|
||||
await asyncio.sleep(random.uniform(0.05, 0.30))
|
||||
if random.random() < 0.10:
|
||||
raise RuntimeError(f"synth-error on {len(text)}-char payload")
|
||||
return b"\x00" * len(text)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# Job model
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
class JobState(str, Enum):
|
||||
QUEUED = "queued"
|
||||
STARTED = "started"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Job:
|
||||
job_id: str
|
||||
text: str
|
||||
voice: str
|
||||
state: JobState = JobState.QUEUED
|
||||
attempt: int = 0
|
||||
result: bytes | None = None
|
||||
error: str | None = None
|
||||
_cancelled: bool = False # internal flag, protected by pipeline lock
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# Pipeline
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
class TTSPipeline:
|
||||
"""Async TTS job pipeline with bounded concurrency and retry."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_concurrency: int = 4,
|
||||
max_queue_size: int = 100,
|
||||
) -> None:
|
||||
self._max_concurrency = max_concurrency
|
||||
self._max_queue_size = max_queue_size
|
||||
|
||||
# Bounded queue enforces backpressure on submit()
|
||||
self._queue: asyncio.Queue[Job] = asyncio.Queue(maxsize=max_queue_size)
|
||||
|
||||
# Limits how many jobs are synthesising simultaneously
|
||||
self._semaphore = asyncio.Semaphore(max_concurrency)
|
||||
|
||||
# All jobs ever submitted (for drain / lookup)
|
||||
self._jobs: dict[str, Job] = {}
|
||||
self._job_counter: int = 0
|
||||
|
||||
# Callbacks registered via on_event()
|
||||
self._callbacks: list[Callable[[str, str], Awaitable[None]]] = []
|
||||
|
||||
# Concurrency-tracking helpers (for tests)
|
||||
self._active_count: int = 0
|
||||
self._active_lock = asyncio.Lock()
|
||||
|
||||
# Internal lock serialises mutations to shared state
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# Background worker that pulls from the queue and dispatches tasks
|
||||
self._worker_task: asyncio.Task[None] | None = None
|
||||
|
||||
# ── lifecycle ────────────────────────────────────────────────────
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Spawn the background worker. Call once before submitting."""
|
||||
if self._worker_task is None:
|
||||
self._worker_task = asyncio.create_task(self._worker_loop())
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Gracefully shut down the worker (for cleanup in tests)."""
|
||||
if self._worker_task is not None:
|
||||
self._worker_task.cancel()
|
||||
try:
|
||||
await self._worker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._worker_task = None
|
||||
|
||||
# ── public API ───────────────────────────────────────────────────
|
||||
|
||||
async def submit(self, text: str, voice: str) -> str:
|
||||
"""Enqueue a job and return its id immediately (non-blocking).
|
||||
|
||||
Raises ``RuntimeError`` with a clear message when the queue is full
|
||||
(backpressure — does *not* block or grow unbounded).
|
||||
"""
|
||||
async with self._lock:
|
||||
self._job_counter += 1
|
||||
job_id = f"job-{self._job_counter:04d}"
|
||||
job = Job(job_id=job_id, text=text, voice=voice)
|
||||
self._jobs[job_id] = job
|
||||
|
||||
# Try to enqueue; on failure roll back the registration
|
||||
try:
|
||||
self._queue.put_nowait(job)
|
||||
except asyncio.QueueFull:
|
||||
async with self._lock:
|
||||
del self._jobs[job_id]
|
||||
raise RuntimeError(
|
||||
f"Queue full (max {self._max_queue_size}); job rejected: {job_id}"
|
||||
) from None
|
||||
|
||||
await self._notify(job, "queued")
|
||||
return job_id
|
||||
|
||||
async def drain(self) -> None:
|
||||
"""Block until every submitted job has reached a terminal state
|
||||
(completed / failed / cancelled) and the internal queue is empty.
|
||||
|
||||
Jobs submitted *after* this call returns are **not** waited for.
|
||||
"""
|
||||
while True:
|
||||
async with self._lock:
|
||||
total = len(self._jobs)
|
||||
done = sum(
|
||||
1 for j in self._jobs.values()
|
||||
if j.state in (JobState.COMPLETED, JobState.FAILED, JobState.CANCELLED)
|
||||
)
|
||||
if total > 0 and total == done and self._queue.empty():
|
||||
return
|
||||
# Yield control so workers can make progress
|
||||
await asyncio.sleep(0.005)
|
||||
|
||||
async def cancel(self, job_id: str) -> None:
|
||||
"""Cancel a job.
|
||||
|
||||
* Queued jobs are removed from processing immediately.
|
||||
* In-flight jobs are marked for cancellation and checked after
|
||||
synthesis completes (or on the next retry boundary).
|
||||
"""
|
||||
async with self._lock:
|
||||
job = self._jobs.get(job_id)
|
||||
if job is None:
|
||||
raise KeyError(f"Unknown job: {job_id}")
|
||||
if job.state in (JobState.COMPLETED, JobState.FAILED):
|
||||
return # already terminal — no-op
|
||||
|
||||
job._cancelled = True
|
||||
if job.state == JobState.QUEUED:
|
||||
# Already in the queue; the worker will see the flag and skip it.
|
||||
job.state = JobState.CANCELLED
|
||||
# If in-flight, the worker checks _cancelled after synthesis.
|
||||
|
||||
await self._notify(job, "cancelled")
|
||||
|
||||
def on_event(self, callback: Callable[[str, str], Awaitable[None]]) -> None:
|
||||
"""Register a lifecycle-event callback.
|
||||
|
||||
Events emitted: ``queued``, ``started``, ``completed``,
|
||||
``failed``, ``cancelled``.
|
||||
"""
|
||||
self._callbacks.append(callback)
|
||||
|
||||
# ── internal ─────────────────────────────────────────────────────
|
||||
|
||||
async def _notify(self, job: Job, event: str) -> None:
|
||||
for cb in self._callbacks:
|
||||
try:
|
||||
await cb(job.job_id, event)
|
||||
except Exception:
|
||||
pass # one bad callback must not break the pipeline
|
||||
|
||||
async def _worker_loop(self) -> None:
|
||||
"""Single dispatcher: pulls jobs from the queue and spawns a task
|
||||
for each. The task itself acquires the concurrency semaphore."""
|
||||
while True:
|
||||
job = await self._queue.get()
|
||||
# task_done is called by _process_job when it finishes (including retries)
|
||||
asyncio.create_task(self._process_job(job))
|
||||
|
||||
async def _process_job(self, job: Job) -> None:
|
||||
"""Run one job with bounded-concurrency synthesis and retry logic."""
|
||||
max_attempts = 3
|
||||
backoffs = [0.1, 0.2, 0.4] # exponential: 0.1 → 0.2 → 0.4 s
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
job.attempt = attempt + 1
|
||||
|
||||
# Respect cancellation before even trying to acquire a slot
|
||||
if job._cancelled and job.state == JobState.CANCELLED:
|
||||
break
|
||||
|
||||
# ── Acquire a concurrency slot for the actual synthesis ──
|
||||
async with self._semaphore:
|
||||
# Double-check cancellation after waiting for a slot
|
||||
if job._cancelled and job.state == JobState.CANCELLED:
|
||||
break
|
||||
|
||||
async with self._active_lock:
|
||||
self._active_count += 1
|
||||
|
||||
job.state = JobState.STARTED
|
||||
await self._notify(job, "started")
|
||||
|
||||
try:
|
||||
result = await mock_synthesize(job.text)
|
||||
|
||||
if job._cancelled and job.state == JobState.CANCELLED:
|
||||
job.state = JobState.CANCELLED
|
||||
else:
|
||||
job.result = result
|
||||
job.state = JobState.COMPLETED
|
||||
await self._notify(job, "completed")
|
||||
|
||||
except Exception as exc:
|
||||
if job._cancelled and job.state == JobState.CANCELLED:
|
||||
job.state = JobState.CANCELLED
|
||||
elif attempt < max_attempts - 1:
|
||||
# Will fall through to backoff below (outside semaphore)
|
||||
pass
|
||||
else:
|
||||
job.error = str(exc)
|
||||
job.state = JobState.FAILED
|
||||
await self._notify(job, "failed")
|
||||
|
||||
finally:
|
||||
async with self._active_lock:
|
||||
self._active_count -= 1
|
||||
|
||||
# ── Retry backoff happens *outside* the semaphore so it
|
||||
# doesn't hold a concurrency slot while sleeping ──────
|
||||
if job.state == JobState.FAILED and attempt < max_attempts - 1:
|
||||
await asyncio.sleep(backoffs[attempt])
|
||||
else:
|
||||
break # completed, cancelled, or exhausted retries
|
||||
|
||||
self._queue.task_done()
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
# Tests
|
||||
# ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
async def main() -> None:
|
||||
print("=" * 60)
|
||||
print("Async TTS Pipeline — Self-Test")
|
||||
print("=" * 60)
|
||||
|
||||
# ── Helper: collect lifecycle events ───────────────────────────
|
||||
events: list[tuple[str, str]] = []
|
||||
|
||||
async def record(event_id: str, event_type: str) -> None:
|
||||
events.append((event_id, event_type))
|
||||
|
||||
# ── Test A: bounded concurrency ────────────────────────────────
|
||||
print("\n[Test A] Bounded concurrency (50 jobs, max 4 simultaneous)")
|
||||
print("-" * 50)
|
||||
|
||||
pipeline_a = TTSPipeline(max_concurrency=4, max_queue_size=100)
|
||||
await pipeline_a.start()
|
||||
pipeline_a.on_event(record)
|
||||
|
||||
peak_concurrency = 0
|
||||
current_active = 0
|
||||
active_lock = asyncio.Lock()
|
||||
|
||||
async def track_events(job_id: str, event_type: str) -> None:
|
||||
nonlocal peak_concurrency, current_active
|
||||
if event_type == "started":
|
||||
async with active_lock:
|
||||
current_active += 1
|
||||
peak_concurrency = max(peak_concurrency, current_active)
|
||||
elif event_type in ("completed", "failed", "cancelled"):
|
||||
async with active_lock:
|
||||
current_active -= 1
|
||||
|
||||
pipeline_a.on_event(track_events)
|
||||
|
||||
ids_a = [await pipeline_a.submit(f"text-{i}", "mock-voice") for i in range(50)]
|
||||
await pipeline_a.drain()
|
||||
|
||||
print(f" Submitted : {len(ids_a)}")
|
||||
print(f" Peak concurrency: {peak_concurrency}")
|
||||
assert peak_concurrency <= 4, f"Concurrency exceeded! Peak was {peak_concurrency}"
|
||||
print(" ✓ Concurrency never exceeded 4")
|
||||
|
||||
# Verify no duplicates / drops
|
||||
states = {jid: next((j.state for j in pipeline_a._jobs.values() if j.job_id == jid))
|
||||
for jid in ids_a}
|
||||
terminal = sum(1 for s in states.values() if s in (JobState.COMPLETED, JobState.FAILED))
|
||||
assert terminal == 50, f"Expected 50 terminal jobs, got {terminal}"
|
||||
print(f" ✓ All 50 jobs reached a terminal state")
|
||||
|
||||
await pipeline_a.stop()
|
||||
|
||||
# ── Test B: backpressure (queue cap = 100) ─────────────────────
|
||||
print("\n[Test B] Backpressure — 101st job must be rejected")
|
||||
print("-" * 50)
|
||||
|
||||
pipeline_b = TTSPipeline(max_concurrency=1, max_queue_size=100)
|
||||
await pipeline_b.start()
|
||||
|
||||
# Fill the queue to capacity by submitting 100 fast jobs
|
||||
ids_b = []
|
||||
for i in range(100):
|
||||
jid = await pipeline_b.submit(f"fill-{i}", "voice")
|
||||
ids_b.append(jid)
|
||||
|
||||
# The 101st must raise immediately (queue is full)
|
||||
rejected = False
|
||||
try:
|
||||
await pipeline_b.submit("overflow", "voice")
|
||||
except RuntimeError as exc:
|
||||
rejected = True
|
||||
print(f" Rejection message: {exc}")
|
||||
|
||||
assert rejected, "Expected RuntimeError for over-full queue"
|
||||
print(" ✓ 101st job was correctly rejected")
|
||||
|
||||
await pipeline_b.drain()
|
||||
print(" ✓ Drain completed after backpressure test")
|
||||
await pipeline_b.stop()
|
||||
|
||||
# ── Test C: cancellation mid-flight + pool health ──────────────
|
||||
print("\n[Test C] Cancel mid-flight, drain, verify pool health")
|
||||
print("-" * 50)
|
||||
|
||||
pipeline_c = TTSPipeline(max_concurrency=4, max_queue_size=100)
|
||||
await pipeline_c.start()
|
||||
|
||||
# Patch mock_synthesize to be slow so we can cancel mid-flight
|
||||
original_synthesize = mock_synthesize
|
||||
|
||||
async def slow_synthesize(text: str) -> bytes:
|
||||
await asyncio.sleep(2.0) # deliberately slow
|
||||
return b"\x00" * len(text)
|
||||
|
||||
import __main__
|
||||
__main__.mock_synthesize = slow_synthesize # temporarily replace
|
||||
|
||||
ids_c = [await pipeline_c.submit(f"cancel-me-{i}", "voice") for i in range(20)]
|
||||
await asyncio.sleep(0.05) # let a few start
|
||||
|
||||
# Cancel one that should be in-flight
|
||||
await pipeline_c.cancel(ids_c[0])
|
||||
print(f" Cancelled job {ids_c[0]}")
|
||||
|
||||
# Also cancel one that is still queued
|
||||
await pipeline_c.cancel(ids_c[1])
|
||||
print(f" Cancelled job {ids_c[1]} (queued)")
|
||||
|
||||
# Restore original synthesizer so remaining jobs finish quickly
|
||||
__main__.mock_synthesize = original_synthesize
|
||||
|
||||
await pipeline_c.drain()
|
||||
|
||||
# Verify no duplicates: every submitted id appears exactly once
|
||||
seen_ids = [e[0] for e in events if e[1] == "queued"]
|
||||
assert len(seen_ids) == len(set(seen_ids)), "Duplicate job IDs detected!"
|
||||
print(" ✓ No duplicate jobs")
|
||||
|
||||
# Verify no drops: every submitted id has a terminal state
|
||||
final_states: dict[str, JobState] = {}
|
||||
for jid in ids_c:
|
||||
job = pipeline_c._jobs[jid]
|
||||
final_states[jid] = job.state
|
||||
|
||||
completed = sum(1 for s in final_states.values() if s == JobState.COMPLETED)
|
||||
cancelled = sum(1 for s in final_states.values() if s == JobState.CANCELLED)
|
||||
failed = sum(1 for s in final_states.values() if s == JobState.FAILED)
|
||||
print(f" Results: {completed} completed, {cancelled} cancelled, {failed} failed")
|
||||
assert completed + cancelled + failed == 20, "Jobs were dropped!"
|
||||
print(" ✓ No jobs dropped or duplicated")
|
||||
|
||||
# Pool must still be usable after drain + cancel
|
||||
jid_new = await pipeline_c.submit("post-drain-job", "voice")
|
||||
assert jid_new.startswith("job-"), "Pipeline unusable after drain"
|
||||
await pipeline_c.drain()
|
||||
assert pipeline_c._jobs[jid_new].state == JobState.COMPLETED
|
||||
print(" ✓ Pipeline still usable after drain + cancellation")
|
||||
|
||||
await pipeline_c.stop()
|
||||
|
||||
# ── Summary ────────────────────────────────────────────────────
|
||||
print("\n" + "=" * 60)
|
||||
print("All tests passed ✓")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,649 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Twitch EventSub → Mock Discord Webhook Bridge
|
||||
==============================================
|
||||
Receives signed Twitch-style webhook events on POST /webhook,
|
||||
verifies HMAC signatures, deduplicates by event_id, rate-limits
|
||||
outbound Discord forwards, and retries on 429 backoff.
|
||||
|
||||
Run: python webhook_bridge.py [--secret <key>] [--port <port>]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import http.server
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import queue
|
||||
import secrets
|
||||
import socketserver
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SECRET = os.environ.get("WEBHOOK_SECRET", secrets.token_hex(16))
|
||||
PORT = int(os.environ.get("WEBHOOK_PORT", "8080"))
|
||||
RATE_LIMIT = 5 # Discord forwards per second
|
||||
IDEMPOTENCY_TTL = 300 # seconds
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||||
datefmt="%H:%M:%S",
|
||||
)
|
||||
log = logging.getLogger("bridge")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mock Discord Webhook Client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class DiscordMock:
|
||||
"""Simulates a downstream Discord webhook endpoint.
|
||||
|
||||
Succeeds ~95 % of the time. Occasionally returns HTTP 429 with a
|
||||
``Retry-After`` header to exercise the backoff path.
|
||||
"""
|
||||
|
||||
def __init__(self, fail_rate: float = 0.05) -> None:
|
||||
self.fail_rate = fail_rate
|
||||
self.forward_count: int = 0
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def send(self, payload: dict[str, Any]) -> tuple[int, dict[str, str]]:
|
||||
"""Return (status_code, headers) simulating Discord's response."""
|
||||
async with self._lock:
|
||||
self.forward_count += 1
|
||||
|
||||
import random
|
||||
|
||||
if random.random() < self.fail_rate:
|
||||
retry_after = random.choice([1, 2, 3])
|
||||
log.warning("Discord returned 429 (Retry-After=%ds)", retry_after)
|
||||
return 429, {"Retry-After": str(retry_after)}
|
||||
|
||||
log.info("Discord forward #%d succeeded", self.forward_count)
|
||||
return 200, {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Idempotency Store
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IdempotencyStore:
|
||||
"""In-memory deduplication store keyed by ``event_id``.
|
||||
|
||||
Entries older than ``IDEMPOTENCY_TTL`` seconds are evicted on access.
|
||||
"""
|
||||
|
||||
def __init__(self, ttl: int = IDEMPOTENCY_TTL) -> None:
|
||||
self._ttl = ttl
|
||||
self._store: dict[str, float] = {} # event_id -> timestamp
|
||||
|
||||
def is_duplicate(self, event_id: str) -> bool:
|
||||
"""Return True if *event_id* was seen within the TTL window."""
|
||||
now = time.monotonic()
|
||||
# Evict stale entries
|
||||
self._store = {k: v for k, v in self._store.items() if now - v < self._ttl}
|
||||
if event_id in self._store:
|
||||
return True
|
||||
self._store[event_id] = now
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Token-Bucket Rate Limiter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TokenBucket:
|
||||
"""Non-blocking token-bucket rate limiter.
|
||||
|
||||
Capacity and refill rate are both *rate* tokens/second.
|
||||
``try_acquire()`` returns True immediately if a token is available,
|
||||
otherwise False. The caller should ``await self.wait()`` to block
|
||||
until a token is available when the bucket is empty.
|
||||
"""
|
||||
|
||||
def __init__(self, rate: float) -> None:
|
||||
self._rate = rate
|
||||
self._tokens = float(rate)
|
||||
self._last = time.monotonic()
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
def _refill(self) -> None:
|
||||
now = time.monotonic()
|
||||
elapsed = now - self._last
|
||||
self._tokens = min(float(self._rate), self._tokens + elapsed * self._rate)
|
||||
self._last = now
|
||||
|
||||
async def try_acquire(self) -> bool:
|
||||
"""Try to consume a token without blocking. Returns True on success."""
|
||||
async with self._lock:
|
||||
self._refill()
|
||||
if self._tokens >= 1.0:
|
||||
self._tokens -= 1.0
|
||||
return True
|
||||
return False
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""Block until a token becomes available, then consume it."""
|
||||
while True:
|
||||
if await self.try_acquire():
|
||||
return
|
||||
# Sleep just enough for one token to refill
|
||||
await asyncio.sleep(1.0 / self._rate)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Webhook Bridge Logic
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class EventContext:
|
||||
"""Holds the parsed event and its raw body for signing."""
|
||||
|
||||
event_id: str
|
||||
event_type: str
|
||||
data: Any
|
||||
raw_body: bytes
|
||||
|
||||
|
||||
class WebhookBridge:
|
||||
"""Core bridge: verify → deduplicate → rate-limit → forward."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
secret: str,
|
||||
discord: DiscordMock,
|
||||
store: IdempotencyStore | None = None,
|
||||
limiter: TokenBucket | None = None,
|
||||
) -> None:
|
||||
self.secret = secret
|
||||
self.discord = discord
|
||||
self.store = store or IdempotencyStore()
|
||||
self.limiter = limiter or TokenBucket(RATE_LIMIT)
|
||||
# Track successful forwards for test assertions
|
||||
self.successful_forwards: list[dict] = []
|
||||
|
||||
# -- HMAC verification --------------------------------------------------
|
||||
|
||||
def verify_signature(self, body: bytes, signature: str) -> bool:
|
||||
"""Constant-time HMAC-SHA256 verification."""
|
||||
expected = hmac.new(
|
||||
self.secret.encode(), body, hashlib.sha256
|
||||
).hexdigest()
|
||||
return hmac.compare_digest(expected, signature)
|
||||
|
||||
# -- Event parsing ------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def parse_event(body: bytes) -> EventContext | None:
|
||||
"""Parse and validate the JSON payload. Returns None on failure."""
|
||||
try:
|
||||
data = json.loads(body)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError) as exc:
|
||||
log.error("Malformed JSON body: %s", exc)
|
||||
return None
|
||||
|
||||
if not isinstance(data, dict):
|
||||
log.error("JSON body is not an object")
|
||||
return None
|
||||
|
||||
event_id = data.get("event_id")
|
||||
event_type = data.get("type")
|
||||
if not event_id or not isinstance(event_id, str):
|
||||
log.error("Missing or invalid 'event_id'")
|
||||
return None
|
||||
if not event_type or not isinstance(event_type, str):
|
||||
log.error("Missing or invalid 'type'")
|
||||
return None
|
||||
|
||||
return EventContext(
|
||||
event_id=event_id,
|
||||
event_type=event_type,
|
||||
data=data.get("data", {}),
|
||||
raw_body=body,
|
||||
)
|
||||
|
||||
# -- Forward pipeline ---------------------------------------------------
|
||||
|
||||
async def forward(self, ctx: EventContext) -> int:
|
||||
"""Run the full pipeline. Returns HTTP status code to send back."""
|
||||
|
||||
# 1. Idempotency check
|
||||
if self.store.is_duplicate(ctx.event_id):
|
||||
log.info("Duplicate event_id=%s — skipping", ctx.event_id)
|
||||
return 200
|
||||
|
||||
# 2. Rate-limit: acquire a token (may block briefly)
|
||||
await self.limiter.wait()
|
||||
|
||||
# 3. Build Discord payload
|
||||
discord_payload = {
|
||||
"content": (
|
||||
f"📨 Twitch Event [{ctx.event_type}]\n"
|
||||
f"ID: {ctx.event_id}\n"
|
||||
f"Data: {json.dumps(ctx.data, default=str)}"
|
||||
),
|
||||
}
|
||||
|
||||
# 4. Send to Discord with single-shot 429 backoff retry
|
||||
status, headers = await self.discord.send(discord_payload)
|
||||
|
||||
if status == 429:
|
||||
retry_after = int(headers.get("Retry-After", "1"))
|
||||
log.warning(
|
||||
"Discord 429 on event_id=%s — backing off %ds",
|
||||
ctx.event_id,
|
||||
retry_after,
|
||||
)
|
||||
await asyncio.sleep(retry_after)
|
||||
# Retry once
|
||||
status, _ = await self.discord.send(discord_payload)
|
||||
|
||||
if status == 200:
|
||||
self.successful_forwards.append(discord_payload)
|
||||
log.info("Forwarded event_id=%s ✓", ctx.event_id)
|
||||
else:
|
||||
log.error("Discord returned %d for event_id=%s", status, ctx.event_id)
|
||||
|
||||
return 200 if status == 200 else 502
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async HTTP Server (pure stdlib, no third-party deps)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _AsyncHandler:
|
||||
"""Minimal async HTTP request handler using asyncio streams."""
|
||||
|
||||
def __init__(self, bridge: WebhookBridge) -> None:
|
||||
self.bridge = bridge
|
||||
|
||||
async def handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
||||
try:
|
||||
# Read request line
|
||||
request_line = await asyncio.wait_for(reader.readline(), timeout=10.0)
|
||||
if not request_line:
|
||||
return
|
||||
parts = request_line.decode().strip().split()
|
||||
if len(parts) < 2:
|
||||
self._send(writer, 400, "Bad Request")
|
||||
return
|
||||
|
||||
method, path = parts[0], parts[1]
|
||||
|
||||
# Read headers
|
||||
headers: dict[str, str] = {}
|
||||
while True:
|
||||
line = await asyncio.wait_for(reader.readline(), timeout=10.0)
|
||||
if line in (b"\r\n", b"\n", b""):
|
||||
break
|
||||
decoded = line.decode().strip()
|
||||
if ":" in decoded:
|
||||
key, val = decoded.split(":", 1)
|
||||
headers[key.strip().lower()] = val.strip()
|
||||
|
||||
# Read body
|
||||
body = b""
|
||||
content_length = int(headers.get("content-length", "0"))
|
||||
if content_length > 0:
|
||||
body = await asyncio.wait_for(reader.readexactly(content_length), timeout=10.0)
|
||||
|
||||
if path == "/webhook" and method == "POST":
|
||||
await self._handle_webhook(reader, writer, body, headers)
|
||||
else:
|
||||
self._send(writer, 404, "Not Found")
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
self._send(writer, 408, "Request Timeout")
|
||||
except Exception as exc:
|
||||
log.exception("Unhandled error in request handler")
|
||||
self._send(writer, 500, "Internal Server Error")
|
||||
finally:
|
||||
writer.close()
|
||||
try:
|
||||
await writer.wait_closed()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _handle_webhook(
|
||||
self, reader, writer, body: bytes, headers: dict[str, str]
|
||||
) -> None:
|
||||
# Signature verification
|
||||
signature = headers.get("x-signature")
|
||||
if not signature:
|
||||
log.warning("Missing X-Signature header")
|
||||
self._send(writer, 401, "Unauthorized")
|
||||
return
|
||||
|
||||
if not self.bridge.verify_signature(body, signature):
|
||||
log.warning("HMAC signature mismatch")
|
||||
self._send(writer, 401, "Unauthorized")
|
||||
return
|
||||
|
||||
# Parse event
|
||||
ctx = self.bridge.parse_event(body)
|
||||
if ctx is None:
|
||||
self._send(writer, 400, "Bad Request — invalid payload")
|
||||
return
|
||||
|
||||
log.info("Received event_id=%s type=%s", ctx.event_id, ctx.event_type)
|
||||
|
||||
# Forward (async pipeline with rate limiting & retry)
|
||||
status = await self.bridge.forward(ctx)
|
||||
self._send(writer, status, "OK" if status == 200 else "Error")
|
||||
|
||||
@staticmethod
|
||||
def _send(writer: asyncio.StreamWriter, status: int, reason: str) -> None:
|
||||
body = f"{status} {reason}\n".encode()
|
||||
response = (
|
||||
b"HTTP/1.1 " + f"{status} {reason}".encode() + b"\r\n"
|
||||
b"Content-Type: text/plain\r\n"
|
||||
b"Content-Length: " + str(len(body)).encode() + b"\r\n"
|
||||
b"Connection: close\r\n"
|
||||
b"\r\n"
|
||||
)
|
||||
writer.write(response + body)
|
||||
|
||||
|
||||
async def run_server(bridge: WebhookBridge, host: str = "127.0.0.1", port: int = PORT) -> None:
|
||||
"""Start the async HTTP server."""
|
||||
server = await asyncio.start_server(
|
||||
lambda r, w: _AsyncHandler(bridge).handle(r, w), host, port
|
||||
)
|
||||
addr = server.sockets[0].getsockname()
|
||||
log.info("Bridge listening on http://%s:%d", addr[0], addr[1])
|
||||
async with server:
|
||||
await server.serve_forever()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Threaded HTTP Server (for test compatibility with urllib/http.client)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _ThreadedServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
|
||||
daemon_threads = True
|
||||
allow_reuse_address = True
|
||||
|
||||
|
||||
class _BridgeHandler(http.server.BaseHTTPRequestHandler):
|
||||
"""Synchronous wrapper around the async bridge for use with http.server."""
|
||||
|
||||
bridge: WebhookBridge = None # type: ignore[assignment]
|
||||
|
||||
def do_POST(self) -> None: # noqa: N802
|
||||
if self.path != "/webhook":
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
self.wfile.write(b"Not Found")
|
||||
return
|
||||
|
||||
body = self.rfile.read(int(self.headers.get("Content-Length", 0)))
|
||||
headers_dict = {k.lower(): v for k, v in self.headers.items()}
|
||||
|
||||
# Signature verification
|
||||
signature = headers_dict.get("x-signature")
|
||||
if not signature:
|
||||
self.send_response(401)
|
||||
self.end_headers()
|
||||
self.wfile.write(b"Unauthorized")
|
||||
return
|
||||
|
||||
if not self.bridge.verify_signature(body, signature):
|
||||
self.send_response(401)
|
||||
self.end_headers()
|
||||
self.wfile.write(b"Unauthorized")
|
||||
return
|
||||
|
||||
ctx = self.bridge.parse_event(body)
|
||||
if ctx is None:
|
||||
self.send_response(400)
|
||||
self.end_headers()
|
||||
self.wfile.write(b"Bad Request")
|
||||
return
|
||||
|
||||
# Run async forward pipeline in a thread-safe way
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
task = loop.run_until_complete(self.bridge.forward(ctx))
|
||||
self.send_response(task)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"OK" if task == 200 else b"Error")
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
def log_message(self, fmt: str, *args: Any) -> None: # noqa: ARG002
|
||||
log.info(fmt, *args)
|
||||
|
||||
|
||||
def start_threaded_server(bridge: WebhookBridge, port: int = PORT) -> _ThreadedServer:
|
||||
"""Start a threaded HTTP server in a background thread. Returns the server."""
|
||||
_BridgeHandler.bridge = bridge
|
||||
srv = _ThreadedServer(("127.0.0.1", port), _BridgeHandler)
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
log.info("Threaded test server on http://%s:%d", "127.0.0.1", port)
|
||||
return srv
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers for tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def make_payload(event_id: str, event_type: str = "chat", data: dict | None = None) -> bytes:
|
||||
"""Create a signed webhook payload."""
|
||||
body_dict = {"event_id": event_id, "type": event_type, "data": data or {}}
|
||||
body = json.dumps(body_dict).encode()
|
||||
sig = hmac.new(SECRET.encode(), body, hashlib.sha256).hexdigest()
|
||||
return body, sig
|
||||
|
||||
|
||||
def post_webhook(url: str, body: bytes, signature: str) -> tuple[int, str]:
|
||||
"""Send a POST request and return (status_code, body_text)."""
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
|
||||
req = urllib.request.Request(
|
||||
url,
|
||||
data=body,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Signature": signature,
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
return resp.status, resp.read().decode()
|
||||
except urllib.error.HTTPError as exc:
|
||||
return exc.code, exc.read().decode()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_tests(port: int = 9099) -> int:
|
||||
"""Run all tests. Returns number of failures."""
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
|
||||
url = f"http://127.0.0.1:{port}/webhook"
|
||||
failures = 0
|
||||
|
||||
def assert_eq(actual, expected, msg: str) -> None:
|
||||
nonlocal failures
|
||||
if actual != expected:
|
||||
print(f" ✗ FAIL: {msg} — expected {expected!r}, got {actual!r}")
|
||||
failures += 1
|
||||
else:
|
||||
print(f" ✓ PASS: {msg}")
|
||||
|
||||
def assert_true(condition: bool, msg: str) -> None:
|
||||
nonlocal failures
|
||||
if not condition:
|
||||
print(f" ✗ FAIL: {msg}")
|
||||
failures += 1
|
||||
else:
|
||||
print(f" ✓ PASS: {msg}")
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Setup fresh bridge for tests
|
||||
# -----------------------------------------------------------------------
|
||||
discord = DiscordMock(fail_rate=0.0) # deterministic: no random failures
|
||||
bridge = WebhookBridge(secret=SECRET, discord=discord)
|
||||
server = start_threaded_server(bridge, port)
|
||||
time.sleep(0.2) # let server start
|
||||
|
||||
try:
|
||||
# -------------------------------------------------------------------
|
||||
print("\n=== Test (a): Valid signed request → forwarded, returns 200 ===")
|
||||
# -------------------------------------------------------------------
|
||||
before = discord.forward_count
|
||||
body, sig = make_payload("evt-a1", "chat", {"message": "hello"})
|
||||
status, text = post_webhook(url, body, sig)
|
||||
assert_eq(status, 200, "HTTP status")
|
||||
assert_true(discord.forward_count == before + 1, "Discord forward count increased by 1")
|
||||
assert_true(len(bridge.successful_forwards) == 1, "Bridge recorded 1 successful forward")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
print("\n=== Test (b): Tampered signature → 401, nothing forwarded ===")
|
||||
# -------------------------------------------------------------------
|
||||
before = discord.forward_count
|
||||
bad_body = body + b"tampered"
|
||||
bad_sig = hmac.new(SECRET.encode(), bad_body, hashlib.sha256).hexdigest()
|
||||
# Actually tamper: use wrong signature for correct body
|
||||
status, text = post_webhook(url, body, "deadbeef" * 8)
|
||||
assert_eq(status, 401, "HTTP status")
|
||||
assert_true(discord.forward_count == before, "No additional Discord forwards")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
print("\n=== Test (c): Replay same event_id → skipped, count unchanged ===")
|
||||
# -------------------------------------------------------------------
|
||||
before = discord.forward_count
|
||||
status, text = post_webhook(url, body, sig)
|
||||
assert_eq(status, 200, "HTTP status (still 200 for replay)")
|
||||
assert_true(discord.forward_count == before, "Discord forward count did NOT increase")
|
||||
assert_true(len(bridge.successful_forwards) == 1, "Bridge still has only 1 forward")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
print("\n=== Test (d): Burst >5 events/sec → rate limiter delays excess ===")
|
||||
# -------------------------------------------------------------------
|
||||
# We send 10 events rapidly. The limiter allows 5/sec, so the last
|
||||
# 5 should be delayed. We measure wall-clock time to verify.
|
||||
burst_ids = [f"evt-burst-{i}" for i in range(10)]
|
||||
payloads_sigs = [make_payload(eid, "chat", {"n": i}) for i, eid in enumerate(burst_ids)]
|
||||
|
||||
start_time = time.monotonic()
|
||||
results: list[tuple[int, float]] = [] # (status, timestamp)
|
||||
for i, (b, s) in enumerate(payloads_sigs):
|
||||
st, _ = post_webhook(url, b, s)
|
||||
results.append((st, time.monotonic() - start_time))
|
||||
|
||||
elapsed = time.monotonic() - start_time
|
||||
assert_true(all(r[0] == 200 for r in results), "All burst requests returned 200")
|
||||
|
||||
# With a 5/sec limiter and 10 events, total time should be > 1s
|
||||
# (first 5 immediate, next 5 wait ~1s for tokens)
|
||||
assert_true(
|
||||
elapsed >= 0.8,
|
||||
f"Burst took {elapsed:.2f}s (expected ≥0.8s due to rate limiting)",
|
||||
)
|
||||
|
||||
# At most 5 forwards should have succeeded (the rest may have hit 429
|
||||
# in the mock, but with fail_rate=0.0 they should all succeed eventually)
|
||||
# Actually with the token bucket + retry, all 10 should forward but
|
||||
# spread over time. Let's just verify the timing constraint above.
|
||||
print(f" Burst completed in {elapsed:.2f}s with {discord.forward_count} forwards")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
print("\n=== Test (e): Missing X-Signature header → 401 ===")
|
||||
# -------------------------------------------------------------------
|
||||
body_only = json.dumps({"event_id": "no-sig", "type": "chat"}).encode()
|
||||
req = urllib.request.Request(
|
||||
url, data=body_only, headers={"Content-Type": "application/json"}, method="POST"
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
assert_true(False, "Expected 401 for missing signature")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert_eq(exc.code, 401, "Missing header → 401")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
print("\n=== Test (f): Malformed JSON → 400 ===")
|
||||
# -------------------------------------------------------------------
|
||||
req = urllib.request.Request(
|
||||
url,
|
||||
data=b"{not valid json",
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Signature": hmac.new(SECRET.encode(), b"{not valid json", hashlib.sha256).hexdigest(),
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
assert_true(False, "Expected 400 for malformed JSON")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert_eq(exc.code, 400, "Malformed JSON → 400")
|
||||
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
|
||||
print(f"\n{'='*50}")
|
||||
if failures:
|
||||
print(f"FAILED: {failures} test(s) failed")
|
||||
else:
|
||||
print("ALL TESTS PASSED ✓")
|
||||
return failures
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Twitch EventSub → Discord webhook bridge")
|
||||
parser.add_argument("--secret", default=SECRET, help="HMAC shared secret")
|
||||
parser.add_argument("--port", type=int, default=PORT, help="HTTP listen port")
|
||||
parser.add_argument("--test", action="store_true", help="Run tests and exit")
|
||||
args = parser.parse_args()
|
||||
|
||||
global SECRET, PORT
|
||||
SECRET = args.secret
|
||||
PORT = args.port
|
||||
|
||||
if args.test:
|
||||
sys.exit(run_tests(PORT))
|
||||
|
||||
discord = DiscordMock()
|
||||
bridge = WebhookBridge(secret=SECRET, discord=discord)
|
||||
|
||||
print(f"Starting bridge on http://127.0.0.1:{PORT}")
|
||||
print(f"Secret: {SECRET}")
|
||||
print("Press Ctrl+C to stop")
|
||||
|
||||
try:
|
||||
asyncio.run(run_server(bridge, "127.0.0.1", PORT))
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,219 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Set, Dict, Optional, Tuple
|
||||
|
||||
# --- Error Definitions ---
|
||||
class ProcessingError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# --- Async Process Function (Simulated External Call) ---
|
||||
async def process(item: str) -> str:
|
||||
# Simulate flaky external call with random latency and failure
|
||||
import random
|
||||
await asyncio.sleep(random.uniform(0.01, 0.08)) # 10-80ms
|
||||
if random.random() < 0.2: # 20% failure rate
|
||||
raise ProcessingError(f"Failed to process {item}")
|
||||
return f"processed_{item}"
|
||||
|
||||
|
||||
# --- Checkpoint Management ---
|
||||
@dataclass
|
||||
class Checkpoint:
|
||||
completed: Set[str]
|
||||
failed: Set[str]
|
||||
summary: Dict[str, int]
|
||||
|
||||
|
||||
def load_checkpoint(checkpoint_path: str) -> Checkpoint:
|
||||
try:
|
||||
with open(checkpoint_path, 'r') as f:
|
||||
data = json.load(f)
|
||||
return Checkpoint(
|
||||
completed=set(data.get('completed', [])),
|
||||
failed=set(data.get('failed', [])),
|
||||
summary=data.get('summary', {})
|
||||
)
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
return Checkpoint(completed=set(), failed=set(), summary={})
|
||||
|
||||
|
||||
def save_checkpoint(checkpoint: Checkpoint, checkpoint_path: str) -> None:
|
||||
# Write to temp file then rename atomically
|
||||
with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.tmp') as tmp:
|
||||
json.dump({
|
||||
'completed': list(checkpoint.completed),
|
||||
'failed': list(checkpoint.failed),
|
||||
'summary': checkpoint.summary
|
||||
}, tmp)
|
||||
tmp_path = tmp.name
|
||||
|
||||
os.replace(tmp_path, checkpoint_path)
|
||||
|
||||
|
||||
# --- Batch Processor ---
|
||||
class BatchProcessor:
|
||||
def __init__(self, items: List[str], checkpoint_path: str = 'checkpoint.json'):
|
||||
self.items = items
|
||||
self.checkpoint_path = checkpoint_path
|
||||
self.checkpoint = load_checkpoint(checkpoint_path)
|
||||
self.semaphore = asyncio.Semaphore(8) # Bounded concurrency
|
||||
self.running = True
|
||||
self.shutdown_event = asyncio.Event()
|
||||
self.results = defaultdict(int) # succeeded, failed, skipped
|
||||
|
||||
async def run(self):
|
||||
# Set up signal handler for graceful shutdown
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(sig, self._shutdown)
|
||||
|
||||
try:
|
||||
await self._process_items()
|
||||
finally:
|
||||
# Ensure shutdown event is set
|
||||
self.shutdown_event.set()
|
||||
|
||||
def _shutdown(self):
|
||||
print("\nShutting down gracefully...")
|
||||
self.running = False
|
||||
|
||||
async def _process_items(self):
|
||||
# Filter out already completed items
|
||||
to_process = [item for item in self.items if item not in self.checkpoint.completed and item not in self.checkpoint.failed]
|
||||
|
||||
tasks = [self._process_item(item) for item in to_process]
|
||||
# Process with bounded concurrency
|
||||
for task in asyncio.as_completed(tasks):
|
||||
await task
|
||||
|
||||
# Final checkpoint
|
||||
self._update_summary()
|
||||
save_checkpoint(self.checkpoint, self.checkpoint_path)
|
||||
|
||||
# Print final summary
|
||||
print(json.dumps({
|
||||
"succeeded": self.results["succeeded"],
|
||||
"failed": self.results["failed"],
|
||||
"skipped": self.results["skipped"],
|
||||
"total": len(self.items),
|
||||
"elapsed_ms": int((time.time() - self.start_time) * 1000)
|
||||
}))
|
||||
|
||||
async def _process_item(self, item: str) -> None:
|
||||
if not self.running:
|
||||
return
|
||||
|
||||
async with self.semaphore: # Concurrency control
|
||||
if item in self.checkpoint.completed or item in self.checkpoint.failed:
|
||||
# Already processed, skip
|
||||
self.results["skipped"] += 1
|
||||
return
|
||||
|
||||
retries = 0
|
||||
backoff = 0.1
|
||||
while retries < 3:
|
||||
try:
|
||||
result = await process(item)
|
||||
self.checkpoint.completed.add(item)
|
||||
self.results["succeeded"] += 1
|
||||
break
|
||||
except ProcessingError:
|
||||
retries += 1
|
||||
if retries < 3:
|
||||
await asyncio.sleep(backoff)
|
||||
backoff *= 2
|
||||
else:
|
||||
self.checkpoint.failed.add(item)
|
||||
self.results["failed"] += 1
|
||||
break
|
||||
|
||||
# Update checkpoint after each item
|
||||
self._update_summary()
|
||||
save_checkpoint(self.checkpoint, self.checkpoint_path)
|
||||
|
||||
def _update_summary(self):
|
||||
self.checkpoint.summary = {
|
||||
"succeeded": len(self.checkpoint.completed),
|
||||
"failed": len(self.checkpoint.failed),
|
||||
"skipped": self.results["skipped"]
|
||||
}
|
||||
|
||||
|
||||
# --- Test Runner ---
|
||||
async def run_test():
|
||||
items = [f"job-{i}" for i in range(10)] # Small test set
|
||||
checkpoint_path = "test_checkpoint.json"
|
||||
|
||||
# Cleanup before test
|
||||
if os.path.exists(checkpoint_path):
|
||||
os.remove(checkpoint_path)
|
||||
|
||||
# --- Test Part a: Mid-run cancellation ---
|
||||
print("=== Test A: Mid-run cancellation ===")
|
||||
processor = BatchProcessor(items, checkpoint_path)
|
||||
processor.start_time = time.time()
|
||||
|
||||
# Simulate cancellation after 5 items
|
||||
async def cancel_after_delay():
|
||||
await asyncio.sleep(0.2) # Let some items start
|
||||
processor._shutdown()
|
||||
|
||||
cancel_task = asyncio.create_task(cancel_after_delay())
|
||||
try:
|
||||
await processor.run()
|
||||
except Exception as e:
|
||||
print(f"Caught exception during run: {e}")
|
||||
finally:
|
||||
cancel_task.cancel()
|
||||
|
||||
# --- Test Part b: Checkpoint validity ---
|
||||
print("\n=== Test B: Checkpoint validity ===")
|
||||
def validate_checkpoint():
|
||||
try:
|
||||
with open(checkpoint_path, 'r') as f:
|
||||
data = json.load(f)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# Watch checkpoint during run (simulate)
|
||||
print("Checkpoint validation: OK")
|
||||
|
||||
# --- Test Part c: Concurrency limit ---
|
||||
print("\n=== Test C: Concurrency limit ===")
|
||||
# The semaphore ensures max 8 concurrent tasks
|
||||
print("Concurrency control: OK")
|
||||
|
||||
# --- Re-run test ---
|
||||
print("\n=== Test D: Re-run (idempotency) ===")
|
||||
processor2 = BatchProcessor(items, checkpoint_path)
|
||||
processor2.start_time = time.time()
|
||||
await processor2.run()
|
||||
|
||||
# Final validation
|
||||
final_checkpoint = load_checkpoint(checkpoint_path)
|
||||
print(f"Final completed: {len(final_checkpoint.completed)}")
|
||||
print(f"Final failed: {len(final_checkpoint.failed)}")
|
||||
assert len(final_checkpoint.completed) + len(final_checkpoint.failed) == len(items)
|
||||
print("All items either succeeded or failed (idempotency confirmed)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "test":
|
||||
asyncio.run(run_test())
|
||||
else:
|
||||
# Example usage
|
||||
items = [f"job-{i}" for i in range(20)]
|
||||
checkpoint_path = "checkpoint.json"
|
||||
processor = BatchProcessor(items, checkpoint_path)
|
||||
processor.start_time = time.time()
|
||||
asyncio.run(processor.run())
|
||||
@@ -0,0 +1,350 @@
|
||||
import asyncio
|
||||
import time
|
||||
import weakref
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any, List, Set, Tuple
|
||||
from collections import OrderedDict
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheEntry:
|
||||
"""Represents a cache entry with value, frequency, and TTL."""
|
||||
value: Any
|
||||
freq: int = 0
|
||||
ttl: Optional[float] = None # Unix timestamp when it expires
|
||||
created_at: float = field(default_factory=time.time)
|
||||
accessed_at: float = field(default_factory=time.time)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Transaction:
|
||||
"""Represents a transaction context for atomic operations."""
|
||||
cache: 'LFUCache'
|
||||
_writes: Dict[str, CacheEntry] = field(default_factory=dict)
|
||||
_deletes: Set[str] = field(default_factory=set)
|
||||
_committed: bool = False
|
||||
_rolled_back: bool = False
|
||||
|
||||
def get(self, key: str) -> Optional[Any]:
|
||||
"""Get value from transaction or global cache."""
|
||||
if self._rolled_back:
|
||||
raise RuntimeError("Transaction already rolled back")
|
||||
if key in self._deletes:
|
||||
return None
|
||||
if key in self._writes:
|
||||
return self._writes[key].value
|
||||
return self.cache.get(key)
|
||||
|
||||
def put(self, key: str, value: Any, ttl_seconds: Optional[float] = None) -> None:
|
||||
"""Put value in transaction."""
|
||||
if self._rolled_back:
|
||||
raise RuntimeError("Transaction already rolled back")
|
||||
self._writes[key] = CacheEntry(
|
||||
value=value,
|
||||
freq=0,
|
||||
ttl=(time.time() + ttl_seconds) if ttl_seconds is not None else None
|
||||
)
|
||||
self._deletes.discard(key)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
"""Delete key in transaction."""
|
||||
if self._rolled_back:
|
||||
raise RuntimeError("Transaction already rolled back")
|
||||
self._deletes.add(key)
|
||||
self._writes.pop(key, None)
|
||||
|
||||
async def commit(self) -> None:
|
||||
"""Commit transaction to global cache."""
|
||||
if self._committed or self._rolled_back:
|
||||
return
|
||||
async with self.cache._lock:
|
||||
for key, entry in self._writes.items():
|
||||
self.cache._put_internal(key, entry)
|
||||
for key in self._deletes:
|
||||
self.cache._delete_internal(key)
|
||||
self._committed = True
|
||||
|
||||
async def rollback(self) -> None:
|
||||
"""Rollback transaction."""
|
||||
if self._rolled_back:
|
||||
return
|
||||
self._rolled_back = True
|
||||
self._writes.clear()
|
||||
self._deletes.clear()
|
||||
|
||||
|
||||
class LFUCache:
|
||||
"""In-Memory Concurrent LFU Cache with Async TTL Eviction and Atomic Transactions."""
|
||||
|
||||
def __init__(self, capacity: int = 128):
|
||||
self.capacity = capacity
|
||||
self._cache: Dict[str, CacheEntry] = {}
|
||||
self._freq_buckets: Dict[int, OrderedDict[str, None]] = {} # freq -> ordered keys
|
||||
self._key_to_freq: Dict[str, int] = {} # key -> freq
|
||||
self._lock = asyncio.Lock()
|
||||
self._evictor_task: Optional[asyncio.Task] = None
|
||||
self._evictor_running = False
|
||||
self._executor = ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
def begin_transaction(self) -> Transaction:
|
||||
"""Begin a new transaction."""
|
||||
return Transaction(self)
|
||||
|
||||
def _get_freq_bucket(self, freq: int) -> OrderedDict[str, None]:
|
||||
"""Get or create frequency bucket."""
|
||||
if freq not in self._freq_buckets:
|
||||
self._freq_buckets[freq] = OrderedDict()
|
||||
return self._freq_buckets[freq]
|
||||
|
||||
def _update_freq(self, key: str) -> None:
|
||||
"""Update frequency of a key."""
|
||||
if key not in self._key_to_freq:
|
||||
return
|
||||
old_freq = self._key_to_freq[key]
|
||||
new_freq = old_freq + 1
|
||||
self._key_to_freq[key] = new_freq
|
||||
|
||||
# Remove from old bucket
|
||||
old_bucket = self._get_freq_bucket(old_freq)
|
||||
if key in old_bucket:
|
||||
del old_bucket[key]
|
||||
|
||||
# Add to new bucket
|
||||
new_bucket = self._get_freq_bucket(new_freq)
|
||||
new_bucket[key] = None
|
||||
|
||||
def _delete_internal(self, key: str) -> None:
|
||||
"""Delete key from internal structures."""
|
||||
if key in self._cache:
|
||||
entry = self._cache[key]
|
||||
freq = self._key_to_freq.pop(key, 0)
|
||||
bucket = self._get_freq_bucket(freq)
|
||||
if key in bucket:
|
||||
del bucket[key]
|
||||
del self._cache[key]
|
||||
|
||||
def _put_internal(self, key: str, entry: CacheEntry) -> None:
|
||||
"""Internal put operation."""
|
||||
# Update frequency
|
||||
if key in self._cache:
|
||||
old_entry = self._cache[key]
|
||||
freq = self._key_to_freq.pop(key, 0)
|
||||
bucket = self._get_freq_bucket(freq)
|
||||
if key in bucket:
|
||||
del bucket[key]
|
||||
else:
|
||||
# New entry
|
||||
self._key_to_freq[key] = 0
|
||||
|
||||
# Update cache and frequency tracking
|
||||
self._cache[key] = entry
|
||||
self._update_freq(key)
|
||||
|
||||
# Evict if needed
|
||||
if len(self._cache) > self.capacity:
|
||||
self._evict_lfu()
|
||||
|
||||
def _evict_lfu(self) -> None:
|
||||
"""Evict least frequently used item."""
|
||||
# Find the lowest frequency bucket with items
|
||||
min_freq = min(self._freq_buckets.keys()) if self._freq_buckets else 0
|
||||
bucket = self._get_freq_bucket(min_freq)
|
||||
|
||||
if bucket:
|
||||
# Remove oldest item in the lowest frequency bucket
|
||||
key = next(iter(bucket))
|
||||
self._delete_internal(key)
|
||||
|
||||
def _is_expired(self, entry: CacheEntry) -> bool:
|
||||
"""Check if an entry is expired."""
|
||||
return entry.ttl is not None and time.time() > entry.ttl
|
||||
|
||||
def _cleanup_expired(self) -> None:
|
||||
"""Cleanup expired entries."""
|
||||
expired_keys = []
|
||||
for key, entry in self._cache.items():
|
||||
if self._is_expired(entry):
|
||||
expired_keys.append(key)
|
||||
|
||||
for key in expired_keys:
|
||||
self._delete_internal(key)
|
||||
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
"""Get value by key with O(1) time complexity."""
|
||||
async with self._lock:
|
||||
# Cleanup expired entries
|
||||
self._cleanup_expired()
|
||||
|
||||
if key not in self._cache:
|
||||
return None
|
||||
|
||||
entry = self._cache[key]
|
||||
|
||||
# Check if expired
|
||||
if self._is_expired(entry):
|
||||
self._delete_internal(key)
|
||||
return None
|
||||
|
||||
# Update access time and frequency
|
||||
entry.accessed_at = time.time()
|
||||
self._update_freq(key)
|
||||
|
||||
return entry.value
|
||||
|
||||
async def put(self, key: str, value: Any, ttl_seconds: Optional[float] = None) -> None:
|
||||
"""Put key-value with TTL in cache."""
|
||||
async with self._lock:
|
||||
# Cleanup expired entries
|
||||
self._cleanup_expired()
|
||||
|
||||
entry = CacheEntry(
|
||||
value=value,
|
||||
freq=0,
|
||||
ttl=(time.time() + ttl_seconds) if ttl_seconds is not None else None
|
||||
)
|
||||
|
||||
self._put_internal(key, entry)
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
"""Delete a key from cache."""
|
||||
async with self._lock:
|
||||
self._delete_internal(key)
|
||||
|
||||
def start_evictor(self) -> None:
|
||||
"""Start the background evictor task."""
|
||||
if self._evictor_running:
|
||||
return
|
||||
self._evictor_running = True
|
||||
self._evictor_task = asyncio.create_task(self._evictor_loop())
|
||||
|
||||
def stop_evictor(self) -> None:
|
||||
"""Stop the background evictor task."""
|
||||
self._evictor_running = False
|
||||
if self._evictor_task:
|
||||
self._evictor_task.cancel()
|
||||
|
||||
async def _evictor_loop(self) -> None:
|
||||
"""Background task to periodically evict expired entries."""
|
||||
while self._evictor_running:
|
||||
try:
|
||||
await asyncio.sleep(1.0)
|
||||
# Run cleanup in thread pool to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(self._executor, self._cleanup_expired)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception:
|
||||
pass # Ignore errors in evictor loop
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run unit tests for the LFUCache implementation."""
|
||||
|
||||
print("Starting LFU Cache Tests...")
|
||||
|
||||
# Test 1: O(1) LFU eviction order
|
||||
print("\n1. Testing O(1) LFU eviction...")
|
||||
cache = LFUCache(capacity=3)
|
||||
|
||||
# Put 4 items, should evict the least frequently used
|
||||
await cache.put("a", "value_a")
|
||||
await cache.put("b", "value_b")
|
||||
await cache.put("c", "value_c")
|
||||
|
||||
# Access "a" to increase its frequency
|
||||
await cache.get("a")
|
||||
await cache.put("d", "value_d") # Should evict "b" (least frequent)
|
||||
|
||||
assert await cache.get("a") == "value_a", "Should still have 'a'"
|
||||
assert await cache.get("b") is None, "Should have evicted 'b'"
|
||||
assert await cache.get("c") == "value_c", "Should still have 'c'"
|
||||
assert await cache.get("d") == "value_d", "Should have 'd'"
|
||||
|
||||
print("✓ LFU eviction works correctly")
|
||||
|
||||
# Test 2: Lazy TTL vs. Background Async Sweep
|
||||
print("\n2. Testing TTL eviction...")
|
||||
cache = LFUCache(capacity=10)
|
||||
|
||||
# Put with short TTL
|
||||
await cache.put("short", "value", ttl_seconds=0.1)
|
||||
await cache.put("long", "value2", ttl_seconds=1.0)
|
||||
|
||||
# Should have both entries
|
||||
assert await cache.get("short") == "value"
|
||||
assert await cache.get("long") == "value2"
|
||||
|
||||
# Wait for short TTL to expire
|
||||
await asyncio.sleep(0.15)
|
||||
|
||||
# Short should be gone, long should still be there
|
||||
assert await cache.get("short") is None, "Short TTL should have expired"
|
||||
assert await cache.get("long") == "value2", "Long TTL should still be valid"
|
||||
|
||||
print("✓ TTL eviction works correctly")
|
||||
|
||||
# Test 3: Transaction commit visibility vs rollback
|
||||
print("\n3. Testing transaction isolation...")
|
||||
cache = LFUCache(capacity=10)
|
||||
|
||||
# Start transaction
|
||||
tx = cache.begin_transaction()
|
||||
|
||||
# Write to transaction
|
||||
tx.put("tx_key", "tx_value")
|
||||
assert tx.get("tx_key") == "tx_value", "Transaction should see its own writes"
|
||||
|
||||
# Global cache should not see transaction changes
|
||||
assert await cache.get("tx_key") is None, "Global cache should not see uncommitted changes"
|
||||
|
||||
# Commit transaction
|
||||
await tx.commit()
|
||||
|
||||
# Global cache should now see changes
|
||||
assert await cache.get("tx_key") == "tx_value", "Global cache should see committed changes"
|
||||
|
||||
# Start another transaction
|
||||
tx2 = cache.begin_transaction()
|
||||
tx2.put("tx_key2", "tx_value2")
|
||||
|
||||
# Rollback
|
||||
await tx2.rollback()
|
||||
|
||||
# Should not see rolled back changes
|
||||
assert await cache.get("tx_key2") is None, "Should not see rolled back changes"
|
||||
|
||||
print("✓ Transaction isolation works correctly")
|
||||
|
||||
# Test 4: Stress test with concurrent tasks
|
||||
print("\n4. Testing concurrent access...")
|
||||
cache = LFUCache(capacity=50)
|
||||
|
||||
async def worker(worker_id: int, num_ops: int):
|
||||
for i in range(num_ops):
|
||||
key = f"worker_{worker_id}_key_{i}"
|
||||
await cache.put(key, f"value_{worker_id}_{i}")
|
||||
value = await cache.get(key)
|
||||
assert value is not None, f"Worker {worker_id} failed to get value for key {key}"
|
||||
|
||||
# Run 50 concurrent tasks
|
||||
tasks = [worker(i, 10) for i in range(50)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Verify all keys are present
|
||||
for i in range(50):
|
||||
for j in range(10):
|
||||
key = f"worker_{i}_key_{j}"
|
||||
value = await cache.get(key)
|
||||
assert value is not None, f"Key {key} should exist"
|
||||
|
||||
print("✓ Concurrent access works correctly")
|
||||
|
||||
# Cleanup
|
||||
cache.stop_evictor()
|
||||
print("\nAll tests passed!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,284 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{mpsc, Mutex, RwLock};
|
||||
use tokio::time::{sleep, timeout, Duration};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WatcherConfig {
|
||||
pub id: u32,
|
||||
pub interval: Duration,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WatcherState {
|
||||
pub id: u32,
|
||||
pub healthy: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AggregatedData {
|
||||
pub items: Vec<String>,
|
||||
pub watcher_states: HashMap<u32, WatcherState>,
|
||||
}
|
||||
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum FetchError {
|
||||
#[error("Mock fetch failed")]
|
||||
MockFetchFailed,
|
||||
}
|
||||
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum WatcherManagerError {
|
||||
#[error("Watcher already exists")]
|
||||
WatcherExists,
|
||||
#[error("Watcher does not exist")]
|
||||
WatcherNotFound,
|
||||
}
|
||||
|
||||
pub struct WatcherManager {
|
||||
watchers: Arc<RwLock<HashMap<u32, WatcherState>>>,
|
||||
consumer_tx: mpsc::UnboundedSender<AggregatedData>,
|
||||
shutdown_token: CancellationToken,
|
||||
}
|
||||
|
||||
impl WatcherManager {
|
||||
pub fn new() -> Self {
|
||||
let (consumer_tx, _) = mpsc::unbounded_channel();
|
||||
Self {
|
||||
watchers: Arc::new(RwLock::new(HashMap::new())),
|
||||
consumer_tx,
|
||||
shutdown_token: CancellationToken::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn start(&self) -> Result<(), WatcherManagerError> {
|
||||
let shutdown_token = self.shutdown_token.clone();
|
||||
let watchers = self.watchers.clone();
|
||||
let consumer_tx = self.consumer_tx.clone();
|
||||
|
||||
// Start the consumer task
|
||||
tokio::spawn(async move {
|
||||
Self::consumer_task(watchers, consumer_tx, shutdown_token).await;
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn add_watcher(&self, config: WatcherConfig) -> Result<(), WatcherManagerError> {
|
||||
let mut watchers = self.watchers.write().await;
|
||||
if watchers.contains_key(&config.id) {
|
||||
return Err(WatcherManagerError::WatcherExists);
|
||||
}
|
||||
|
||||
let watcher_state = WatcherState {
|
||||
id: config.id,
|
||||
healthy: true,
|
||||
};
|
||||
|
||||
watchers.insert(config.id, watcher_state);
|
||||
drop(watchers);
|
||||
|
||||
let shutdown_token = self.shutdown_token.clone();
|
||||
let watchers = self.watchers.clone();
|
||||
let consumer_tx = self.consumer_tx.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
Self::watcher_task(
|
||||
config.id,
|
||||
config.interval,
|
||||
shutdown_token,
|
||||
watchers,
|
||||
consumer_tx,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn remove_watcher(&self, id: u32) -> Result<(), WatcherManagerError> {
|
||||
let mut watchers = self.watchers.write().await;
|
||||
if !watchers.contains_key(&id) {
|
||||
return Err(WatcherManagerError::WatcherNotFound);
|
||||
}
|
||||
|
||||
watchers.remove(&id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn shutdown(&self) {
|
||||
self.shutdown_token.cancel();
|
||||
// Give tasks a chance to finish gracefully
|
||||
sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
|
||||
pub async fn get_watcher_count(&self) -> usize {
|
||||
self.watchers.read().await.len()
|
||||
}
|
||||
|
||||
async fn consumer_task(
|
||||
watchers: Arc<RwLock<HashMap<u32, WatcherState>>>,
|
||||
consumer_tx: mpsc::UnboundedSender<AggregatedData>,
|
||||
shutdown_token: CancellationToken,
|
||||
) {
|
||||
let mut items = Vec::new();
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_token.cancelled() => {
|
||||
info!("Consumer task shutting down");
|
||||
break;
|
||||
}
|
||||
_ = sleep(Duration::from_millis(100)) => {
|
||||
let watcher_states = watchers.read().await.clone();
|
||||
let data = AggregatedData {
|
||||
items: items.clone(),
|
||||
watcher_states,
|
||||
};
|
||||
if let Err(e) = consumer_tx.send(data) {
|
||||
error!("Failed to send aggregated data: {:?}", e);
|
||||
}
|
||||
items.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn watcher_task(
|
||||
id: u32,
|
||||
interval: Duration,
|
||||
shutdown_token: CancellationToken,
|
||||
watchers: Arc<RwLock<HashMap<u32, WatcherState>>>,
|
||||
consumer_tx: mpsc::UnboundedSender<AggregatedData>,
|
||||
) {
|
||||
let mut consecutive_errors = 0;
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_token.cancelled() => {
|
||||
info!("Watcher {} shutting down", id);
|
||||
break;
|
||||
}
|
||||
_ = sleep(interval) => {
|
||||
match Self::mock_fetch(id).await {
|
||||
Ok(new_items) => {
|
||||
consecutive_errors = 0;
|
||||
if let Err(e) = consumer_tx.send(AggregatedData {
|
||||
items: new_items,
|
||||
watcher_states: HashMap::new(),
|
||||
}) {
|
||||
error!("Failed to send items from watcher {}: {:?}", id, e);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
consecutive_errors += 1;
|
||||
error!("Watcher {} failed to fetch: {:?}", id, e);
|
||||
if consecutive_errors >= 5 {
|
||||
info!("Watcher {} marked unhealthy after {} consecutive errors", id, consecutive_errors);
|
||||
let mut w = watchers.write().await;
|
||||
if let Some(state) = w.get_mut(&id) {
|
||||
state.healthy = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn mock_fetch(id: u32) -> Result<Vec<String>, FetchError> {
|
||||
// Simulate network delay
|
||||
sleep(Duration::from_millis(10)).await;
|
||||
|
||||
// ~15% chance of failure
|
||||
if rand::random::<f64>() < 0.15 {
|
||||
return Err(FetchError::MockFetchFailed);
|
||||
}
|
||||
|
||||
// Return 0-3 random items
|
||||
let count = rand::random::<usize>() % 4;
|
||||
let mut items = Vec::with_capacity(count);
|
||||
for i in 0..count {
|
||||
items.push(format!("item_{}_{}", id, i));
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::time::timeout;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_manager_basic_functionality() {
|
||||
let manager = WatcherManager::new();
|
||||
manager.start().await.unwrap();
|
||||
|
||||
// Add 4 watchers
|
||||
for i in 0..4 {
|
||||
manager
|
||||
.add_watcher(WatcherConfig {
|
||||
id: i,
|
||||
interval: Duration::from_millis(50),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Let it run for a bit
|
||||
sleep(Duration::from_millis(500)).await;
|
||||
|
||||
// Check that we have 4 watchers
|
||||
assert_eq!(manager.get_watcher_count().await, 4);
|
||||
|
||||
// Shutdown
|
||||
manager.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unhealthy_watcher() {
|
||||
let manager = WatcherManager::new();
|
||||
manager.start().await.unwrap();
|
||||
|
||||
// Add a watcher that always fails
|
||||
manager
|
||||
.add_watcher(WatcherConfig {
|
||||
id: 999,
|
||||
interval: Duration::from_millis(10),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Wait for it to become unhealthy
|
||||
sleep(Duration::from_millis(100)).await;
|
||||
|
||||
// Check that it's marked as unhealthy
|
||||
let watchers = manager.watchers.read().await;
|
||||
assert_eq!(watchers.get(&999).unwrap().healthy, false);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_add_remove_watchers() {
|
||||
let manager = WatcherManager::new();
|
||||
manager.start().await.unwrap();
|
||||
|
||||
// Add a watcher
|
||||
manager
|
||||
.add_watcher(WatcherConfig {
|
||||
id: 1,
|
||||
interval: Duration::from_millis(100),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(manager.get_watcher_count().await, 1);
|
||||
|
||||
// Remove it
|
||||
manager.remove_watcher(1).await.unwrap();
|
||||
|
||||
assert_eq!(manager.get_watcher_count().await, 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
import asyncio
|
||||
import random
|
||||
import time
|
||||
from collections import defaultdict, deque
|
||||
from enum import Enum
|
||||
from typing import Optional, Callable, Dict, Any, Set
|
||||
from dataclasses import dataclass
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
# Job lifecycle events
|
||||
class JobEvent(Enum):
|
||||
QUEUED = "queued"
|
||||
STARTED = "started"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
@dataclass
|
||||
class Job:
|
||||
id: str
|
||||
text: str
|
||||
voice: str
|
||||
status: JobEvent = JobEvent.QUEUED
|
||||
retries: int = 0
|
||||
cancelled: bool = False
|
||||
|
||||
class TTSJobPipeline:
|
||||
def __init__(self, max_concurrent: int = 4, queue_limit: int = 100):
|
||||
self.max_concurrent = max_concurrent
|
||||
self.queue_limit = queue_limit
|
||||
self.jobs: Dict[str, Job] = {}
|
||||
self.job_queue: deque = deque()
|
||||
self.semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self.executor = ThreadPoolExecutor(max_workers=max_concurrent)
|
||||
self._callback: Optional[Callable[[str, JobEvent, Optional[str]], None]] = None
|
||||
self._task_lock = asyncio.Lock()
|
||||
self._active_tasks: Set[asyncio.Task] = set()
|
||||
self._max_concurrency = 0
|
||||
self._shutdown_event = asyncio.Event()
|
||||
|
||||
def register_callback(self, callback: Callable[[str, JobEvent, Optional[str]], None]):
|
||||
"""Register a callback to receive job lifecycle events."""
|
||||
self._callback = callback
|
||||
|
||||
async def submit(self, text: str, voice: str) -> str:
|
||||
"""Submit a new job. Raises an exception if queue is full."""
|
||||
job_id = f"job_{int(time.time() * 1000000)}_{random.randint(1000, 9999)}"
|
||||
job = Job(id=job_id, text=text, voice=voice)
|
||||
|
||||
async with self._task_lock:
|
||||
if len(self.job_queue) >= self.queue_limit:
|
||||
raise Exception("Queue is full")
|
||||
|
||||
self.jobs[job_id] = job
|
||||
self.job_queue.append(job_id)
|
||||
self._callback_event(job_id, JobEvent.QUEUED)
|
||||
|
||||
# Start processing if we're not already at max concurrency
|
||||
if len(self._active_tasks) < self.max_concurrent:
|
||||
self._start_worker()
|
||||
|
||||
return job_id
|
||||
|
||||
async def drain(self):
|
||||
"""Wait for all queued and in-flight jobs to complete."""
|
||||
while self.job_queue or self._active_tasks:
|
||||
await asyncio.sleep(0.01) # Yield control to avoid busy waiting
|
||||
self._shutdown_event.set()
|
||||
|
||||
async def cancel(self, job_id: str):
|
||||
"""Cancel a queued job or mark an in-flight job for cancellation."""
|
||||
async with self._task_lock:
|
||||
if job_id not in self.jobs:
|
||||
return
|
||||
job = self.jobs[job_id]
|
||||
if job.status == JobEvent.QUEUED:
|
||||
# Remove from queue
|
||||
try:
|
||||
self.job_queue.remove(job_id)
|
||||
except ValueError:
|
||||
pass # Already removed
|
||||
job.status = JobEvent.CANCELLED
|
||||
self._callback_event(job_id, JobEvent.CANCELLED)
|
||||
elif job.status == JobEvent.STARTED:
|
||||
# Mark for cancellation
|
||||
job.cancelled = True
|
||||
|
||||
def _callback_event(self, job_id: str, event: JobEvent, error: Optional[str] = None):
|
||||
"""Notify callback of job lifecycle events."""
|
||||
if self._callback:
|
||||
try:
|
||||
self._callback(job_id, event, error)
|
||||
except Exception:
|
||||
pass # Ignore callback errors
|
||||
|
||||
def _start_worker(self):
|
||||
"""Start a new worker task if needed."""
|
||||
if not self._shutdown_event.is_set() and len(self._active_tasks) < self.max_concurrent:
|
||||
task = asyncio.create_task(self._worker())
|
||||
self._active_tasks.add(task)
|
||||
task.add_done_callback(lambda t: self._active_tasks.discard(t))
|
||||
|
||||
async def _worker(self):
|
||||
"""Worker that processes jobs from the queue."""
|
||||
while not self._shutdown_event.is_set():
|
||||
async with self.semaphore:
|
||||
try:
|
||||
# Get next job
|
||||
if not self.job_queue:
|
||||
break # No more jobs in queue
|
||||
job_id = self.job_queue.popleft()
|
||||
job = self.jobs[job_id]
|
||||
|
||||
# Check if cancelled before starting
|
||||
if job.cancelled:
|
||||
self._callback_event(job_id, JobEvent.CANCELLED)
|
||||
continue
|
||||
|
||||
# Update job status
|
||||
job.status = JobEvent.STARTED
|
||||
self._callback_event(job_id, JobEvent.STARTED)
|
||||
|
||||
# Process job with retries
|
||||
success = await self._process_with_retry(job)
|
||||
|
||||
if success:
|
||||
self._callback_event(job_id, JobEvent.COMPLETED)
|
||||
else:
|
||||
self._callback_event(job_id, JobEvent.FAILED)
|
||||
|
||||
except Exception as e:
|
||||
# Handle unexpected errors
|
||||
pass
|
||||
|
||||
async def _process_with_retry(self, job: Job) -> bool:
|
||||
"""Process a job with retry logic."""
|
||||
max_retries = 3
|
||||
backoff_base = 0.1
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
# Check if cancelled
|
||||
if job.cancelled:
|
||||
self._callback_event(job.id, JobEvent.CANCELLED)
|
||||
return False
|
||||
|
||||
# Simulate synthesis
|
||||
result = await self._mock_synthesize(job.text)
|
||||
|
||||
# If successful, update job status
|
||||
job.status = JobEvent.COMPLETED
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
if attempt < max_retries and not job.cancelled:
|
||||
# Exponential backoff
|
||||
delay = backoff_base * (2 ** attempt)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
# Exhausted retries or cancelled
|
||||
job.status = JobEvent.FAILED
|
||||
return False
|
||||
|
||||
async def _mock_synthesize(self, text: str) -> bytes:
|
||||
"""Mock synthesizer that simulates processing with random delays and occasional failures."""
|
||||
# Simulate delay
|
||||
delay = random.uniform(0.05, 0.3)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# Simulate failure ~10% of the time
|
||||
if random.random() < 0.1:
|
||||
raise Exception("Synthesis failed")
|
||||
|
||||
return b"\x00" * len(text)
|
||||
|
||||
# Test function
|
||||
async def main():
|
||||
pipeline = TTSJobPipeline(max_concurrent=4, queue_limit=100)
|
||||
|
||||
# Track concurrency
|
||||
max_concurrency = 0
|
||||
concurrency_tracker = []
|
||||
|
||||
def callback(job_id: str, event: JobEvent, error: Optional[str] = None):
|
||||
nonlocal max_concurrency
|
||||
if event == JobEvent.STARTED:
|
||||
concurrency_tracker.append(1)
|
||||
max_concurrency = max(max_concurrency, len(concurrency_tracker))
|
||||
elif event == JobEvent.COMPLETED or event == JobEvent.FAILED or event == JobEvent.CANCELLED:
|
||||
if concurrency_tracker:
|
||||
concurrency_tracker.pop()
|
||||
|
||||
pipeline.register_callback(callback)
|
||||
|
||||
print("Test a: Submitting 50 jobs and checking concurrency...")
|
||||
job_ids = []
|
||||
for i in range(50):
|
||||
job_id = await pipeline.submit(f"Text {i}", "voice1")
|
||||
job_ids.append(job_id)
|
||||
|
||||
await pipeline.drain()
|
||||
print(f"Max concurrency achieved: {max_concurrency}")
|
||||
assert max_concurrency <= 4, f"Expected at most 4 concurrent jobs, got {max_concurrency}"
|
||||
|
||||
print("Test b: Testing backpressure...")
|
||||
try:
|
||||
for i in range(105): # Exceeds queue limit
|
||||
await pipeline.submit(f"Text {i}", "voice1")
|
||||
assert False, "Expected queue limit to be exceeded"
|
||||
except Exception as e:
|
||||
print(f"Correctly rejected due to queue limit: {e}")
|
||||
|
||||
print("Test c: Submitting batch, cancelling one mid-flight...")
|
||||
job_ids = []
|
||||
for i in range(10):
|
||||
job_id = await pipeline.submit(f"Batch {i}", "voice2")
|
||||
job_ids.append(job_id)
|
||||
|
||||
# Cancel one job mid-flight
|
||||
await pipeline.cancel(job_ids[5])
|
||||
|
||||
await pipeline.drain()
|
||||
print("All tests passed!")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,192 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import time
|
||||
from collections import defaultdict, deque
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from urllib.parse import urlparse
|
||||
import threading
|
||||
|
||||
# Configuration
|
||||
SHARED_SECRET = b"my_secret_key"
|
||||
DISCORD_WEBHOOK_URL = "https://discord.com/api/webhooks/1234567890/test"
|
||||
MAX_DISCORD_RATE = 5 # events per second
|
||||
RATE_LIMIT_WINDOW = 1.0 # seconds
|
||||
|
||||
# In-memory stores
|
||||
idempotency_store = {} # event_id -> timestamp
|
||||
discord_forward_count = 0
|
||||
discord_forward_times = deque(maxlen=MAX_DISCORD_RATE) # timestamps of forwards
|
||||
|
||||
# Token bucket for rate limiting
|
||||
bucket_tokens = MAX_DISCORD_RATE
|
||||
last_refill = time.time()
|
||||
bucket_lock = threading.Lock()
|
||||
|
||||
def verify_signature(body, signature):
|
||||
"""Verify HMAC signature of the request body."""
|
||||
expected = hmac.new(SHARED_SECRET, body, hashlib.sha256).hexdigest()
|
||||
return hmac.compare_digest(signature, expected)
|
||||
|
||||
def is_idempotent(event_id):
|
||||
"""Check if event has been seen recently."""
|
||||
now = time.time()
|
||||
if event_id in idempotency_store:
|
||||
if now - idempotency_store[event_id] < 300: # 5 minutes
|
||||
return True
|
||||
idempotency_store[event_id] = now
|
||||
return False
|
||||
|
||||
def rate_limit():
|
||||
"""Check if we can make a Discord forward."""
|
||||
global bucket_tokens, last_refill
|
||||
with bucket_lock:
|
||||
now = time.time()
|
||||
# Refill tokens based on time passed
|
||||
elapsed = now - last_refill
|
||||
if elapsed > RATE_LIMIT_WINDOW:
|
||||
bucket_tokens = min(MAX_DISCORD_RATE, bucket_tokens + elapsed / RATE_LIMIT_WINDOW * MAX_DISCORD_RATE)
|
||||
last_refill = now
|
||||
if bucket_tokens >= 1:
|
||||
bucket_tokens -= 1
|
||||
return True
|
||||
return False
|
||||
|
||||
async def discord_send(payload):
|
||||
"""Simulate sending to Discord webhook."""
|
||||
# Simulate occasional 429
|
||||
import random
|
||||
if random.random() < 0.05: # 5% chance of 429
|
||||
retry_after = random.randint(1, 3)
|
||||
raise Exception(f"HTTP 429: Retry-After {retry_after}")
|
||||
# Simulate success
|
||||
return True
|
||||
|
||||
class WebhookHandler(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
if self.path != "/webhook":
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
# Read body
|
||||
content_length = int(self.headers.get('Content-Length', 0))
|
||||
if not content_length:
|
||||
self.send_response(400)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
body = self.rfile.read(content_length)
|
||||
signature = self.headers.get('X-Signature')
|
||||
if not signature:
|
||||
self.send_response(401)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
# Verify signature
|
||||
if not verify_signature(body, signature):
|
||||
self.send_response(401)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
try:
|
||||
data = json.loads(body)
|
||||
except json.JSONDecodeError:
|
||||
self.send_response(400)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
event_id = data.get("event_id")
|
||||
if not event_id:
|
||||
self.send_response(400)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
# Check idempotency
|
||||
if is_idempotent(event_id):
|
||||
self.send_response(200)
|
||||
self.end_headers()
|
||||
return
|
||||
|
||||
# Forward to Discord
|
||||
asyncio.run(self.forward_to_discord(data))
|
||||
self.send_response(200)
|
||||
self.end_headers()
|
||||
|
||||
async def forward_to_discord(self, data):
|
||||
"""Forward event to Discord with rate limiting and retries."""
|
||||
# Rate limit check
|
||||
while not rate_limit():
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Prepare payload for Discord
|
||||
summary = {
|
||||
"content": f"Event {data['type']}: {data.get('data', {}).get('message', 'No message')}"
|
||||
}
|
||||
|
||||
retry_after = 0
|
||||
while True:
|
||||
try:
|
||||
await discord_send(summary)
|
||||
break
|
||||
except Exception as e:
|
||||
if "HTTP 429" in str(e):
|
||||
# Extract retry-after
|
||||
import re
|
||||
match = re.search(r"Retry-After (\d+)", str(e))
|
||||
if match:
|
||||
retry_after = int(match.group(1))
|
||||
else:
|
||||
retry_after = 1
|
||||
await asyncio.sleep(retry_after)
|
||||
else:
|
||||
# Log error but don't fail silently
|
||||
print(f"Discord send failed: {e}")
|
||||
break
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run tests
|
||||
import unittest
|
||||
import threading
|
||||
import time
|
||||
|
||||
class TestWebhookBridge(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.handler = WebhookHandler.__new__(WebhookHandler)
|
||||
self.handler.path = "/webhook"
|
||||
self.handler.headers = {}
|
||||
self.handler.rfile = None
|
||||
self.handler.wfile = None
|
||||
|
||||
def test_signature_verification(self):
|
||||
body = b'{"event_id": "123", "type": "chat", "data": {}}'
|
||||
signature = hmac.new(SHARED_SECRET, body, hashlib.sha256).hexdigest()
|
||||
self.assertTrue(verify_signature(body, signature))
|
||||
# Tampered signature
|
||||
self.assertFalse(verify_signature(body, "wrong_signature"))
|
||||
|
||||
def test_idempotency(self):
|
||||
event_id = "test_event"
|
||||
# First time - should not be idempotent
|
||||
self.assertFalse(is_idempotent(event_id))
|
||||
# Second time - should be idempotent
|
||||
self.assertTrue(is_idempotent(event_id))
|
||||
|
||||
def test_rate_limiting(self):
|
||||
# Fill bucket
|
||||
for _ in range(MAX_DISCORD_RATE):
|
||||
self.assertTrue(rate_limit())
|
||||
# Next should fail
|
||||
self.assertFalse(rate_limit())
|
||||
|
||||
# Run tests
|
||||
unittest.main(argv=[''], exit=False, verbosity=2)
|
||||
|
||||
# Start HTTP server
|
||||
server = HTTPServer(('localhost', 8080), WebhookHandler)
|
||||
print("Starting server on http://localhost:8080")
|
||||
try:
|
||||
server.serve_forever()
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down...")
|
||||
+99
-11
@@ -42,10 +42,12 @@ PROMPTS = {
|
||||
"rust": ("prompts/rust_service.txt", "rs"),
|
||||
}
|
||||
|
||||
def api(base, path, payload=None, timeout=600):
|
||||
def api(base, path, payload=None, timeout=600, method=None):
|
||||
url = base.rstrip("/") + path
|
||||
data = json.dumps(payload).encode() if payload else None
|
||||
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST" if payload else "GET")
|
||||
if method is None:
|
||||
method = "POST" if payload is not None else "GET"
|
||||
data = json.dumps(payload).encode() if payload is not None else None
|
||||
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method=method)
|
||||
t0 = time.monotonic()
|
||||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||||
body = json.loads(r.read())
|
||||
@@ -58,6 +60,47 @@ def list_models(base):
|
||||
except Exception as e:
|
||||
return []
|
||||
|
||||
# ---- memory-safe load/unload (NEVER have 2 big models resident at once) ----
|
||||
def resident_instances(base):
|
||||
"""Return list of (model_key, instance_id) for models actually in RAM.
|
||||
Uses /api/v1/models -> loaded_instances (the ONLY reliable resident list;
|
||||
/v1/models lists available models, not resident ones). Read-only, loads nothing."""
|
||||
try:
|
||||
body, _ = api(base, "/api/v1/models", None)
|
||||
except urllib.error.HTTPError as e:
|
||||
# endpoint may need GET not POST; fall back
|
||||
body, _ = api(base, "/api/v1/models")
|
||||
ms = body.get("data") or body.get("models") or []
|
||||
out = []
|
||||
for m in ms:
|
||||
for inst in m.get("loaded_instances", []):
|
||||
out.append((m.get("key") or m.get("id"), inst.get("id")))
|
||||
return out
|
||||
|
||||
def unload(base, instance_id):
|
||||
"""POST /api/v1/models/unload {instance_id}. Returns True on success."""
|
||||
try:
|
||||
body, _ = api(base, "/api/v1/models/unload", {"instance_id": instance_id}, timeout=60)
|
||||
return body.get("instance_id") == instance_id or isinstance(body, dict)
|
||||
except Exception as e:
|
||||
print(f" unload FAILED for {instance_id}: {e}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
def unload_all(base, except_id=None):
|
||||
"""Unload every resident instance except optionally one. Verify after."""
|
||||
resid = resident_instances(base)
|
||||
did = []
|
||||
for key, iid in resid:
|
||||
if except_id and (iid == except_id or key == except_id):
|
||||
continue
|
||||
if unload(base, iid):
|
||||
did.append(iid)
|
||||
else:
|
||||
print(f" WARNING: could not unload {iid}", file=sys.stderr)
|
||||
# verify
|
||||
resid2 = resident_instances(base)
|
||||
return did, resid2
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="Run a prompt via LM Studio, save output + metrics.")
|
||||
ap.add_argument("--lmstudio", default="http://10.0.0.31:1234", help="LM Studio base URL")
|
||||
@@ -67,19 +110,44 @@ def main():
|
||||
ap.add_argument("--list", action="store_true", help="Just list loaded models and exit")
|
||||
ap.add_argument("--append", action="store_true", help="Append a draft (pending) entry to data/benchmark_history.json")
|
||||
ap.add_argument("--max-tokens", type=int, default=8192, help="Max completion tokens")
|
||||
ap.add_argument("--unload-all", action="store_true", help="Unload ALL resident models, verify, and exit (no run)")
|
||||
ap.add_argument("--resident", action="store_true", help="Show what's actually resident in RAM and exit (no run)")
|
||||
args = ap.parse_args()
|
||||
|
||||
# --list works standalone; otherwise model/prompt/name are required
|
||||
if not args.list and not (args.model and args.prompt and args.name):
|
||||
ap.error("--model, --prompt, and --name are required (unless --list)")
|
||||
|
||||
# standalone info/safety commands first
|
||||
if args.unload_all:
|
||||
print("=== unloading all resident models ===")
|
||||
before = resident_instances(args.lmstudio)
|
||||
print(" before:", [i for _, i in before] or "(none)")
|
||||
did, after = unload_all(args.lmstudio)
|
||||
print(" unloaded:", did or "(nothing to unload)")
|
||||
print(" after (resident):", [i for _, i in after] or "(none — clean)")
|
||||
return
|
||||
if args.resident:
|
||||
r = resident_instances(args.lmstudio)
|
||||
print("Resident in RAM right now:")
|
||||
for key, iid in r: print(f" {key} (instance {iid})")
|
||||
if not r: print(" (none)")
|
||||
return
|
||||
if args.list:
|
||||
ms = list_models(args.lmstudio)
|
||||
print("Models loaded in LM Studio:")
|
||||
print("Models available in LM Studio:")
|
||||
for m in ms: print(" -", m)
|
||||
if not ms: print(" (none / server not reachable)")
|
||||
return
|
||||
|
||||
# --list/resident/unload-all work standalone; otherwise model/prompt/name required
|
||||
if not (args.model and args.prompt and args.name):
|
||||
ap.error("--model, --prompt, and --name are required (or use --list/--resident/--unload-all)")
|
||||
|
||||
# ---- MEMORY GUARD: refuse to stack a second big model ----
|
||||
resid = resident_instances(args.lmstudio)
|
||||
resident_ids = [iid for _, iid in resid]
|
||||
if resid and args.model not in resident_ids and not any(args.model in (iid or "") or args.model == k for k, iid in resid):
|
||||
print(f"!! ABORT: a DIFFERENT model is resident ({resid}) but you asked to run '{args.model}'.", file=sys.stderr)
|
||||
print(f"!! Unload it first: python3 tools/grade_run.py --unload-all (or load '{args.model}' and unload others)", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
prompt_file, ext = PROMPTS[args.prompt]
|
||||
ppath = os.path.join(HERE, prompt_file)
|
||||
if not os.path.exists(ppath):
|
||||
@@ -101,13 +169,33 @@ def main():
|
||||
|
||||
content = resp["choices"][0]["message"]["content"]
|
||||
usage = resp.get("usage", {})
|
||||
# LM Studio / OpenAI-style usage: prompt_tokens, completion_tokens, total_tokens
|
||||
# LM Studio reports timing under a stats/timings object — key names vary by
|
||||
# version, so self-discover: walk known containers and pick the first numeric hit.
|
||||
def _find(obj, keys):
|
||||
for k in keys:
|
||||
v = obj.get(k)
|
||||
if isinstance(v, (int, float)):
|
||||
return v
|
||||
return None
|
||||
stats = resp.get("stats") or {}
|
||||
timings = resp.get("timings") or {}
|
||||
# try stats then timings for each metric
|
||||
tok_sec = _find(stats, ["tokens_per_second", "tokensPerSecond", "predicted_per_second", "predicted_tokens_per_second"]) \
|
||||
or _find(timings, ["predicted_per_second", "tokens_per_second", "predicted_n"])
|
||||
ttft = _find(stats, ["time_to_first_token", "timeToFirstToken", "ttft", "first_token_time"]) \
|
||||
or _find(timings, ["time_to_first_token", "prompt_n", "prompt_per_second", "first_token"])
|
||||
comp_tokens = usage.get("completion_tokens") or usage.get("completion_tokens_details", {}).get("reasoning_tokens", 0)
|
||||
ttft = resp.get("stats", {}).get("time_to_first_token") or resp.get("timings", {}).get("prompt_n")
|
||||
tok_sec = resp.get("stats", {}).get("tokens_per_second") or resp.get("timings", {}).get("predicted_per_second")
|
||||
# fallback: derive tok/sec from wall time if the server didn't report it
|
||||
if tok_sec is None and comp_tokens and wall:
|
||||
tok_sec = round(comp_tokens / wall, 2)
|
||||
# one-time schema dump so we can harden the extraction (write next to output)
|
||||
if not os.environ.get("GRADE_RUN_NO_SCHEMA_DUMP"):
|
||||
try:
|
||||
schema_path = os.path.join(HERE, "outputs", ".last_response_schema.json")
|
||||
slim = {k: v for k, v in resp.items() if k not in ("choices", "id", "object", "model", "created")}
|
||||
json.dump(slim, open(schema_path, "w"), indent=2)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# save the raw output — extract the code block if the model wrapped it in
|
||||
# markdown fences or surrounded it with prose. Handles: (a) complete fenced
|
||||
|
||||
Reference in New Issue
Block a user