.NET 7 TokenBucketLimiter：如何获取实例调用TryReplenish补充令牌？

我正在使用.NET 7新增的Rate Limiting功能，希望通过TokenBucketLimiter限制用户对API的访问：每个用户初始拥有10次调用权限，且支持用户“付费”获取额外调用次数。据了解可通过ReplenishingRateLimiter的TryReplenish()方法实现，但不清楚如何获取已在Startup中注册的限流器实例。

Startup中的注册代码如下：

services.AddRateLimiter(options =>
{
    options.GlobalLimiter = PartitionedRateLimiter.Create<HttpContext, string>(resource =>
    {
        var request = resource.Request.Headers["token"].ToString();
        return RateLimitPartition.GetTokenBucketLimiter(request, _ => new TokenBucketRateLimiterOptions
        {
            TokenLimit = 10,
            ReplenishmentPeriod = TimeSpan.Zero,
            TokensPerPeriod = 10,
            QueueLimit = 0,
            AutoReplenishment = false,
        });
    });

    options.RejectionStatusCode = 429;
});

希望实现的测试端点代码：

[HttpGet("replenish")]
public async Task<IActionResult> DummyEndpoint()
{
    var token = Request.Headers["token"].ToString();
    var limiter = //somehow get my limiter based on the token used as partition key;
    limiter.TryReplenish();
    return Ok();
}

请问这种方式是否可行，或是我理解有误，需要采用其他实现方式？

你的思路方向是对的，但直接获取分区限流器实例的方式需要调整——PartitionedRateLimiter内部的分区实例不会直接暴露给外部获取。以下是可行的实现方案：

1. 自定义限流器存储 #

放弃直接从全局分区限流器中获取实例，改为自己维护每个用户（token）对应的TokenBucketLimiter实例：

• 先注册一个单例字典用于存储限流器：

services.AddSingleton<ConcurrentDictionary<string, TokenBucketLimiter>>();

• 修改Startup中的限流器注册逻辑，从自定义字典中获取或创建实例：

services.AddRateLimiter(options =>
{
    var limiterStore = services.BuildServiceProvider().GetRequiredService<ConcurrentDictionary<string, TokenBucketLimiter>>();
    
    options.GlobalLimiter = PartitionedRateLimiter.Create<HttpContext, string>(resource =>
    {
        var token = resource.Request.Headers["token"].ToString();
        var limiter = limiterStore.GetOrAdd(token, key => new TokenBucketLimiter(new TokenBucketRateLimiterOptions
        {
            TokenLimit = 10,
            ReplenishmentPeriod = TimeSpan.Zero,
            TokensPerPeriod = 10,
            QueueLimit = 0,
            AutoReplenishment = false,
        }));
        return RateLimitPartition.GetTokenBucketLimiter(_ => limiter);
    });

    options.RejectionStatusCode = 429;
});

• 在端点中直接从字典获取实例并调用TryReplenish：

private readonly ConcurrentDictionary<string, TokenBucketLimiter> _limiterStore;

public YourController(ConcurrentDictionary<string, TokenBucketLimiter> limiterStore)
{
    _limiterStore = limiterStore;
}

[HttpGet("replenish")]
public IActionResult DummyEndpoint()
{
    var token = Request.Headers["token"].ToString();
    if (_limiterStore.TryGetValue(token, out var limiter))
    {
        // 可指定补充的令牌数量，默认TryReplenish()会补充TokensPerPeriod配置的数量
        limiter.TryReplenish(10); // 示例：补充10次调用权限
    }
    return Ok();
}

2. 关键注意事项 #

• 要处理token为空的场景，避免空键引发异常；
• 如果是分布式部署，内存字典无法跨实例共享，需改用Redis等分布式缓存存储用户令牌剩余量，此时可能需要自定义限流器逻辑，而非直接使用TokenBucketLimiter；
• TryReplenish()默认补充TokensPerPeriod配置的数量，也可传入指定参数灵活控制补充的调用次数。

3. 原思路说明 #

原思路中直接获取分区内限流器实例不可行，因为PartitionedRateLimiter的内部分区是封装状态，没有提供公开的获取方法。通过自定义存储的方式，既能满足全局限流需求，又能灵活操作指定用户的限流器实例。

内容的提问来源于stack exchange，提问作者Paweł Wąsowski