import time
import threading
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
class LeakyBucket4LLM:
    def __init__(self, capacity: int, rate: float):
        """
        漏桶算法（适配大模型限流）
        :param capacity: 漏桶最大容量（最大等待推理请求数）
        :param rate: 漏水速率（每秒处理的大模型请求数）
        """
        self.capacity = capacity
        self.rate = rate
        self.water = 0  # 当前等待的请求数
        self.last_time = time.time()
        self.lock = threading.Lock()  # 线程锁，保证高并发下数据安全
    def allow_request(self) -> bool:
        """判断是否允许处理大模型推理请求"""
        with self.lock:
            now = time.time()
            # 计算时间差，生成可漏出的水量
            time_diff = now - self.last_time
            leak_water = time_diff * self.rate
            
            # 更新剩余水量（不能小于0）
            self.water = max(0.0, self.water - leak_water)
            self.last_time = now
            # 桶未满：允许入队，水量+1
            if self.water < self.capacity:
                self.water += 1
                return True
            # 桶已满：拒绝请求
            return False
# ==================== 漏桶算法可视化 ====================
def simulate_leaky_bucket(capacity: int, rate: float, request_times: list):
    """模拟漏桶算法，记录完整工作过程"""
    if not request_times:
        return [], [], [], []
    
    max_time = max(request_times) + 6
    dt = 0.01
    times = np.arange(0, max_time, dt)
    
    water_levels = []
    input_events = []
    output_events = []
    rejected_events = []
    current_water = 0
    request_idx = 0
    
    for t in times:
        while request_idx < len(request_times) and abs(request_times[request_idx] - t) < dt/2:
            if current_water < capacity:
                current_water += 1
                input_events.append((t, 1))
            else:
                rejected_events.append((t, 1))
            request_idx += 1
        
        leak_amount = rate * dt
        leaked = min(current_water, leak_amount)
        current_water -= leaked
        if leaked > 0:
            output_events.append((t, leaked))
        water_levels.append(current_water)
    
    return times, water_levels, input_events, output_events, rejected_events
