import time
import threading
from collections import defaultdict

class FixedWindowRateLimiter:
    """
    固定窗口限流器
    缺陷：窗口邊界處可能出現(xiàn) 2 倍流量突刺
    """
    
    def __init__(self, max_requests: int, window_seconds: int):
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self._counts = defaultdict(int)
        self._window_starts = defaultdict(float)
        self._lock = threading.Lock()
    
    def is_allowed(self, key: str = "default") -> bool:
        now = time.time()
        
        with self._lock:
            window_start = self._window_starts[key]
            
            # 當(dāng)前時(shí)間超過(guò)窗口，重置計(jì)數(shù)
            if now - window_start >= self.window_seconds:
                self._window_starts[key] = now
                self._counts[key] = 0
            
            if self._counts[key] < self.max_requests:
                self._counts[key] += 1
                return True
            
            return False


# 演示固定窗口的邊界問(wèn)題
limiter = FixedWindowRateLimiter(max_requests=10, window_seconds=60)

# 假設(shè)窗口為 [0s, 60s)，在第 59s 發(fā)送 10 個(gè)請(qǐng)求（全部通過(guò)）
# 然后在第 61s（新窗口開(kāi)始）再發(fā)送 10 個(gè)請(qǐng)求（全部通過(guò)）
# 實(shí)際上在 [59s, 61s] 這 2 秒內(nèi)通過(guò)了 20 個(gè)請(qǐng)求——是上限的 2 倍！
print("固定窗口在邊界處存在 2 倍流量突刺風(fēng)險(xiǎn)")

import time
import threading
from typing import Optional

class TokenBucketRateLimiter:
    """
    令牌桶限流器
    
    特性：
    - 允許短時(shí)突發(fā)（桶滿時(shí)可消費(fèi)全部令牌）
    - 長(zhǎng)期速率受 rate 控制
    - 線程安全
    """
    
    def __init__(self, rate: float, capacity: int):
        """
        :param rate: 令牌補(bǔ)充速率（個(gè)/秒）
        :param capacity: 桶的最大容量（突發(fā)上限）
        """
        self.rate = rate
        self.capacity = capacity
        self._tokens = float(capacity)   # 初始滿桶
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()
    
    def _refill(self):
        """根據(jù)時(shí)間流逝補(bǔ)充令牌（惰性計(jì)算，不需要后臺(tái)線程）"""
        now = time.monotonic()
        elapsed = now - self._last_refill
        new_tokens = elapsed * self.rate
        self._tokens = min(self.capacity, self._tokens + new_tokens)
        self._last_refill = now
    
    def acquire(self, tokens: int = 1) -> bool:
        """
        嘗試獲取令牌
        :param tokens: 需要的令牌數(shù)（支持批量消費(fèi)）
        :return: True=放行, False=拒絕
        """
        with self._lock:
            self._refill()
            
            if self._tokens >= tokens:
                self._tokens -= tokens
                return True
            return False
    
    def acquire_with_wait(self, tokens: int = 1, 
                          timeout: float = None) -> bool:
        """
        等待直到獲取到令牌（阻塞版本）
        :param timeout: 最長(zhǎng)等待時(shí)間（秒），None 表示永久等待
        """
        start = time.monotonic()
        
        while True:
            with self._lock:
                self._refill()
                if self._tokens >= tokens:
                    self._tokens -= tokens
                    return True
            
            # 計(jì)算需要等待多久才能獲得足夠令牌
            with self._lock:
                deficit = tokens - self._tokens
                wait_time = deficit / self.rate
            
            if timeout is not None:
                elapsed = time.monotonic() - start
                if elapsed + wait_time > timeout:
                    return False
            
            time.sleep(min(wait_time, 0.01))  # 最多睡 10ms 避免過(guò)度等待
    
    @property
    def current_tokens(self) -> float:
        """查看當(dāng)前令牌數(shù)（調(diào)試用）"""
        with self._lock:
            self._refill()
            return self._tokens


