Tower Load Shed

立即拒绝请求如果内部服务尚未准备就绪。这也被称为负载卸载(load-shedding)。

许可证

本项目使用MIT许可证。

贡献

除非您明确声明，否则任何有意提交给Tower的贡献都应被视为MIT许可，无需任何附加条款或条件。

安装

在项目目录中运行以下Cargo命令：

cargo add tower-load-shed

或者在Cargo.toml中添加以下行：

tower-load-shed = "0.3.0"

完整示例代码

use tower::ServiceBuilder;
use tower_load_shed::LoadShedLayer;
use std::time::Duration;
use tokio::time::sleep;

#[tokio::main]
async fn main() {
    // 创建一个服务构建器并添加负载卸载层
    let svc = ServiceBuilder::new()
        // 添加负载卸载中间件
        .layer(LoadShedLayer::new())
        // 模拟一个可能会过载的服务
        .service_fn(|req: String| async move {
            // 模拟处理时间
            sleep(Duration::from_millis(100)).await;
            Ok::<_, ()>(format!("Processed: {}", req))
        });

    // 模拟正常请求
    match svc.ready().await {
        Ok(mut ready_svc) => {
            if let Ok(resp) = ready_svc.call("Hello".to_string()).await {
                println!("Response: {}", resp);
            }
        }
        Err(_) => println!("Service overloaded, request rejected"),
    }

    // 模拟服务过载情况
    // 在实际应用中，过载可能是由于并发请求过多或资源不足造成的
    for _ in 0..100 {
        if let Err(_) = svc.ready().await {
            println!("Service is overloaded, rejecting request");
            break;
        }
    }
}

工作原理

1. LoadShedLayer会检查内部服务是否准备好处理请求
2. 如果服务未准备好(如并发请求过多)，立即返回错误
3. 防止服务因过载而崩溃，保证系统稳定性
4. 适用于高并发场景下的流量控制

使用场景

• 微服务架构中的API网关
• 高并发Web服务
• 需要保护后端服务的代理层
• 任何需要防止过载的Rust网络服务

注意事项

• 负载卸载是一种"快速失败"策略，不是排队机制
• 适用于短期过载情况，长期过载需要扩容解决方案
• 可以与其他Tower中间件(如限流、超时)组合使用

扩展示例代码

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tower::{Service, ServiceBuilder, ServiceExt};
use tower_load_shed::LoadShedLayer;
use tokio::sync::Semaphore;

// 自定义服务结构体
struct MyService {
    counter: Arc<AtomicUsize>,
    semaphore: Arc<Semaphore>,
}

impl Service<String> for MyService {
    type Response = String;
    type Error = String;
    type Future = tokio::task::JoinHandle<Result<String, String>>;

    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
        // 检查信号量是否可用
        match self.semaphore.try_acquire() {
            Ok(_) => std::task::Poll::Ready(Ok(())),
            Err(_) => std::task::Poll::Ready(Err("Service busy".to_string())),
        }
    }

    fn call(&mut self, req: String) -> Self::Future {
        let counter = Arc::clone(&self.counter);
        let permit = self.semaphore.try_acquire().unwrap();
        
        tokio::spawn(async move {
            // 模拟处理时间
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
            let count = counter.fetch_add(1, Ordering::SeqCst);
            drop(permit);
            Ok(format!("Request {} processed: {}", count, req))
        })
    }
}

#[tokio::main]
async fn main() {
    // 创建共享状态
    let counter = Arc::new(AtomicUsize::new(0));
    let semaphore = Arc::new(Semaphore::new(10)); // 限制并发数为10
    
    // 构建服务
    let service = MyService {
        counter: Arc::clone(&counter),
        semaphore: Arc::clone(&semaphore),
    };

    // 添加负载卸载中间件
    let svc = ServiceBuilder::new()
        .layer(LoadShedLayer::new())
        .service(service);

    // 模拟并发请求
    let mut handles = vec![];
    for i in 0..20 {
        let mut svc = svc.clone();
        handles.push(tokio::spawn(async move {
            match svc.ready().await {
                Ok(mut ready_svc) => {
                    ready_svc.call(format!("Request-{}", i)).await
                }
                Err(_) => Err("Service overloaded".to_string()),
            }
        }));
    }

    // 收集结果
    for handle in handles {
        match handle.await {
            Ok(Ok(res)) => println!("Success: {}", res),
            Ok(Err(e)) => println!("Error: {}", e),
            Err(e) => println!("Task failed: {}", e),
        }
    }
}

这个扩展示例演示了：

1. 自定义服务实现Service trait
2. 使用信号量限制并发请求数
3. 结合LoadShedLayer实现负载卸载
4. 模拟高并发场景下的请求处理
5. 统计请求处理计数

当并发请求超过信号量限制时，LoadShedLayer会立即拒绝请求，防止系统过载。