def plot_bucket_level(capacity: int, rate: float, times, water_levels, input_events, rejected_events, total_requests: int):
    """图1：漏桶堆积曲线"""
    fig, ax = plt.subplots(figsize=(14, 7))
    fig.patch.set_facecolor('#f8f9fa')
    ax.set_facecolor('#ffffff')
    
    # 蓝色填充区域：桶内请求堆积
    ax.fill_between(times, water_levels, alpha=0.4, color='#3498db')
    ax.plot(times, water_levels, 'b-', linewidth=2, label='桶内请求数')
    
    # 红色虚线：容量上限
    ax.axhline(y=capacity, color='#e74c3c', linestyle='--', linewidth=2, label=f'容量上限 ({capacity})')
    # 绿色点线：处理速率
    ax.axhline(y=rate, color='#27ae60', linestyle=':', linewidth=2, label=f'处理速率 ({rate}/s)')
    
    # 红色阴影区域：桶满时期
    overflow_mask = np.array(water_levels) >= capacity
    if np.any(overflow_mask):
        start_idx = np.where(overflow_mask)[0][0]
        end_idx = np.where(overflow_mask)[0][-1]
        ax.axvspan(times[start_idx], times[end_idx], alpha=0.15, color='red')
        ax.text((times[start_idx] + times[end_idx]) / 2, capacity + 0.5,
                 '桶满! 请求被拒绝', ha='center', fontsize=11, color='#e74c3c', fontweight='bold')
    
    # 竖线标记请求到达
    for t, _ in input_events:
        ax.axvline(x=t, color='#3498db', alpha=0.4, linewidth=1)
    # X标记拒绝的请求
    for t, _ in rejected_events:
        ax.scatter(t, capacity + 0.5, color='#e74c3c', s=80, marker='x', zorder=5)
    
    ax.scatter([], [], color='#e74c3c', s=80, marker='x', label=f'被拒绝({len(rejected_events)}个)')
    
    # 添加解释注解框
    annotation_text = (
        "【图例说明】\n"
        "• 蓝色区域：桶内堆积的请求数量\n"
        "• 红色虚线：漏桶容量上限，超出则拒绝\n"
        "• 绿色点线：漏桶处理速率（匀速输出）\n"
        "• 红色阴影：桶满时期，新请求被拒绝"
    )
    ax.text(0.02, 0.98, annotation_text, transform=ax.transAxes, fontsize=9,
            verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    ax.set_xlabel('时间 (秒)', fontsize=12)
    ax.set_ylabel('请求数', fontsize=12)
    ax.set_title(f'漏桶堆积过程 (容量={capacity}, 速率={rate}/s)\n{total_requests}个请求，每0.1秒到达1个', 
                 fontsize=14, fontweight='bold', color='#2c3e50', pad=10)
    ax.legend(loc='upper right')
    ax.set_xlim(0, max(times))
    ax.set_ylim(0, min(capacity * 2, 15))
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('图1-漏桶堆积曲线.png', dpi=300, bbox_inches='tight', facecolor='#f8f9fa')
    plt.show()
    print("✅ 图1已保存: 图1-漏桶堆积曲线.png")
def plot_flow_comparison(capacity: int, rate: float, times, input_events, output_events, rejected_events, total_requests: int):
    """图2：流量整形对比"""
    fig, ax = plt.subplots(figsize=(14, 7))
    fig.patch.set_facecolor('#f8f9fa')
    ax.set_facecolor('#ffffff')
    
    input_rate = np.zeros_like(times)
    for t, _ in input_events:
        mask = (times >= t - 0.5) & (times < t + 0.5)
        input_rate[mask] += 1
    
    output_rate = np.zeros_like(times)
    for t, amount in output_events:
        mask = (times >= t - 0.2) & (times < t + 0.2)
        output_rate[mask] += amount * 5
    
    # 红色填充：输入速率
    ax.fill_between(times, input_rate, alpha=0.3, color='#e74c3c')
    ax.plot(times, input_rate, 'r-', linewidth=2, label='输入速率（突发）')
    # 绿色线：输出速率
    ax.plot(times, output_rate, 'g-', linewidth=2.5, label='输出速率（恒定）')
    
    # 绿色虚线：固定输出速率
    ax.axhline(y=rate, color='#27ae60', linestyle='--', linewidth=2, alpha=0.7)
    ax.text(max(times) * 0.85, rate + 0.3, f'固定输出={rate}/s', fontsize=10, color='#27ae60')
    
    # X标记拒绝的请求
    for t, _ in rejected_events:
        ax.axvline(x=t, color='#e74c3c', alpha=0.5, linewidth=2, linestyle=':')
        ax.scatter(t, rate + 1, color='#e74c3c', s=100, marker='x', zorder=5)
    
    if rejected_events:
        ax.scatter([], [], color='#e74c3c', s=100, marker='x', label=f'被拒绝({len(rejected_events)}个)')
    
    # 添加解释注解框
    annotation_text = (
        "【图例说明】\n"
        "• 红色曲线：请求到达速率（突发）\n"
        "• 绿色曲线：漏桶输出速率（恒定）\n"
        "• 漏桶将突发流量整形为均匀输出\n"
        "• 桶满时请求被拒绝（X标记）"
    )
    ax.text(0.02, 0.98, annotation_text, transform=ax.transAxes, fontsize=9,
            verticalalignment='top', bbox=dict(boxstyle='round', facecolor='honeydew', alpha=0.8))
    
    ax.set_xlabel('时间 (秒)', fontsize=12)
    ax.set_ylabel('请求速率 (req/s)', fontsize=12)
    ax.set_title(f'流量整形效果对比 (容量={capacity}, 速率={rate}/s)\n{total_requests}个请求：突发流量被整形为恒定输出', 
                 fontsize=14, fontweight='bold', color='#2c3e50', pad=10)
    ax.legend(loc='upper right')
    ax.set_xlim(0, max(times))
    ax.set_ylim(0, max(max(input_rate) if len(input_rate) > 0 else 1, rate + 3))
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('图2-流量整形对比.png', dpi=300, bbox_inches='tight', facecolor='#f8f9fa')
    plt.show()
    print("✅ 图2已保存: 图2-流量整形对比.png")
def visualize_leaky_bucket():
    """可视化漏桶算法"""
    # 使用示例中的参数: 容量=5, 速率=2
    capacity = 5
    rate = 2
    total_requests = 20  # 20个请求
    
    # 模拟20个请求，每0.1秒到达1个
    request_times = [i * 0.1 for i in range(total_requests)]
    
    times, water_levels, input_events, output_events, rejected_events = simulate_leaky_bucket(
        capacity, rate, request_times
    )
    
    # 绘制两个独立的图
    plot_bucket_level(capacity, rate, times, water_levels, input_events, rejected_events, total_requests)
    plot_flow_comparison(capacity, rate, times, input_events, output_events, rejected_events, total_requests)
# ==================== 大模型场景测试 ====================
# 配置：漏桶容量=5（最多等待5个请求），每秒处理2个大模型推理请求
llm_leaky_bucket = LeakyBucket4LLM(capacity=5, rate=2)
# 记录入队和被处理的请求
request_log = []  # [(req_id, status, relative_time), ...]
log_lock = threading.Lock()
def llm_inference_request(req_id: int, start_time: float):
    """模拟大模型推理请求"""
    if llm_leaky_bucket.allow_request():
        with log_lock:
            request_log.append((req_id, '入队', time.time() - start_time))
    else:
        with log_lock:
            request_log.append((req_id, '拒绝', time.time() - start_time))
if __name__ == "__main__":
    # 先显示可视化图
    visualize_leaky_bucket()
    
    # 测试限流功能
    total_requests = 20
    print("\n" + "="*50)
    print(f"开始模拟{total_requests}个并发大模型请求...")
    print("漏桶配置: 容量=5, 速率=2/s (每0.5秒处理1个请求)")
    print("="*50)
    
    start_time = time.time()
    
    # 发送20个请求
    for i in range(total_requests):
        threading.Thread(target=llm_inference_request, args=(i, start_time)).start()
        time.sleep(0.1)
    
    # 等待所有请求处理完成
    time.sleep(3)
    
    # 按时间顺序打印日志
    request_log.sort(key=lambda x: x[2])
    for req_id, status, rel_time in request_log:
        print(f"[{rel_time:.2f}s] 推理请求{req_id}：{'已入队，等待匀速处理' if status=='入队' else '限流拒绝，桶已满'}")
    
    # 统计结果
    queued = sum(1 for r in request_log if r[1] == '入队')
    rejected = sum(1 for r in request_log if r[1] == '拒绝')
    
    print("\n" + "="*50)
    print(f"请求统计: 共{total_requests}个请求")
    print(f"  - 入队: {queued}个")
    print(f"  - 拒绝: {rejected}个")
    print("="*50)

import time
import threading
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
class TokenBucket:
    """令牌桶算法"""
    def __init__(self, capacity: int, rate: float):
        """
        令牌桶算法实现
        :param capacity: 桶的最大容量（最大令牌数）
        :param rate: 令牌生成速率（每秒生成的令牌数）
        """
        self.capacity = capacity
        self.rate = rate
        self.tokens = float(capacity)  # 初始令牌数（满桶）
        self.last_time = time.time()
        self.lock = threading.Lock()
    
    def consume_token(self) -> bool:
        """
        尝试消耗一个令牌
        :return: 是否成功获取令牌
        """
        with self.lock:
            now = time.time()
            # 计算新增的令牌数
            time_diff = now - self.last_time
            new_tokens = time_diff * self.rate
            
            # 更新令牌数（不超过容量）
            self.tokens = min(self.capacity, self.tokens + new_tokens)
            self.last_time = now
            
            # 尝试消耗令牌
            if self.tokens >= 1:
                self.tokens -= 1
                return True
            return False
    
    def get_token_count(self) -> float:
        """获取当前令牌数"""
        with self.lock:
            return self.tokens
# ==================== 令牌桶可视化 ====================
def simulate_token_bucket(capacity: int, rate: float, request_times: list):
    """模拟令牌桶算法，记录完整工作过程"""
    if not request_times:
        return [], [], [], []
    
    max_time = max(request_times) + 6
    dt = 0.01
    times = np.arange(0, max_time, dt)
    
    token_levels = []
    consume_events = []
    generate_events = []
    rejected_events = []
    current_tokens = float(capacity)
    next_generate_time = 0
    
    for t in times:
        # 生成令牌
        if abs(t - next_generate_time) < dt/2 and current_tokens < capacity:
            current_tokens = min(capacity, current_tokens + 1)
            generate_events.append((t, 1))
            next_generate_time += 1 / rate
        
        # 尝试消耗令牌
        for req_t in request_times:
            if abs(t - req_t) < dt/2:
                if current_tokens >= 1:
                    current_tokens -= 1
                    consume_events.append((t, 1))
                else:
                    rejected_events.append((t, 1))
                break
        
        token_levels.append(current_tokens)
    
    return times, token_levels, generate_events, consume_events, rejected_events
def plot_bucket_tokens(capacity: int, rate: float, times, token_levels, generate_events, consume_events, rejected_events, total_requests: int):
    """图1：令牌桶令牌变化"""
    fig, ax = plt.subplots(figsize=(14, 7))
    fig.patch.set_facecolor('#f8f9fa')
    ax.set_facecolor('#ffffff')
    
    # 填充区域表示令牌数
    ax.fill_between(times, token_levels, alpha=0.3, color='#f39c12')
    ax.plot(times, token_levels, color='#e67e22', linewidth=2, label='可用令牌数')
    
    # 容量上限
    ax.axhline(y=capacity, color='#e74c3c', linestyle='--', linewidth=2, label=f'容量上限 ({capacity})')
    # 生成速率线
    ax.axhline(y=rate, color='#27ae60', linestyle=':', linewidth=2, label=f'生成速率 ({rate}/s)')
    
    # 令牌生成标记
    for t, _ in generate_events:
        ax.axvline(x=t, color='#27ae60', alpha=0.3, linewidth=1)
        idx = int(t / 0.01)
        if idx < len(token_levels):
            ax.scatter(t, token_levels[idx] + 0.3, color='#27ae60', s=50, marker='^', zorder=5)
    
    # 令牌消耗标记
    for t, _ in consume_events:
        ax.axvline(x=t, color='#3498db', alpha=0.3, linewidth=1)
        idx = int(t / 0.01)
        if idx < len(token_levels):
            ax.scatter(t, token_levels[idx] - 0.3, color='#3498db', s=50, marker='v', zorder=5)
    
    # 拒绝标记
    for t, _ in rejected_events:
        ax.scatter(t, -0.5, color='#e74c3c', s=100, marker='x', zorder=5)
    
    ax.scatter([], [], color='#27ae60', s=50, marker='^', label=f'生成令牌({len(generate_events)}次)')
    ax.scatter([], [], color='#3498db', s=50, marker='v', label=f'消耗令牌({len(consume_events)}次)')
    ax.scatter([], [], color='#e74c3c', s=100, marker='x', label=f'拒绝请求({len(rejected_events)}次)')
    
    # 注解框
    annotation_text = (
        "【图例说明】\n"
        "  橙色曲线：桶内剩余令牌数\n"
        "  红色虚线：容量上限，超出则丢弃\n"
        "  绿色点线：令牌生成速率\n"
        "  ^ 标记：令牌生成\n"
        "  v 标记：令牌消耗\n"
        "  x 标记：无令牌时被拒绝"
    )
    ax.text(0.02, 0.98, annotation_text, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.9))
    
    ax.set_xlabel('时间 (秒)', fontsize=12)
    ax.set_ylabel('令牌数', fontsize=12)
    ax.set_title(f'令牌桶令牌变化 (容量={capacity}, 速率={rate}/s)\n{total_requests}个请求，每0.2秒到达1个', 
                 fontsize=14, fontweight='bold', color='#2c3e50', pad=10)
    ax.legend(loc='upper right')
    ax.set_xlim(0, max(times))
    ax.set_ylim(-1, min(capacity * 2, 15))
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('图1-令牌桶变化.png', dpi=300, bbox_inches='tight', facecolor='#f8f9fa')
    plt.show()
    print("[OK] 图1已保存: 图1-令牌桶变化.png")