# ── 使用示例 ──────────────────────────────────────────────
def demo_token_bucket():
    # 每秒補(bǔ)充 5 個(gè)令牌，桶容量 10
    limiter = TokenBucketRateLimiter(rate=5, capacity=10)
    
    print("=== 令牌桶演示 ===")
    print(f"初始令牌數(shù)：{limiter.current_tokens:.1f}")
    
    # 模擬突發(fā)請(qǐng)求：一次性消費(fèi) 8 個(gè)令牌
    results = []
    for i in range(12):
        allowed = limiter.acquire()
        results.append("?" if allowed else "?")
        print(f"請(qǐng)求 {i+1:2d}：{'放行' if allowed else '拒絕'} "
              f"（剩余令牌：{limiter.current_tokens:.2f}）")
    
    # 等待 2 秒后令牌恢復(fù)
    print("\n等待 2 秒，令牌恢復(fù)中...")
    time.sleep(2)
    print(f"2 秒后令牌數(shù)：{limiter.current_tokens:.1f}（補(bǔ)充了約 10 個(gè)）")
    print(f"再次請(qǐng)求：{'放行' if limiter.acquire() else '拒絕'}")

demo_token_bucket()

import time
import threading
from collections import deque
from dataclasses import dataclass
from typing import Any, Optional

@dataclass
class Request:
    """封裝請(qǐng)求及其元數(shù)據(jù)"""
    data: Any
    arrive_time: float
    key: str = "default"

class LeakyBucketRateLimiter:
    """
    漏桶限流器
    
    特性：
    - 嚴(yán)格恒定的輸出速率
    - 超出桶容量的請(qǐng)求直接丟棄
    - 天然實(shí)現(xiàn)請(qǐng)求整形（Traffic Shaping）
    """
    
    def __init__(self, rate: float, capacity: int):
        """
        :param rate: 漏出速率（請(qǐng)求數(shù)/秒）
        :param capacity: 桶容量（最大排隊(duì)數(shù)）
        """
        self.rate = rate
        self.capacity = capacity
        self._interval = 1.0 / rate  # 兩次漏出之間的間隔
        
        self._queue: deque = deque()
        self._lock = threading.Lock()
        self._last_leak_time = time.monotonic()
        
        # 啟動(dòng)后臺(tái)漏出線程
        self._running = True
        self._worker = threading.Thread(target=self._leak_worker, daemon=True)
        self._worker.start()
    
    def _leak_worker(self):
        """后臺(tái)工作線程：以固定速率處理請(qǐng)求"""
        while self._running:
            time.sleep(self._interval)
            
            with self._lock:
                if self._queue:
                    request = self._queue.popleft()
                    wait_time = time.monotonic() - request.arrive_time
                    print(f"  ?? 漏出處理：key={request.key}, "
                          f"等待={wait_time*1000:.1f}ms")
    
    def submit(self, data: Any, key: str = "default") -> bool:
        """
        提交請(qǐng)求到漏桶
        :return: True=入隊(duì)成功, False=桶已滿（丟棄）
        """
        with self._lock:
            if len(self._queue) >= self.capacity:
                print(f"  ?? 桶已滿，丟棄請(qǐng)求：key={key}")
                return False
            
            request = Request(data=data, arrive_time=time.monotonic(), key=key)
            self._queue.append(request)
            print(f"  ?? 請(qǐng)求入隊(duì)：key={key}, 隊(duì)列深度={len(self._queue)}")
            return True
    
    def stop(self):
        self._running = False


# ── 使用示例 ──────────────────────────────────────────────
def demo_leaky_bucket():
    print("\n=== 漏桶演示（2 req/s，桶容量 5）===")
    limiter = LeakyBucketRateLimiter(rate=2, capacity=5)
    
    # 模擬突發(fā)：瞬間提交 8 個(gè)請(qǐng)求
    for i in range(8):
        limiter.submit(f"request-{i}", key=f"user_{i % 3}")
    
    # 等待漏桶處理完畢
    time.sleep(4)
    limiter.stop()
    print("漏桶處理完成，輸出速率嚴(yán)格為 2 req/s")

demo_leaky_bucket()

import time
import threading
from collections import deque
from typing import Dict

