Gemma 4 26B-A4B head-to-head vs Qwen (4 prompts, via grade_run.py API)
Ran via tools/grade_run.py against LM Studio (no clipboard). Results: TTS: 80 Minor Flaws PASSES (real N-worker concurrency) <- Qwen 49, didn't parse Rust: 72 Minor Flaws COMPILES CLEAN (0 errs w/ deps) <- Qwen 50, 7 real errors Webhook: 55 Critical uses forbidden aiohttp (won't run) <- Qwen 75, passed Automation: 48 Critical SyntaxError (global-after-assign) <- first run for both DECISIVE head-to-head: Gemma generalizes where Qwen fails (TTS, Rust), but Qwen beats it on stdlib-discipline prompts (webhook). The two are COMPLEMENTARY local offloads, not redundant. Fixed: grade_run.py extractor (markdown/prose wrapping, multi-fence lang selection), TTFT-null handling in generator. TTFT capture from LM Studio API still needs the right stats key (left null + noted). Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Claude
@@ -509,6 +509,132 @@
|
||||
"Tests are mildly timing-flaky: 5% random 429 in discord_send + a tight '>1.0s' threshold; no tests for the 400 (bad JSON) or missing-signature 401 paths even though the code handles them."
|
||||
],
|
||||
"patch_code": "# FIX 1 (clock): use monotonic for TTL too.\nclass IdempotencyStore:\n def is_seen(self, event_id):\n now = time.monotonic()\n ...\n def mark(self, event_id):\n self.store[event_id] = time.monotonic()\n\n# FIX 2 (body cap): reject oversized bodies before reading.\nMAX_BODY = 64 * 1024\ncontent_length = int(headers.get('content-length', 0))\nif content_length > MAX_BODY:\n writer.write(b'HTTP/1.1 413 Payload Too Large\\r\\nContent-Length: 0\\r\\n\\r\\n'); await writer.drain(); return\nbody = await reader.readexactly(content_length) if 0 < content_length <= MAX_BODY else b''\n\n# FIX 3 (leak): bound forward_timestamps (deque maxlen=N) or drop it if unused.\nfrom collections import deque\nself.forward_timestamps: deque = deque(maxlen=1000)\n\n# FIX 4 (reason phrase): use a fixed map.\nREASON = {200:'OK',400:'Bad Request',401:'Unauthorized',404:'Not Found',502:'Bad Gateway',500:'Internal Server Error'}\nresponse = f'HTTP/1.1 {status} {REASON.get(status,\"OK\")}\\r\\n...'\n\n# FIX 5: add tests for the 400 (malformed JSON) and missing-signature 401 paths."
|
||||
},
|
||||
{
|
||||
"id": "gemma4-26b-a4b-8bit-mlx-tts",
|
||||
"prompt_id": "tts",
|
||||
"timestamp": "2026-07-29T01:30:00Z",
|
||||
"model_name": "Gemma 4 26B-A4B",
|
||||
"quant": "8-bit MLX",
|
||||
"param_size": "26B-A4B (MoE)",
|
||||
"format": "mlx",
|
||||
"lang": "py",
|
||||
"tok_sec": 49.48,
|
||||
"total_tokens": 6876,
|
||||
"ttft_sec": null,
|
||||
"ttft_note": "LM Studio API TTFT key not captured by grade_run.py yet",
|
||||
"filename": "outputs/gemma4-26b-a4b-8bit-mlx-tts.py",
|
||||
"tests_pass": true,
|
||||
"total_score": 80,
|
||||
"breakdown": {
|
||||
"complexity": 16,
|
||||
"concurrency": 18,
|
||||
"error_handling": 16,
|
||||
"resource_safety": 15,
|
||||
"test_integrity": 15
|
||||
},
|
||||
"verdict": "Minor Logic Flaws",
|
||||
"best_for": "DECISIVE: passes the TTS pipeline that Qwen 6-bit couldn't even parse. Uses REAL bounded concurrency (N worker tasks + Semaphore) \u2014 exactly Qwen's fatal flaw avoided. Gemma generalizes to multi-task orchestration where Qwen fails. Strong offload candidate for queue/pipeline work.",
|
||||
"critical_bugs": [
|
||||
"LM Studio API TTFT not captured by grade_run.py (key-name mismatch) \u2014 needs a fix to the script.",
|
||||
"Detail audit pending a full read, but tests pass and the concurrency model is correct (N create_task workers + Semaphore, proper drain via queue.join + running-count, cancel + gather)."
|
||||
],
|
||||
"patch_code": "# TTFT capture: LM Studio returns timing under 'stats' or 'timings' with keys like\n# 'time_to_first_token' / 'prompt_progress' \u2014 grade_run.py should dump resp['stats'] raw\n# once to find the right key, then parse it."
|
||||
},
|
||||
{
|
||||
"id": "gemma4-26b-a4b-8bit-mlx-rust",
|
||||
"prompt_id": "rust",
|
||||
"timestamp": "2026-07-29T01:30:00Z",
|
||||
"model_name": "Gemma 4 26B-A4B",
|
||||
"quant": "8-bit MLX",
|
||||
"param_size": "26B-A4B (MoE)",
|
||||
"format": "mlx",
|
||||
"lang": "rs",
|
||||
"tok_sec": 58.56,
|
||||
"total_tokens": 4972,
|
||||
"ttft_sec": null,
|
||||
"ttft_note": "LM Studio API TTFT key not captured by grade_run.py yet",
|
||||
"filename": "outputs/gemma4-26b-a4b-8bit-mlx-rust.rs",
|
||||
"tests_pass": false,
|
||||
"total_score": 72,
|
||||
"breakdown": {
|
||||
"ownership": 17,
|
||||
"concurrency": 15,
|
||||
"error_handling": 14,
|
||||
"cancellation": 13,
|
||||
"test_integrity": 13
|
||||
},
|
||||
"verdict": "Minor Logic Flaws",
|
||||
"best_for": "DECISIVE: Rust COMPILES CLEAN (0 errors with deps declared) where Qwen had 7 real compile errors (mpsc::bounded hallucination, ownership moves). The code is structurally sound Rust. Deductions only for undeclared deps + markdown wrapping. Gemma is the better Rust offload pick by a wide margin.",
|
||||
"critical_bugs": [
|
||||
"Uses tokio-util and rand crates WITHOUT declaring them in the dependency block (prompt said assume tokio/serde/thiserror only). Compiles clean once added \u2014 so it's a deps-list omission, not a code bug.",
|
||||
"Wrapped output in markdown fences (```toml and ```rust) \u2014 required extraction to get runnable code. A submission-hygiene issue, not a logic one.",
|
||||
"Full logic audit pending; 0 compile errors with deps declared is the headline."
|
||||
],
|
||||
"patch_code": "# FIX 1 (deps): add to Cargo.toml:\n# tokio-util = \"0.7\"\n# rand = \"0.8\"\n# FIX 2: grade_run.py extractor now prefers the ```rust fence; this won't recur."
|
||||
},
|
||||
{
|
||||
"id": "gemma4-26b-a4b-8bit-mlx-webhook",
|
||||
"prompt_id": "webhook",
|
||||
"timestamp": "2026-07-29T01:30:00Z",
|
||||
"model_name": "Gemma 4 26B-A4B",
|
||||
"quant": "8-bit MLX",
|
||||
"param_size": "26B-A4B (MoE)",
|
||||
"format": "mlx",
|
||||
"lang": "py",
|
||||
"tok_sec": 59.21,
|
||||
"total_tokens": 5565,
|
||||
"ttft_sec": null,
|
||||
"ttft_note": "LM Studio API TTFT key not captured by grade_run.py yet",
|
||||
"filename": "outputs/gemma4-26b-a4b-8bit-mlx-webhook.py",
|
||||
"tests_pass": false,
|
||||
"total_score": 55,
|
||||
"breakdown": {
|
||||
"schema_io": 14,
|
||||
"transport": 12,
|
||||
"error_handling": 13,
|
||||
"state_safety": 13,
|
||||
"test_integrity": 13
|
||||
},
|
||||
"verdict": "Critical Bugs",
|
||||
"best_for": "Violated the stdlib-only constraint: used aiohttp (2 imports) for the HTTP server, so it won't run as-is. Weaker than Qwen's 75 on the same prompt (Qwen used stdlib and passed). Not the offload pick for webhook work; Qwen is.",
|
||||
"critical_bugs": [
|
||||
"FATAL for the prompt: uses aiohttp (external dep) despite 'stdlib only' requirement. ModuleNotFoundError on import \u2014 won't run as delivered.",
|
||||
"Chose aiohttp reasoning that stdlib http.server is synchronous/blocks the loop \u2014 a fair architectural point, but it violated the explicit constraint instead of using asyncio.start_server (stdlib, async) like Qwen did.",
|
||||
"Wrapped in markdown + prose preamble (needed extraction)."
|
||||
],
|
||||
"patch_code": "# FIX: replace aiohttp server with asyncio.start_server (stdlib, fully async) \u2014 exactly\n# what the Qwen webhook submission did. Then the logic (rate limit/idempotency/429) can run."
|
||||
},
|
||||
{
|
||||
"id": "gemma4-26b-a4b-8bit-mlx-automation",
|
||||
"prompt_id": "automation",
|
||||
"timestamp": "2026-07-29T01:30:00Z",
|
||||
"model_name": "Gemma 4 26B-A4B",
|
||||
"quant": "8-bit MLX",
|
||||
"param_size": "26B-A4B (MoE)",
|
||||
"format": "mlx",
|
||||
"lang": "py",
|
||||
"tok_sec": 59.25,
|
||||
"total_tokens": 5499,
|
||||
"ttft_sec": null,
|
||||
"ttft_note": "LM Studio API TTFT key not captured by grade_run.py yet",
|
||||
"filename": "outputs/gemma4-26b-a4b-8bit-mlx-automation.py",
|
||||
"tests_pass": false,
|
||||
"total_score": 48,
|
||||
"breakdown": {
|
||||
"idempotency": 14,
|
||||
"retry_backoff": 13,
|
||||
"checkpointing": 13,
|
||||
"signal_handling": 10,
|
||||
"test_integrity": 13
|
||||
},
|
||||
"verdict": "Critical Bugs",
|
||||
"best_for": "Doesn't run: SyntaxError (global processor declared after assignment in the signal-handler setup). A real but localized bug. First automation result for either model, so no head-to-head yet.",
|
||||
"critical_bugs": [
|
||||
"FATAL: line 270 'global processor' declared AFTER processor is assigned earlier in the same function -> SyntaxError. The whole module fails to parse.",
|
||||
"Signal-handler design (global instance for SIGINT access) is the root cause \u2014 globals-after-assignment is a classic Python footgun the model walked into."
|
||||
],
|
||||
"patch_code": "# FIX: move 'global processor' to the FIRST line of the function that assigns it,\n# before any assignment to processor."
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -582,10 +582,12 @@ def render_detail(m, data):
|
||||
patch = m.get("patch_code", "")
|
||||
metrics_block = ""
|
||||
if m.get("tok_sec") is not None:
|
||||
ttft = m.get("ttft_sec")
|
||||
ttft_str = f"{ttft:.2f}s" if isinstance(ttft, (int, float)) else "—"
|
||||
metrics_block = f"""
|
||||
<div class="mini"><span class="mlbl">tok/sec</span><span class="mval mono" style="color:var(--mag)">{m['tok_sec']:.2f}</span></div>
|
||||
<div class="mini"><span class="mlbl">tokens</span><span class="mval mono">{m.get('total_tokens') or '—'}</span></div>
|
||||
<div class="mini"><span class="mlbl">TTFT</span><span class="mval mono">{m.get('ttft_sec'):.2f}s</span> </div>"""
|
||||
<div class="mini"><span class="mlbl">TTFT</span><span class="mval mono">{ttft_str}</span> </div>"""
|
||||
else:
|
||||
metrics_block = '<div class="mini" style="grid-column:1/-1"><span class="mlbl">runtime</span><span class="mval mono" style="color:var(--blue)">CLOUD — not measured</span></div>'
|
||||
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import signal
|
||||
import time
|
||||
from dataclasses import dataclass, asdict, field
|
||||
from pathlib import Path
|
||||
from typing import Set, List, Dict, Any
|
||||
|
||||
# --- Domain Exceptions ---
|
||||
|
||||
class ProcessingError(Exception):
|
||||
"""Raised when the flaky external call fails."""
|
||||
pass
|
||||
|
||||
# --- Models ---
|
||||
|
||||
@dataclass
|
||||
class Checkpoint:
|
||||
completed_ids: Set[str] = field(default_factory=set)
|
||||
failed_ids: Set[str] = field(default_factory=set)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"completed_ids": list(self.completed_ids),
|
||||
"failed_ids": list(self.failed_ids)
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Checkpoint":
|
||||
return cls(
|
||||
completed_ids=set(data.get("completed_ids", [])),
|
||||
failed_ids=set(data.get("failed_ids", []))
|
||||
)
|
||||
|
||||
# --- Core Logic ---
|
||||
|
||||
class BatchProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
items: List[str],
|
||||
checkpoint_path: str = "checkpoint.json",
|
||||
max_concurrency: int = 8,
|
||||
max_retries: int = 3
|
||||
):
|
||||
self.items = items
|
||||
self.checkpoint_path = Path(checkpoint_path)
|
||||
self.semaphore = asyncio.Semaphore(max_concurrency)
|
||||
self.max_retries = max_retries
|
||||
self.checkpoint = Checkpoint()
|
||||
self.shutdown_event = asyncio.Event()
|
||||
self.start_time = 0.0
|
||||
self._pending_tasks: Set[asyncio.Task] = set()
|
||||
|
||||
def _load_checkpoint(self):
|
||||
if self.checkpoint_path.exists():
|
||||
try:
|
||||
with open(self.checkpoint_path, "r") as f:
|
||||
data = json.load(f)
|
||||
self.checkpoint = Checkpoint.from_dict(data)
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
# If corrupt, we start fresh or could raise error.
|
||||
# Requirement says "crash must leave valid checkpoint", so we assume it's valid.
|
||||
pass
|
||||
|
||||
def _save_checkpoint(self):
|
||||
"""Atomic write using a temporary file and os.replace."""
|
||||
temp_path = self.checkpoint_path.with_suffix(".tmp")
|
||||
try:
|
||||
with open(temp_path, "w") as f:
|
||||
json.dump(self.checkpoint.to_dict(), f)
|
||||
# os.replace is atomic on both POSIX and Windows (for existing files)
|
||||
os.replace(temp_path, self.checkpoint_path)
|
||||
except Exception as e:
|
||||
print(f"Failed to save checkpoint: {e}")
|
||||
|
||||
async def _process_item_with_retry(self, item: str, process_func):
|
||||
"""Handles the retry logic and backoff for a single item."""
|
||||
async with self.semaphore:
|
||||
# Check if we are shutting down before starting a new item
|
||||
if self.shutdown_event.is_set():
|
||||
return
|
||||
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
await process_func(item)
|
||||
# Success path
|
||||
self.checkpoint.completed_ids.add(item)
|
||||
self._save_checkpoint()
|
||||
return "succeeded"
|
||||
except ProcessingError:
|
||||
if attempt < self.max_retries:
|
||||
# Exponential backoff: 0.1, 0.2, 0.4...
|
||||
wait_time = 0.1 * (2 ** attempt)
|
||||
await asyncio.sleep(wait_time)
|
||||
else:
|
||||
# Terminal failure after retries
|
||||
self.checkpoint.failed_ids.add(item)
|
||||
self._save_checkpoint()
|
||||
return "failed"
|
||||
except asyncio.CancelledError:
|
||||
# If the task is cancelled (e.g. during shutdown),
|
||||
# we don't record it as completed or failed.
|
||||
raise
|
||||
|
||||
async def run(self, process_func) -> Dict[str, Any]:
|
||||
self._load_checkpoint()
|
||||
self.start_time = time.perf_counter()
|
||||
|
||||
# Identify items to skip
|
||||
to_process = [
|
||||
i for i in self.items
|
||||
if i not in self.checkpoint.completed_ids and i not in self.checkpoint.failed_ids
|
||||
]
|
||||
skipped_count = len(self.items) - len(to_process) - len(self.checkpoint.failed_ids)
|
||||
# Note: skipped is items already in checkpoint.completed_ids
|
||||
skipped_count = len([i for i in self.items if i in self.checkpoint.completed_ids])
|
||||
|
||||
# We use a set to track tasks so we can await them during graceful shutdown
|
||||
tasks = []
|
||||
for item in to_process:
|
||||
if self.shutdown_event.is_set():
|
||||
break
|
||||
task = asyncio.create_task(self._process_item_with_retry(item, process_func))
|
||||
tasks.append(task)
|
||||
self._pending_tasks.add(task)
|
||||
task.add_done_callback(self._pending_tasks.discard)
|
||||
|
||||
# Wait for all spawned tasks to finish
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
elapsed_ms = int((time.perf_counter() - self.start_time) * 1000)
|
||||
|
||||
# Final Summary Calculation
|
||||
# We must re-calculate based on the final state of checkpoint + what was skipped
|
||||
succeeded = len(self.checkpoint.completed_ids)
|
||||
failed = len(self.checkpoint.failed_ids)
|
||||
# Items that were already done before we started
|
||||
already_done = len([i for i in self.items if i in self.checkpoint.completed_ids and i not in to_process])
|
||||
# Actually, the requirement asks for a summary of this run's context.
|
||||
# Let's provide the absolute state:
|
||||
return {
|
||||
"succeeded": succeeded,
|
||||
"failed": failed,
|
||||
"skipped": skipped_count,
|
||||
"total": len(self.items),
|
||||
"elapsed_ms": elapsed_ms
|
||||
}
|
||||
|
||||
# --- Mock External Service ---
|
||||
|
||||
async def flaky_process(item: str) -> str:
|
||||
"""The provided flaky function."""
|
||||
# Simulate random latency 10-80ms
|
||||
await asyncio.sleep(random.uniform(0.01, 0.08))
|
||||
# 80% success rate
|
||||
if random.random() < 0.2:
|
||||
raise ProcessingError(f"Failed to process {item}")
|
||||
return f"Result for {item}"
|
||||
|
||||
# --- Tests ---
|
||||
|
||||
async def run_tests():
|
||||
print("\n--- Starting Tests ---")
|
||||
checkpoint_file = "test_checkpoint.json"
|
||||
if os.path.exists(checkpoint_file):
|
||||
os.remove(checkpoint_file)
|
||||
|
||||
# 1. Concurrency Test
|
||||
print("Test: Bounded Concurrency (Max 8)...", end="", flush=True)
|
||||
active_count = 0
|
||||
max_observed = 0
|
||||
lock = asyncio.Lock()
|
||||
|
||||
async def concurrency_mock(item: str):
|
||||
nonlocal active_count, max_observed
|
||||
async with lock:
|
||||
active_count += 1
|
||||
max_observed = max(max_observed, active_count)
|
||||
await asyncio.sleep(0.05)
|
||||
async with lock:
|
||||
active_count -= 1
|
||||
|
||||
proc = BatchProcessor(["item-" + str(i) for i in range(20)], checkpoint_path=checkpoint_file, max_concurrency=8)
|
||||
await proc.run(concurrency_mock)
|
||||
assert max_observed <= 8, f"Concurrency exceeded! Observed {max_observed}"
|
||||
print("PASSED")
|
||||
|
||||
# 2. Integrity Test (Watcher)
|
||||
print("Test: Checkpoint Integrity (JSON Validity)...", end="", flush=True)
|
||||
stop_watcher = asyncio.Event()
|
||||
|
||||
async def integrity_watcher():
|
||||
while not stop_watcher.is_set():
|
||||
if os.path.exists(checkpoint_file):
|
||||
try:
|
||||
with open(checkpoint_file, "r") as f:
|
||||
json.load(f)
|
||||
except Exception as e:
|
||||
raise AssertionError(f"Corrupt JSON detected: {e}")
|
||||
await asyncio.sleep(0.005)
|
||||
|
||||
watcher_task = asyncio.create_task(integrity_watcher())
|
||||
proc = BatchProcessor(["item-" + str(i) for i in range(50)], checkpoint_path=checkpoint_file)
|
||||
await proc.run(flaky_process)
|
||||
stop_watcher.set()
|
||||
await watcher_task
|
||||
print("PASSED")
|
||||
|
||||
# 3. Resumability / Idempotency Test
|
||||
print("Test: Resumability & Idempotency...", end="", flush=True)
|
||||
items = [f"resumable-{i}" for i in range(100)]
|
||||
checkpoint_file = "resumable.json"
|
||||
if os.path.exists(checkpoint_file): os.remove(checkpoint_file)
|
||||
|
||||
# Run 1: Start and interrupt mid-way
|
||||
proc = BatchProcessor(items, checkpoint_path=checkpoint_file)
|
||||
|
||||
# We simulate an interruption by cancelling the main task after a short delay
|
||||
main_task = asyncio.create_task(proc.run(flaky_process))
|
||||
await asyncio.sleep(0.2)
|
||||
proc.shutdown_event.set() # Signal shutdown
|
||||
# We don't cancel the task immediately to allow in-flight items to finish (graceful)
|
||||
await main_task
|
||||
|
||||
# Run 2: Restart
|
||||
proc_retry = BatchProcessor(items, checkpoint_path=checkpoint_file)
|
||||
final_summary = await proc_retry.run(flaky_process)
|
||||
|
||||
# Assertions
|
||||
# 1. All items must eventually be in completed or failed sets
|
||||
all_processed = proc_retry.checkpoint.completed_ids | proc_retry.checkpoint.failed_ids
|
||||
assert len(all_processed) == len(items), f"Not all items processed. Got {len(all_processed)}/{len(items)}"
|
||||
# 2. Checkpoint must be valid and contain all items
|
||||
assert final_summary["total"] == 100
|
||||
print("PASSED")
|
||||
|
||||
# Cleanup
|
||||
for f in [checkpoint_file, "test_checkpoint.json", "resumable.json"]:
|
||||
if os.path.exists(f): os.remove(f)
|
||||
print("--- All Tests Passed Successfully ---\n")
|
||||
|
||||
# --- Main Entry Point ---
|
||||
|
||||
async def main():
|
||||
# Setup signal handling for SIGINT (Ctrl+C)
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
# We'll use a global processor instance to allow the signal handler access
|
||||
processor = None
|
||||
|
||||
def handle_sigint():
|
||||
if processor:
|
||||
print("\n[!] SIGINT received. Shutting down gracefully (finishing in-flight tasks)...")
|
||||
processor.shutdown_event.set()
|
||||
|
||||
# Register signal handler
|
||||
try:
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(sig, handle_sigint)
|
||||
except NotImplementedError:
|
||||
# signal.add_signal_handler is not implemented on Windows in some Python versions
|
||||
pass
|
||||
|
||||
# Input Data
|
||||
work_items = [f"job-{i:03d}" for i in range(100)]
|
||||
|
||||
global processor
|
||||
processor = BatchProcessor(work_items, max_concurrency=8)
|
||||
|
||||
try:
|
||||
summary = await processor.run(flaky_process)
|
||||
print(json.dumps(summary))
|
||||
except KeyboardInterrupt:
|
||||
# Fallback for environments where signal handler doesn't catch it
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"Fatal error: {e}")
|
||||
finally:
|
||||
# If we are running tests, don't exit yet
|
||||
import sys
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "--test":
|
||||
pass
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "--test":
|
||||
asyncio.run(run_tests())
|
||||
else:
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
# Final catch for the top-level loop
|
||||
sys.exit(0)
|
||||
@@ -0,0 +1,317 @@
|
||||
use rand::Rng;
|
||||
use serde::Serialize;
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
use thiserror::Error;
|
||||
use tokio::sync::{mpsc, RwLock};
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
// --- Types and Errors ---
|
||||
|
||||
#[derive(Debug, Error, Clone)]
|
||||
pub enum WatcherError {
|
||||
#[error("Failed to fetch data from source")]
|
||||
FetchError,
|
||||
#[error("Channel closed unexpectedly")]
|
||||
ChannelClosed,
|
||||
#[error("Watcher exceeded max consecutive failures")]
|
||||
Unhealthy,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct WatchItem {
|
||||
pub watcher_id: u32,
|
||||
pub data: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct AggregatedReport {
|
||||
pub total_items: usize,
|
||||
pub items: Vec<WatchItem>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct WatcherId(u32);
|
||||
|
||||
// --- Mock Source ---
|
||||
|
||||
/// Simulates an external API/Source.
|
||||
/// Returns 0-3 items or fails ~15% of the time.
|
||||
async fn mock_fetch(id: u32, force_fail: bool) -> Result<Vec<String>, WatcherError> {
|
||||
// Simulate network latency
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
let mut rng = rand::thread_rng();
|
||||
if force_fail || rng.gen_bool(0.15) {
|
||||
return Err(WatcherError::FetchError);
|
||||
}
|
||||
|
||||
let count = rng.gen_range(0..=3);
|
||||
Ok((0..count).map(|i| format!("val-{}-{}", id, i)).collect())
|
||||
}
|
||||
|
||||
// --- Watcher Logic ---
|
||||
|
||||
struct Watcher {
|
||||
id: WatcherId,
|
||||
tx: mpsc::Sender<WatchItem>,
|
||||
token: CancellationToken,
|
||||
force_fail: bool,
|
||||
}
|
||||
|
||||
impl Watcher {
|
||||
async fn run(self) -> Result<(), WatcherError> {
|
||||
let mut consecutive_failures = 0;
|
||||
const MAX_FAILURES: u32 = 5;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// 1. Handle Shutdown Signal
|
||||
_ = self.token.cancelled() => {
|
||||
println!("[Watcher {:?}] Shutting down gracefully", self.id);
|
||||
return Ok(());
|
||||
}
|
||||
// 2. Periodic Polling
|
||||
_ = tokio::time::sleep(Duration::from_millis(200)) => {
|
||||
match mock_fetch(self.id.0, self.force_fail).await {
|
||||
Ok(items) => {
|
||||
consecutive_failures = 0; // Reset on success
|
||||
for item in items {
|
||||
let msg = WatchItem { watcher_id: self.id.0, data: item };
|
||||
// Backpressure: .send() awaits if the channel is full.
|
||||
// This prevents a fast watcher from overwhelming the consumer.
|
||||
if self.tx.send(msg).await.is_err() {
|
||||
return Err(WatcherError::ChannelClosed);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
consecutive_failures += 1;
|
||||
eprintln!("[Watcher {:?}] Error ({}): {}", self.id, consecutive_failures, e);
|
||||
|
||||
if consecutive_failures >= MAX_FAILURES {
|
||||
eprintln!("[Watcher {:?}] Marked UNHEALTHY. Terminating.", self.id);
|
||||
return Err(WatcherError::Unhealthy);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Manager Logic ---
|
||||
|
||||
pub struct WatcherManager {
|
||||
/// Map of active watcher handles. Protected by RwLock for concurrent add/remove.
|
||||
watchers: Arc<RwLock<HashMap<WatcherId, JoinHandle<()>>>>,
|
||||
/// Channel to send items to the consumer.
|
||||
tx: mpsc::Sender<WatchItem>,
|
||||
/// Token to signal all tasks to stop.
|
||||
shutdown_token: CancellationToken,
|
||||
}
|
||||
|
||||
impl WatcherManager {
|
||||
pub fn new(buffer_size: usize) -> (Self, mpsc::Receiver<WatchItem>, CancellationToken) {
|
||||
let (tx, rx) = mpsc::channel(buffer_size);
|
||||
let token = CancellationToken::new();
|
||||
|
||||
let manager = Self {
|
||||
watchers: Arc::new(RwLock::new(HashMap::new())),
|
||||
tx,
|
||||
shutdown_token: token.clone(),
|
||||
};
|
||||
|
||||
(manager, rx, token)
|
||||
}
|
||||
|
||||
pub async fn add_watcher(&self, id: u32, force_fail: bool) {
|
||||
let watcher_id = WatcherId(id);
|
||||
let watcher = Watcher {
|
||||
id: watcher_id,
|
||||
tx: self.tx.clone(),
|
||||
token: self.shutdown_token.clone(),
|
||||
force_fail,
|
||||
};
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
if let Err(e) = watcher.run().await {
|
||||
eprintln!("[Manager] Watcher {:?} exited with error: {}", id, e);
|
||||
}
|
||||
});
|
||||
|
||||
let mut watchers = self.watchers.write().await;
|
||||
watchers.insert(watcher_id, handle);
|
||||
println!("[Manager] Added watcher {}", id);
|
||||
}
|
||||
|
||||
pub async fn remove_watcher(&self, id: u32) {
|
||||
let watcher_id = WatcherId(id);
|
||||
let mut watchers = self.watchers.write().await;
|
||||
if let Some(handle) = watchers.remove(&watcher_id) {
|
||||
// Abort the task immediately if it's still running
|
||||
handle.abort();
|
||||
println!("[Manager] Removed watcher {}", id);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_active_count(&self) -> usize {
|
||||
self.watchers.read().await.len()
|
||||
}
|
||||
|
||||
pub async fn shutdown(self) {
|
||||
println!("[Manager] Initiating shutdown...");
|
||||
self.shutdown_token.cancel();
|
||||
|
||||
// Wait for all handles to finish
|
||||
let mut watchers = self.watchers.write().await;
|
||||
for (id, handle) in watchers.drain() {
|
||||
let _ = handle.await;
|
||||
println!("[Manager] Joined watcher {:?}", id.0);
|
||||
}
|
||||
println!("[Manager] All watchers stopped.");
|
||||
}
|
||||
}
|
||||
|
||||
// --- Consumer Logic ---
|
||||
|
||||
async fn run_consumer(mut rx: mpsc::Receiver<WatchItem>, token: CancellationToken) -> AggregatedReport {
|
||||
let mut items = Vec::new();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = token.cancelled() => {
|
||||
println!("[Consumer] Shutdown signal received.");
|
||||
break;
|
||||
}
|
||||
msg = rx.recv() => {
|
||||
match msg {
|
||||
Some(item) => items.push(item),
|
||||
None => break, // Channel closed
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
AggregatedReport {
|
||||
total_items: items.len(),
|
||||
items,
|
||||
}
|
||||
}
|
||||
|
||||
// --- Tests ---
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_lifecycle_and_aggregation() {
|
||||
let (manager, rx, token) = WatcherManager::new(100);
|
||||
let manager = Arc::new(manager);
|
||||
|
||||
// Spawn consumer
|
||||
let consumer_token = token.clone();
|
||||
let consumer_handle = tokio::spawn(async move {
|
||||
run_consumer(rx, consumer_token).await
|
||||
});
|
||||
|
||||
// Add 4 watchers
|
||||
for i in 1..=4 {
|
||||
manager.add_watcher(i, false).await;
|
||||
}
|
||||
|
||||
// Let them run for a bit
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
|
||||
// Shutdown
|
||||
let m = Arc::try_unwrap(manager).ok().expect("Manager should have no other owners");
|
||||
m.shutdown().await;
|
||||
token.cancel(); // Ensure consumer stops
|
||||
|
||||
let report = consumer_handle.await.unwrap();
|
||||
|
||||
assert!(report.total_items > 0, "Consumer should have received items");
|
||||
println!("Test A Passed: Received {} items", report.total_items);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unhealthy_watcher() {
|
||||
let (manager, rx, token) = WatcherManager::new(10);
|
||||
let manager = Arc::new(manager);
|
||||
|
||||
// Spawn consumer to prevent channel blockage
|
||||
let ct = token.clone();
|
||||
tokio::spawn(async move { run_consumer(rx, ct).await });
|
||||
|
||||
// Add a watcher that always fails
|
||||
manager.add_watcher(99, true).await;
|
||||
|
||||
// Wait for it to hit 5 failures (5 * 200ms + buffer)
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
|
||||
// The watcher should have exited.
|
||||
// Note: In this implementation, the manager doesn't automatically remove
|
||||
// dead handles from the map unless we implement a reaper, but we check if it's still "active"
|
||||
// by observing that the task is no longer running.
|
||||
let count = manager.get_active_count().await;
|
||||
assert!(count >= 1); // It's still in the map, but the task is dead.
|
||||
// In a real system, we'd have a background reaper task to clean the HashMap.
|
||||
println!("Test B Passed: Unhealthy watcher handled.");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_concurrent_mutation() {
|
||||
let (manager, rx, token) = WatcherManager::new(10);
|
||||
let manager = Arc::new(manager);
|
||||
tokio::spawn(async move { run_consumer(rx, token).await });
|
||||
|
||||
let mut handles = vec![];
|
||||
|
||||
// Hammer the manager with concurrent adds and removes
|
||||
for i in 0..50 {
|
||||
let m = manager.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
m.add_watcher(i, false).await;
|
||||
if i % 2 == 0 {
|
||||
m.remove_watcher(i).await;
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
for h in handles {
|
||||
h.await.unwrap();
|
||||
}
|
||||
|
||||
let final_count = manager.get_active_count().await;
|
||||
println!("Test C Passed: Concurrent mutation finished with {} watchers", final_count);
|
||||
assert!(final_count <= 50);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("Starting Service...");
|
||||
let (manager, rx, token) = WatcherManager::new(100);
|
||||
|
||||
// Spawn consumer
|
||||
let ct = token.clone();
|
||||
let consumer_handle = tokio::spawn(async move {
|
||||
run_consumer(rx, ct).await
|
||||
});
|
||||
|
||||
// Add some watchers
|
||||
manager.add_watcher(1, false).await;
|
||||
manager.add_watcher(2, false).await;
|
||||
manager.add_watcher(3, true).await; // This one will fail
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
|
||||
// Shutdown
|
||||
manager.shutdown().await;
|
||||
token.cancel(); // Signal consumer
|
||||
|
||||
let report = consumer_handle.await.unwrap();
|
||||
println!("Final Report: {:?}", report);
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
import asyncio
|
||||
import random
|
||||
import uuid
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum, auto
|
||||
from typing import Callable, Dict, Optional, Set
|
||||
|
||||
|
||||
class JobStatus(Enum):
|
||||
QUEUED = auto()
|
||||
STARTED = auto()
|
||||
COMPLETED = auto()
|
||||
FAILED = auto()
|
||||
CANCELLED = auto()
|
||||
|
||||
|
||||
@dataclass
|
||||
class Job:
|
||||
id: str
|
||||
text: str
|
||||
voice: str
|
||||
status: JobStatus = JobStatus.QUEUED
|
||||
retries_attempted: int = 0
|
||||
cancelled: bool = False
|
||||
|
||||
|
||||
async def mock_synthesize(text: str) -> bytes:
|
||||
"""
|
||||
Simulates a TTS engine.
|
||||
Sleeps 50-300ms, fails ~10% of the time.
|
||||
"""
|
||||
await asyncio.sleep(random.uniform(0.05, 0.3))
|
||||
if random.random() < 0.1:
|
||||
raise Exception("Synthetic Synthesis Error")
|
||||
return b"\x00" * len(text)
|
||||
|
||||
|
||||
class TTSPipeline:
|
||||
def __init__(self, max_concurrency: int = 4, max_queue_size: int = 100):
|
||||
self.max_concurrency = max_concurrency
|
||||
self.max_queue_size = max_queue_size
|
||||
|
||||
self._queue: asyncio.Queue[str] = asyncio.Queue()
|
||||
self._jobs: Dict[str, Job] = {}
|
||||
self._workers: Set[asyncio.Task] = set()
|
||||
self._callback: Optional[Callable[[str, JobStatus], None]] = None
|
||||
|
||||
self._semaphore = asyncio.Semaphore(max_concurrency)
|
||||
self._running_count = 0 # Tracks jobs in queue + jobs in flight
|
||||
self._lock = asyncio.Lock() # Protects shared state updates
|
||||
|
||||
# Metrics for testing
|
||||
self._current_active_count = 0
|
||||
self._peak_concurrency = 0
|
||||
|
||||
def set_callback(self, cb: Callable[[str, JobStatus], None]):
|
||||
self._callback = cb
|
||||
|
||||
def _emit(self, job_id: str, status: JobStatus):
|
||||
if self._callback:
|
||||
self._callback(job_id, status)
|
||||
|
||||
async def start(self):
|
||||
"""Initialize the worker pool."""
|
||||
for _ in range(self.max_concurrency):
|
||||
worker = asyncio.create_task(self._worker())
|
||||
self._workers.add(worker)
|
||||
|
||||
async def stop(self):
|
||||
"""Gracefully shut down workers."""
|
||||
for worker in self._workers:
|
||||
worker.cancel()
|
||||
await asyncio.gather(*self._workers, return_exceptions=True)
|
||||
self._workers.clear()
|
||||
|
||||
async def submit(self, text: str, voice: str) -> str:
|
||||
"""Enqueue a job. Returns job_id immediately."""
|
||||
if self._queue.qsize() >= self.max_queue_size:
|
||||
raise BufferError("Pipeline queue is full (backpressure)")
|
||||
|
||||
job_id = str(uuid.uuid4())
|
||||
job = Job(id=job_id, text=text, voice=voice)
|
||||
|
||||
async with self._lock:
|
||||
self._jobs[job_id] = job
|
||||
self._running_count += 1
|
||||
|
||||
await self._queue.put(job_id)
|
||||
self._emit(job_id, JobStatus.QUEUED)
|
||||
return job_id
|
||||
|
||||
async def cancel(self, job_id: str):
|
||||
"""Cancel a queued or in-flight job."""
|
||||
async with self._lock:
|
||||
if job_id in self._jobs:
|
||||
job = self._jobs[job_id]
|
||||
if job.status not in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
|
||||
job.cancelled = True
|
||||
|
||||
async def drain(self):
|
||||
"""Wait until all queued and in-flight jobs are finished."""
|
||||
# 1. Wait for the queue to be empty (all items processed by workers)
|
||||
await self._queue.join()
|
||||
# 2. Wait for the running count to hit zero (all workers finished their logic)
|
||||
while True:
|
||||
async with self._lock:
|
||||
if self._running_count == 0:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def _worker(self):
|
||||
while True:
|
||||
job_id = await self._queue.get()
|
||||
try:
|
||||
async with self._lock:
|
||||
job = self._jobs.get(job_id)
|
||||
|
||||
if not job:
|
||||
continue
|
||||
|
||||
# Check if cancelled while sitting in queue
|
||||
if job.cancelled:
|
||||
await self._finalize_job(job, JobStatus.CANCELLED)
|
||||
continue
|
||||
|
||||
# Process job with bounded concurrency and retries
|
||||
async with self._semaphore:
|
||||
await self._process_with_retries(job)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
finally:
|
||||
self._queue.task_done()
|
||||
|
||||
async def _process_with_retries(self, job: Job):
|
||||
max_retries = 3
|
||||
backoff = 0.1
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
# Check cancellation before starting/retrying
|
||||
if job.cancelled:
|
||||
await self._finalize_job(job, JobStatus.CANCELLED)
|
||||
return
|
||||
|
||||
try:
|
||||
# Track concurrency for testing purposes
|
||||
async with self._lock:
|
||||
self._current_active_count += 1
|
||||
self._peak_concurrency = max(self._peak_concurrency, self._current_active_count)
|
||||
|
||||
self._emit(job.id, JobStatus.STARTED)
|
||||
|
||||
# The actual work
|
||||
await mock_synthesize(job.text)
|
||||
|
||||
# Check cancellation after the await (in case it was cancelled during synthesis)
|
||||
if job.cancelled:
|
||||
await self._finalize_job(job, JobStatus.CANCELLED)
|
||||
else:
|
||||
await self._finalize_job(job, JobStatus.COMPLETED)
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
if job.cancelled:
|
||||
await self._finalize_job(job, JobStatus.CANCELLED)
|
||||
return
|
||||
|
||||
if attempt < max_retries:
|
||||
# Exponential backoff
|
||||
await asyncio.sleep(backoff * (2 ** attempt))
|
||||
else:
|
||||
# Exhausted retries
|
||||
await self._finalize_job(job, JobStatus.FAILED)
|
||||
finally:
|
||||
async with self._lock:
|
||||
self._current_active_count -= 1
|
||||
|
||||
async def _finalize_job(self, job: Job, status: JobStatus):
|
||||
"""Updates job state and decrements the global running counter."""
|
||||
async with self._lock:
|
||||
if job.status not in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
|
||||
job.status = status
|
||||
self._running_count -= 1
|
||||
self._emit(job.id, status)
|
||||
|
||||
|
||||
# ==========================================
|
||||
# TEST SUITE
|
||||
# ==========================================
|
||||
|
||||
async def main():
|
||||
print("--- Starting TTS Pipeline Tests ---\n")
|
||||
pipeline = TTSPipeline(max_concurrency=4, max_queue_size=100)
|
||||
await pipeline.start()
|
||||
|
||||
# Event tracking for assertions
|
||||
events = []
|
||||
def callback(job_id, status):
|
||||
events.append((job_id, status))
|
||||
|
||||
pipeline.set_callback(callback)
|
||||
|
||||
# --- Test A: Bounded Concurrency & Drain ---
|
||||
print("Test A: Submitting 50 jobs and checking concurrency...")
|
||||
job_ids = []
|
||||
for i in range(50):
|
||||
jid = await pipeline.submit(f"Text {i}", "en-US")
|
||||
job_ids.append(jid)
|
||||
|
||||
await pipeline.drain()
|
||||
|
||||
assert pipeline._peak_concurrency <= 4, f"Concurrency exceeded! Peak: {pipeline._peak_concurrency}"
|
||||
print(f" [PASS] Max concurrency was {pipeline._peak_concurrency}/4")
|
||||
|
||||
# --- Test B: Backpressure ---
|
||||
print("\nTest B: Testing backpressure (100 job cap)...")
|
||||
# Fill the queue to 100. Note: drain() cleared previous jobs, so current count is 0.
|
||||
# We submit 100 to fill it.
|
||||
for i in range(100):
|
||||
await pipeline.submit("Backpressure test", "en-US")
|
||||
|
||||
try:
|
||||
await pipeline.submit("The breaking job", "en-US")
|
||||
print(" [FAIL] Pipeline accepted 101st job without error.")
|
||||
except BufferError:
|
||||
print(" [PASS] Pipeline correctly rejected 101st job with BufferError.")
|
||||
|
||||
# Clear the queue for next test
|
||||
await pipeline.drain()
|
||||
|
||||
# --- Test C: Cancellation & Recovery ---
|
||||
print("\nTest C: Mid-flight cancellation and recovery...")
|
||||
# Submit a batch
|
||||
batch_ids = []
|
||||
for i in range(10):
|
||||
batch_ids.append(await pipeline.submit("Cancel me", "en-US"))
|
||||
|
||||
# Wait a tiny bit to ensure jobs are picked up by workers (in-flight)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Cancel the first job in the batch
|
||||
target_id = batch_ids[0]
|
||||
await pipeline.cancel(target_id)
|
||||
print(f" Cancelled job: {target_id}")
|
||||
|
||||
await pipeline.drain()
|
||||
|
||||
# Verify the cancelled job is marked correctly in events
|
||||
cancelled_events = [e for e in events if e[0] == target_id and e[1] == JobStatus.CANCELLED]
|
||||
assert len(cancelled_events) > 0, "Job was not recorded as CANCELLED"
|
||||
print(" [PASS] Job cancellation detected.")
|
||||
|
||||
# Verify pool is still usable
|
||||
print(" Verifying pool usability...")
|
||||
new_job = await pipeline.submit("Post-cancel job", "en-US")
|
||||
await pipeline.drain()
|
||||
|
||||
# Check if the new job completed
|
||||
new_job_events = [e for e in events if e[0] == new_job and e[1] == JobStatus.COMPLETED]
|
||||
# Note: It might be FAILED if the 10% error hit, but it shouldn't be stuck.
|
||||
assert any(e[1] in (JobStatus.COMPLETED, JobStatus.FAILED) for e in events if e[0] == new_job), "Pool stuck!"
|
||||
print(" [PASS] Pool recovered and processed new jobs.")
|
||||
|
||||
await pipeline.stop()
|
||||
print("\n--- All Tests Passed Successfully ---")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
@@ -0,0 +1,276 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import time
|
||||
import random
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Optional
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
# --- CONFIGURATION ---
|
||||
SHARED_SECRET = "super_secret_twitch_key"
|
||||
DISCORD_WEBHOOK_URL = "https://discord.com/api/webhooks/mock"
|
||||
RATE_LIMIT_CAPACITY = 5
|
||||
RATE_LIMIT_REFILL_PER_SEC = 5
|
||||
IDEMPOTENCY_TTL_SECONDS = 300 # 5 minutes
|
||||
|
||||
# --- MODELS & STATE ---
|
||||
|
||||
@dataclass
|
||||
class DiscordResponse:
|
||||
status: int
|
||||
retry_after: Optional[float] = None
|
||||
|
||||
class TokenBucket:
|
||||
"""Implements the Token Bucket algorithm for rate limiting."""
|
||||
def __init__(self, capacity: float, refill_rate: float):
|
||||
self.capacity = capacity
|
||||
self.refill_rate = refill_rate
|
||||
self.tokens = capacity
|
||||
self.last_update = time.monotonic()
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
async def consume(self):
|
||||
async with self.lock:
|
||||
while self.tokens < 1:
|
||||
# Refill tokens based on time passed
|
||||
now = time.monotonic()
|
||||
passed = now - self.last_update
|
||||
self.tokens = min(self.capacity, self.tokens + (passed * self.refill_rate))
|
||||
self.last_update = now
|
||||
|
||||
if self.tokens < 1:
|
||||
# Wait for enough time to pass to get at least 1 token
|
||||
wait_time = (1 - self.tokens) / self.refill_rate
|
||||
await asyncio.sleep(wait_time)
|
||||
|
||||
self.tokens -= 1
|
||||
|
||||
class IdempotencyStore:
|
||||
"""In-memory store to prevent duplicate processing."""
|
||||
def __init__(self, ttl: int):
|
||||
self.store: Dict[str, float] = {}
|
||||
self.ttl = ttl
|
||||
|
||||
def is_duplicate(self, event_id: str) -> bool:
|
||||
now = time.time()
|
||||
if event_id in self.store:
|
||||
if now - self.store[event_id] < self.ttl:
|
||||
return True
|
||||
self.store[event_id] = now
|
||||
return False
|
||||
|
||||
async def cleanup_loop(self):
|
||||
"""Background task to evict stale entries."""
|
||||
while True:
|
||||
await asyncio.sleep(60)
|
||||
now = time.time()
|
||||
expired = [k for k, v in self.store.items() if now - v > self.ttl]
|
||||
for k in expired:
|
||||
del self.store[k]
|
||||
|
||||
# --- MOCK DISCORD SERVICE ---
|
||||
|
||||
async def discord_send(payload: dict) -> DiscordResponse:
|
||||
"""
|
||||
Mocks a Discord webhook.
|
||||
Succeeds ~95% of the time. Returns 429 occasionally.
|
||||
"""
|
||||
# Simulate network latency
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
rand = random.random()
|
||||
if rand < 0.05: # 5% chance of rate limit
|
||||
return DiscordResponse(status=429, retry_after=1.0)
|
||||
elif rand < 0.07: # 2% chance of other error
|
||||
return DiscordResponse(status=500)
|
||||
|
||||
# Success
|
||||
return DiscordResponse(status=204)
|
||||
|
||||
# --- CORE BRIDGE LOGIC ---
|
||||
|
||||
class WebhookBridge:
|
||||
def __init__(self):
|
||||
self.idempotency = IdempotencyStore(IDEMPOTENCY_TTL_SECONDS)
|
||||
self.limiter = TokenBucket(RATE_LIMIT_CAPACITY, RATE_LIMIT_REFILL_PER_SEC)
|
||||
self.queue = asyncio.Queue()
|
||||
self.forward_count = 0 # For testing assertions
|
||||
|
||||
async def worker(self):
|
||||
"""Background worker that processes the queue and respects rate limits."""
|
||||
while True:
|
||||
payload = await self.queue.get()
|
||||
try:
|
||||
await self._process_forward(payload)
|
||||
except Exception as e:
|
||||
print(f"[Worker Error] {e}")
|
||||
finally:
|
||||
self.queue.task_done()
|
||||
|
||||
async def _process_forward(self, payload: dict):
|
||||
"""Handles the actual POST to Discord with retry logic."""
|
||||
await self.limiter.consume()
|
||||
|
||||
resp = await discord_send(payload)
|
||||
|
||||
if resp.status == 429:
|
||||
print(f"[Rate Limit] Discord returned 429. Retrying after {resp.retry_after}s")
|
||||
await asyncio.sleep(resp.retry_after or 1)
|
||||
# Retry once
|
||||
resp = await discord_send(payload)
|
||||
|
||||
if resp.status in (200, 204):
|
||||
self.forward_count += 1
|
||||
else:
|
||||
print(f"[Error] Discord failed with status {resp.status}")
|
||||
|
||||
async def handle_webhook(self, request: web.Request) -> web.Response:
|
||||
"""HTTP Handler for POST /webhook."""
|
||||
# 1. Verify Signature
|
||||
signature_header = request.headers.get("X-Signature")
|
||||
if not signature_header:
|
||||
return web.Response(text="Missing X-Signature", status=401)
|
||||
|
||||
body = await request.read()
|
||||
expected_sig = hmac.new(
|
||||
SHARED_SECRET.encode(),
|
||||
body,
|
||||
hashlib.sha256
|
||||
).hexdigest()
|
||||
|
||||
if not hmac.compare_digest(expected_sig, signature_header):
|
||||
return web.Response(text="Invalid Signature", status=401)
|
||||
|
||||
# 2. Parse JSON
|
||||
try:
|
||||
data = json.loads(body)
|
||||
event_id = data.get("event_id")
|
||||
if not event_id:
|
||||
return web.Response(text="Missing event_id", status=400)
|
||||
except json.JSONDecodeError:
|
||||
return web.Response(text="Invalid JSON", status=400)
|
||||
|
||||
# 3. Idempotency Check
|
||||
if self.idempotency.is_duplicate(event_id):
|
||||
return web.Response(text="Duplicate event", status=200)
|
||||
|
||||
# 4. Enqueue for forwarding
|
||||
await self.queue.put(data)
|
||||
return web.Response(text="Accepted", status=202)
|
||||
|
||||
# --- TEST SUITE ---
|
||||
|
||||
async def run_tests(bridge: WebhookBridge, server_url: str):
|
||||
print("\n--- Starting Tests ---")
|
||||
async with aiohttp.ClientSession() as session:
|
||||
|
||||
# Helper to sign requests
|
||||
def sign(payload_dict):
|
||||
body = json.dumps(payload_dict).encode()
|
||||
sig = hmac.new(SHARED_SECRET.encode(), body, hashlib.sha256).hexdigest()
|
||||
return body, sig
|
||||
|
||||
# a) Valid Request
|
||||
print("Test A: Valid signature...")
|
||||
payload_a = {"event_id": "evt_1", "type": "chat", "data": {"msg": "hi"}}
|
||||
body, sig = sign(payload_a)
|
||||
async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": sig}) as r:
|
||||
assert r.status == 202
|
||||
await asyncio.sleep(0.5) # Wait for worker
|
||||
assert bridge.forward_count == 1
|
||||
print("✅ Passed")
|
||||
|
||||
# b) Tampered Signature
|
||||
print("Test B: Tampered signature...")
|
||||
payload_b = {"event_id": "evt_2", "type": "chat"}
|
||||
body, _ = sign(payload_b)
|
||||
async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": "wrong_sig"}) as r:
|
||||
assert r.status == 401
|
||||
print("✅ Passed")
|
||||
|
||||
# c) Idempotency (Replay)
|
||||
print("Test C: Replay event_id...")
|
||||
payload_c = {"event_id": "evt_1", "type": "chat"} # Same ID as Test A
|
||||
body, sig = sign(payload_c)
|
||||
async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": sig}) as r:
|
||||
assert r.status == 200 # Returns 200 for duplicates per requirements
|
||||
await asyncio.sleep(0.5)
|
||||
assert bridge.forward_count == 1 # Count should NOT have increased
|
||||
print("✅ Passed")
|
||||
|
||||
# d) Rate Limiting (Burst)
|
||||
print("Test D: Burst > 5 events/sec...")
|
||||
bridge.forward_count = 0 # Reset count for this test
|
||||
start_time = time.monotonic()
|
||||
|
||||
# Send 10 events rapidly
|
||||
tasks = []
|
||||
for i in range(10):
|
||||
p = {"event_id": f"burst_{i}", "type": "chat"}
|
||||
b, s = sign(p)
|
||||
tasks.append(session.post(f"{server_url}/webhook", data=b, headers={"X-Signature": s}))
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Wait for all to be processed by worker
|
||||
await bridge.queue.join()
|
||||
duration = time.monotonic() - start_time
|
||||
|
||||
# 10 events at 5/sec should take at least ~1.8-2.0 seconds
|
||||
# (First 5 instant, next 5 delayed by 1s each)
|
||||
print(f" Burst of 10 took {duration:.2f}s")
|
||||
assert duration >= 1.0
|
||||
print("✅ Passed")
|
||||
|
||||
print("--- All Tests Passed Successfully ---\n")
|
||||
|
||||
# --- MAIN ENTRYPOINT ---
|
||||
|
||||
async def main():
|
||||
bridge = WebhookBridge()
|
||||
app = web.Application()
|
||||
app.router.add_post('/webhook', bridge.handle_webhook)
|
||||
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, 'localhost', 8080)
|
||||
await site.start()
|
||||
|
||||
print("Server started at http://localhost:8080")
|
||||
print("Press Ctrl+C to stop.")
|
||||
|
||||
# Start background tasks
|
||||
worker_task = asyncio.create_task(bridge.worker())
|
||||
cleanup_task = asyncio.create_task(bridge.idempotency.cleanup_loop())
|
||||
|
||||
# Run tests automatically
|
||||
try:
|
||||
await run_tests(bridge, "http://localhost:8080")
|
||||
except AssertionError as e:
|
||||
print(f"❌ Test Failed!")
|
||||
raise e
|
||||
except Exception as e:
|
||||
print(f"❌ Error during tests: {e}")
|
||||
raise e
|
||||
|
||||
# Keep running until interrupted
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(3600)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
finally:
|
||||
worker_task.cancel()
|
||||
cleanup_task.cancel()
|
||||
await runner.cleanup()
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
+19
-7
@@ -109,16 +109,28 @@ def main():
|
||||
if tok_sec is None and comp_tokens and wall:
|
||||
tok_sec = round(comp_tokens / wall, 2)
|
||||
|
||||
# save the raw output
|
||||
# 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
|
||||
# block, (b) prose preamble + fence, (c) bare code + trailing stray fence/prose.
|
||||
out_path = os.path.join(HERE, "outputs", f"{args.name}.{ext}")
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
# strip common markdown fences if present
|
||||
import re as _re
|
||||
text = content.strip()
|
||||
if text.startswith("```"):
|
||||
lines = text.split("\n")
|
||||
if lines[0].startswith("```"): lines = lines[1:]
|
||||
if lines and lines[-1].strip() == "```": lines = lines[:-1]
|
||||
text = "\n".join(lines)
|
||||
pairs = list(_re.finditer(r"```(?:[a-zA-Z0-9_+-]*)?\n(.*?)```", text, _re.S))
|
||||
if pairs:
|
||||
text = pairs[0].group(1).strip()
|
||||
else:
|
||||
out, started = [], False
|
||||
for l in text.split("\n"):
|
||||
if l.strip().startswith("```"):
|
||||
break # stray fence ends the code
|
||||
if not started and _re.match(r"^\s*(#!|import |from |async |def |class |use |pub )", l):
|
||||
started = True
|
||||
if started:
|
||||
if _re.match(r"^\s*(- |\* |\d+\. |> |# )", l) and not _re.match(r"^\s*(import |from |def |class |async |return |if |for |while |try|except|with )", l):
|
||||
break # trailing markdown prose
|
||||
out.append(l)
|
||||
text = "\n".join(out).strip() if out else text
|
||||
open(out_path, "w").write(text)
|
||||
print(f"✓ saved {len(text)} bytes -> {os.path.relpath(out_path, HERE)}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user