def plot_request_comparison(capacity: int, rate: float, times, generate_events, consume_events, rejected_events, total_requests: int):
    """图2：请求处理对比"""
    fig, ax = plt.subplots(figsize=(14, 7))
    fig.patch.set_facecolor('#f8f9fa')
    ax.set_facecolor('#ffffff')
    
    # 生成令牌速率
    generate_rate = np.zeros_like(times)
    for t, _ in generate_events:
        mask = (times >= t - 0.5) & (times < t + 0.5)
        generate_rate[mask] += 1
    
    # 消耗令牌速率
    consume_rate = np.zeros_like(times)
    for t, _ in consume_events:
        mask = (times >= t - 0.5) & (times < t + 0.5)
        consume_rate[mask] += 1
    
    # 固定生成速率线
    ax.axhline(y=rate, color='#27ae60', linestyle='--', linewidth=2, alpha=0.7, label=f'固定生成速率={rate}/s')
    
    # 填充：生成
    ax.fill_between(times, generate_rate, alpha=0.3, color='#27ae60')
    ax.plot(times, generate_rate, color='#27ae60', linewidth=2, label='令牌生成速率')
    
    # 填充：消耗
    ax.fill_between(times, consume_rate, alpha=0.3, color='#3498db')
    ax.plot(times, consume_rate, color='#3498db', linewidth=2, label='令牌消耗速率')
    
    # 拒绝标记
    for t, _ in rejected_events:
        ax.axvline(x=t, color='#e74c3c', alpha=0.5, linewidth=2, linestyle=':')
        ax.scatter(t, rate + 1, color='#e74c3c', s=100, marker='x', zorder=5)
    
    if rejected_events:
        ax.scatter([], [], color='#e74c3c', s=100, marker='x', label=f'被拒绝({len(rejected_events)}次)')
    
    # 注解框
    annotation_text = (
        "【令牌桶特点】\n"
        "  令牌按固定速率持续生成\n"
        "  允许突发：令牌积累后快速消费\n"
        "  无令牌时请求被拒绝"
    )
    ax.text(0.02, 0.98, annotation_text, transform=ax.transAxes, fontsize=10,
            verticalalignment='top', bbox=dict(boxstyle='round', facecolor='honeydew', alpha=0.9))
    
    ax.set_xlabel('时间 (秒)', fontsize=12)
    ax.set_ylabel('速率 (个/秒)', fontsize=12)
    ax.set_title(f'请求处理对比 (容量={capacity}, 速率={rate}/s)\n{total_requests}个请求：允许突发流量，先积累后输出', 
                 fontsize=14, fontweight='bold', color='#2c3e50', pad=10)
    ax.legend(loc='upper right')
    ax.set_xlim(0, max(times))
    ax.set_ylim(0, max(max(generate_rate) if len(generate_rate) > 0 else 1, 
                       max(consume_rate) if len(consume_rate) > 0 else 1, rate + 2))
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('图2-请求处理对比.png', dpi=300, bbox_inches='tight', facecolor='#f8f9fa')
    plt.show()
    print("[OK] 图2已保存: 图2-请求处理对比.png")