class SlidingWindowLogLimiter:
    """
    滑動(dòng)窗口日志限流器（精確版）
    記錄每個(gè)請(qǐng)求的時(shí)間戳，精確統(tǒng)計(jì)窗口內(nèi)的請(qǐng)求數(shù)
    
    適合：低流量、高精度要求的場(chǎng)景
    缺陷：內(nèi)存占用隨請(qǐng)求量線性增長(zhǎng)
    """
    
    def __init__(self, max_requests: int, window_seconds: float):
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        # key → deque of timestamps
        self._logs: Dict[str, deque] = {}
        self._lock = threading.Lock()
    
    def is_allowed(self, key: str = "default") -> bool:
        now = time.monotonic()
        window_start = now - self.window_seconds
        
        with self._lock:
            if key not in self._logs:
                self._logs[key] = deque()
            
            log = self._logs[key]
            
            # 清除窗口外的舊記錄
            while log and log[0] <= window_start:
                log.popleft()
            
            # 判斷窗口內(nèi)請(qǐng)求數(shù)
            if len(log) < self.max_requests:
                log.append(now)
                return True
            
            return False
    
    def current_count(self, key: str = "default") -> int:
        """當(dāng)前窗口內(nèi)的請(qǐng)求數(shù)"""
        now = time.monotonic()
        window_start = now - self.window_seconds
        
        with self._lock:
            if key not in self._logs:
                return 0
            return sum(1 for t in self._logs[key] if t > window_start)


class SlidingWindowCounterLimiter:
    """
    滑動(dòng)窗口計(jì)數(shù)器限流器（高效版）
    將窗口劃分為多個(gè)小格子，用計(jì)數(shù)器近似統(tǒng)計(jì)
    
    適合：高流量、對(duì)精度要求不苛刻的場(chǎng)景
    Redis 分布式限流的主流方案
    """
    
    def __init__(self, max_requests: int, window_seconds: float, 
                 precision: int = 10):
        """
        :param precision: 將窗口劃分為多少個(gè)子格子（越多越精確，越耗內(nèi)存）
        """
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self.precision = precision
        self.slot_duration = window_seconds / precision
        
        # 每個(gè) key 的滑動(dòng)計(jì)數(shù)器：{slot_index: count}
        self._counters: Dict[str, Dict[int, int]] = {}
        self._lock = threading.Lock()
    
    def _current_slot(self) -> int:
        """當(dāng)前時(shí)間對(duì)應(yīng)的格子索引"""
        return int(time.monotonic() / self.slot_duration)
    
    def is_allowed(self, key: str = "default") -> bool:
        current_slot = self._current_slot()
        # 有效的格子范圍
        valid_slots = set(range(current_slot - self.precision + 1, 
                                current_slot + 1))
        
        with self._lock:
            if key not in self._counters:
                self._counters[key] = {}
            
            counter = self._counters[key]
            
            # 清除過(guò)期格子
            expired = [s for s in counter if s not in valid_slots]
            for s in expired:
                del counter[s]
            
            # 統(tǒng)計(jì)窗口內(nèi)總請(qǐng)求數(shù)
            total = sum(counter.values())
            
            if total < self.max_requests:
                counter[current_slot] = counter.get(current_slot, 0) + 1
                return True
            
            return False


# ── 對(duì)比演示 ──────────────────────────────────────────────
def demo_sliding_window():
    print("\n=== 滑動(dòng)窗口演示（10 req/10s）===")
    
    log_limiter = SlidingWindowLogLimiter(max_requests=10, window_seconds=10)
    counter_limiter = SlidingWindowCounterLimiter(max_requests=10, 
                                                   window_seconds=10, 
                                                   precision=10)
    
    # 模擬在 5 秒內(nèi)發(fā)送 15 個(gè)請(qǐng)求
    results_log = []
    results_counter = []
    
    for i in range(15):
        results_log.append(log_limiter.is_allowed("user_1"))
        results_counter.append(counter_limiter.is_allowed("user_1"))
        time.sleep(0.2)
    
    print(f"滑動(dòng)日志結(jié)果：{['?' if r else '?' for r in results_log]}")
    print(f"計(jì)數(shù)器結(jié)果：  {['?' if r else '?' for r in results_counter]}")
    print(f"日志版窗口內(nèi)請(qǐng)求數(shù)：{log_limiter.current_count('user_1')}")

demo_sliding_window()

from abc import ABC, abstractmethod
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
import time

class RateLimiter(ABC):
    """限流器抽象基類"""
    
    @abstractmethod
    def is_allowed(self, key: str) -> bool:
        pass
    
    @abstractmethod
    def get_info(self, key: str) -> dict:
        """返回限流狀態(tài)信息（用于響應(yīng)頭）"""
        pass


class RateLimitMiddleware:
    """
    FastAPI 限流中間件
    支持按 IP、用戶 ID、API Key 等維度限流
    """
    
    def __init__(self, app, limiter: RateLimiter, 
                 key_func=None, exclude_paths: list = None):
        self.app = app
        self.limiter = limiter
        self.key_func = key_func or self._default_key_func
        self.exclude_paths = exclude_paths or ["/health", "/docs"]
    
    @staticmethod
    def _default_key_func(request: Request) -> str:
        """默認(rèn)按客戶端 IP 限流"""
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        return request.client.host
    
    async def __call__(self, scope, receive, send):
        if scope["type"] == "http":
            request = Request(scope, receive)
            
            # 排除不限流的路徑
            if request.url.path not in self.exclude_paths:
                key = self.key_func(request)
                
                if not self.limiter.is_allowed(key):
                    response = JSONResponse(
                        status_code=429,
                        content={
                            "error": "Too Many Requests",
                            "message": "請(qǐng)求過(guò)于頻繁，請(qǐng)稍后再試",
                            "retry_after": 1
                        },
                        headers={
                            "Retry-After": "1",
                            "X-RateLimit-Limit": "100",
                        }
                    )
                    await response(scope, receive, send)
                    return
        
        await self.app(scope, receive, send)


# 構(gòu)建應(yīng)用
app = FastAPI(title="限流演示 API")

# 選擇限流算法（可熱切換）
# 方案1：令牌桶（允許突發(fā)）
rate_limiter = TokenBucketRateLimiter(rate=10, capacity=20)
# 方案2：滑動(dòng)窗口（精確統(tǒng)計(jì)）
# rate_limiter = SlidingWindowLogLimiter(max_requests=100, window_seconds=60)

app.add_middleware(RateLimitMiddleware, limiter=rate_limiter)


@app.get("/api/data")
async def get_data(request: Request):
    client_ip = request.client.host
    return {
        "message": "請(qǐng)求成功",
        "client": client_ip,
        "timestamp": time.time()
    }

@app.get("/health")
async def health():
    return {"status": "ok"}

import redis
import time

class RedisTokenBucketLimiter:
    """
    基于 Redis 的分布式令牌桶限流器
    使用 Lua 腳本保證原子性
    """
    
    # Lua 腳本：原子化地補(bǔ)充令牌并判斷是否放行
    LUA_SCRIPT = """
    local key = KEYS[1]
    local rate = tonumber(ARGV[1])
    local capacity = tonumber(ARGV[2])
    local now = tonumber(ARGV[3])
    local requested = tonumber(ARGV[4])
    
    local last_tokens = tonumber(redis.call('hget', key, 'tokens') or capacity)
    local last_refill = tonumber(redis.call('hget', key, 'last_refill') or now)
    
    -- 補(bǔ)充令牌
    local elapsed = now - last_refill
    local new_tokens = math.min(capacity, last_tokens + elapsed * rate)
    
    local allowed = 0
    if new_tokens >= requested then
        new_tokens = new_tokens - requested
        allowed = 1
    end
    
    redis.call('hset', key, 'tokens', new_tokens)
    redis.call('hset', key, 'last_refill', now)
    redis.call('expire', key, math.ceil(capacity / rate) * 2)
    
    return {allowed, math.floor(new_tokens)}
    """
    
    def __init__(self, redis_url: str, rate: float, capacity: int):
        self.redis = redis.from_url(redis_url)
        self.rate = rate
        self.capacity = capacity
        self._script = self.redis.register_script(self.LUA_SCRIPT)
    
    def is_allowed(self, key: str, tokens: int = 1) -> tuple[bool, int]:
        """
        :return: (是否允許, 剩余令牌數(shù))
        """
        now = time.time()
        result = self._script(
            keys=[f"ratelimit:{key}"],
            args=[self.rate, self.capacity, now, tokens]
        )
        allowed = bool(result[0])
        remaining = int(result[1])
        return allowed, remaining


# 使用示例（需要 Redis 服務(wù)）
# redis_limiter = RedisTokenBucketLimiter(
#     redis_url="redis://localhost:6379",
#     rate=10,
#     capacity=50
# )
# allowed, remaining = redis_limiter.is_allowed("user:12345")
# print(f"放行：{allowed}，剩余令牌：{remaining}")