def visualize_token_bucket():
    """令牌桶可视化"""
    capacity = 8
    rate = 2
    total_requests = 15
    
    # 模拟15个请求，每0.2秒到达1个（突发）
    request_times = [i * 0.2 for i in range(total_requests)]
    
    times, token_levels, generate_events, consume_events, rejected_events = simulate_token_bucket(
        capacity, rate, request_times
    )
    
    # 绘制图表
    plot_bucket_tokens(capacity, rate, times, token_levels, generate_events, consume_events, rejected_events, total_requests)
    plot_request_comparison(capacity, rate, times, generate_events, consume_events, rejected_events, total_requests)
# ==================== 大模型场景测试 ====================
# 配置：令牌桶容量=8（最多8个令牌），每秒生成2个令牌
llm_token_bucket = TokenBucket(capacity=8, rate=2)
# 记录日志
request_log = []
log_lock = threading.Lock()
def llm_inference_request(req_id: int, start_time: float):
    """模拟大模型推理请求"""
    if llm_token_bucket.consume_token():
        with log_lock:
            request_log.append((req_id, '通过', time.time() - start_time))
        print(f"[令牌桶] 推理请求{req_id}：获取令牌成功，处理中...")
        time.sleep(0.3)  # 模拟推理耗时
    else:
        with log_lock:
            request_log.append((req_id, '拒绝', time.time() - start_time))
        print(f"[令牌桶] 推理请求{req_id}：无令牌可用，请求被拒绝")
if __name__ == "__main__":
    visualize_token_bucket()
    
    # 测试限流功能
    total_requests = 15
    print("\n" + "=" * 60)
    print(f"开始模拟{total_requests}个并发大模型请求...")
    print("令牌桶配置: 容量=8, 生成速率=2/s (每0.5秒生成1个令牌)")
    print("=" * 60)
    
    start_time = time.time()
    
    # 发送15个请求（突发）
    for i in range(total_requests):
        threading.Thread(target=llm_inference_request, args=(i, start_time)).start()
        time.sleep(0.1)  # 每0.1秒发送一个
    
    # 等待所有请求处理完成
    time.sleep(5)
    
    # 按时间顺序打印日志
    request_log.sort(key=lambda x: x[2])
    print("\n" + "-" * 50)
    print("请求处理日志:")
    print("-" * 50)
    for req_id, status, rel_time in request_log:
        status_text = "通过" if status == '通过' else "拒绝"
        print(f"[{rel_time:.2f}s] 推理请求{req_id}：{status_text}")
    
    # 统计结果
    passed = sum(1 for r in request_log if r[1] == '通过')
    rejected = sum(1 for r in request_log if r[1] == '拒绝')
    
    print("\n" + "=" * 60)
    print(f"请求统计: 共{total_requests}个请求")
    print(f"  - 通过: {passed}个")
    print(f"  - 拒绝: {rejected}个")
    print("=" * 60)