如何持久化.NET RateLimiting中TokenBucketRateLimiter的令牌？

问题本质 #

默认的TokenBucketRateLimiter依赖内存存储令牌状态，服务重启后所有状态会被清空。要实现持久化，需将令牌桶核心状态（剩余令牌数、最后补充时间）存储到PostgreSQL中，替代内存存储。

方案一：自定义PostgreSQL持久化令牌桶 #

1. 创建PostgreSQL存储表 #

在PostgreSQL中创建用于存储用户令牌状态的表，保存每个用户的令牌桶运行数据：

CREATE TABLE IF NOT EXISTS user_rate_limits (
    username VARCHAR(255) PRIMARY KEY,
    remaining_tokens INT NOT NULL,
    last_replenished_at TIMESTAMPTZ NOT NULL,
    token_limit INT NOT NULL,
    replenishment_period INTERVAL NOT NULL,
    tokens_per_period INT NOT NULL
);

2. 实现PostgreSQL令牌桶逻辑 #

创建自定义的PostgreSqlTokenBucketRateLimiter类，继承RateLimiter，通过Npgsql操作数据库完成令牌的获取、补充和状态更新：

public class PostgreSqlTokenBucketRateLimiter : RateLimiter
{
    private readonly string _username;
    private readonly TokenBucketRateLimiterOptions _options;
    private readonly NpgsqlDataSource _dataSource;

    public PostgreSqlTokenBucketRateLimiter(string username, TokenBucketRateLimiterOptions options, NpgsqlDataSource dataSource)
    {
        _username = username;
        _options = options;
        _dataSource = dataSource;
    }

    protected override async ValueTask<RateLimitLease> AcquireAsync(int permitCount, CancellationToken cancellationToken)
    {
        using var conn = await _dataSource.OpenConnectionAsync(cancellationToken);
        using var tx = await conn.BeginTransactionAsync(cancellationToken);

        try
        {
            // 读取或初始化用户令牌状态
            var state = await GetOrInitializeTokenState(conn, cancellationToken);
            
            // 计算周期内可补充的令牌数
            var now = DateTimeOffset.UtcNow;
            var timeSinceLastReplenish = now - state.LastReplenishedAt;
            var tokensToAdd = (int)(timeSinceLastReplenish.TotalSeconds / _options.ReplenishmentPeriod.TotalSeconds) * _options.TokensPerPeriod;
            
            var newRemainingTokens = Math.Min(state.RemainingTokens + tokensToAdd, _options.TokenLimit);
            var newLastReplenished = now;

            // 判断是否允许获取令牌
            if (newRemainingTokens >= permitCount)
            {
                newRemainingTokens -= permitCount;
                await UpdateTokenState(conn, newRemainingTokens, newLastReplenished, cancellationToken);
                await tx.CommitAsync(cancellationToken);
                return RateLimitLease.Success;
            }
            else
            {
                await tx.RollbackAsync(cancellationToken);
                return RateLimitLease.Failed;
            }
        }
        catch
        {
            await tx.RollbackAsync(cancellationToken);
            throw;
        }
    }

    private async Task<TokenState> GetOrInitializeTokenState(NpgsqlConnection conn, CancellationToken cancellationToken)
    {
        var cmd = new NpgsqlCommand(@"
            SELECT remaining_tokens, last_replenished_at, token_limit, replenishment_period, tokens_per_period
            FROM user_rate_limits
            WHERE username = @Username", conn);
        cmd.Parameters.AddWithValue("@Username", _username);

        using var reader = await cmd.ExecuteReaderAsync(cancellationToken);
        if (await reader.ReadAsync(cancellationToken))
        {
            return new TokenState
            {
                RemainingTokens = reader.GetInt32(0),
                LastReplenishedAt = reader.GetDateTimeOffset(1),
                TokenLimit = reader.GetInt32(2),
                ReplenishmentPeriod = reader.GetTimeSpan(3),
                TokensPerPeriod = reader.GetInt32(4)
            };
        }
        else
        {
            var initialState = new TokenState
            {
                RemainingTokens = _options.TokenLimit,
                LastReplenishedAt = DateTimeOffset.UtcNow,
                TokenLimit = _options.TokenLimit,
                ReplenishmentPeriod = _options.ReplenishmentPeriod,
                TokensPerPeriod = _options.TokensPerPeriod
            };

            var insertCmd = new NpgsqlCommand(@"
                INSERT INTO user_rate_limits (username, remaining_tokens, last_replenished_at, token_limit, replenishment_period, tokens_per_period)
                VALUES (@Username, @RemainingTokens, @LastReplenishedAt, @TokenLimit, @ReplenishmentPeriod, @TokensPerPeriod)", conn);
            insertCmd.Parameters.AddWithValue("@Username", _username);
            insertCmd.Parameters.AddWithValue("@RemainingTokens", initialState.RemainingTokens);
            insertCmd.Parameters.AddWithValue("@LastReplenishedAt", initialState.LastReplenishedAt);
            insertCmd.Parameters.AddWithValue("@TokenLimit", initialState.TokenLimit);
            insertCmd.Parameters.AddWithValue("@ReplenishmentPeriod", initialState.ReplenishmentPeriod);
            insertCmd.Parameters.AddWithValue("@TokensPerPeriod", initialState.TokensPerPeriod);

            await insertCmd.ExecuteNonQueryAsync(cancellationToken);
            return initialState;
        }
    }

    private async Task UpdateTokenState(NpgsqlConnection conn, int remainingTokens, DateTimeOffset lastReplenishedAt, CancellationToken cancellationToken)
    {
        var cmd = new NpgsqlCommand(@"
            UPDATE user_rate_limits
            SET remaining_tokens = @RemainingTokens, last_replenished_at = @LastReplenishedAt
            WHERE username = @Username", conn);
        cmd.Parameters.AddWithValue("@Username", _username);
        cmd.Parameters.AddWithValue("@RemainingTokens", remainingTokens);
        cmd.Parameters.AddWithValue("@LastReplenishedAt", lastReplenishedAt);

        await cmd.ExecuteNonQueryAsync(cancellationToken);
    }

    protected override RateLimitLease Acquire(int permitCount)
    {
        throw new NotImplementedException("同步获取令牌逻辑未实现，生产环境建议优先使用异步方法");
    }

    private class TokenState
    {
        public int RemainingTokens { get; set; }
        public DateTimeOffset LastReplenishedAt { get; set; }
        public int TokenLimit { get; set; }
        public TimeSpan ReplenishmentPeriod { get; set; }
        public int TokensPerPeriod { get; set; }
    }
}

3. 注册自定义令牌桶策略 #

在Program.cs中替换原有内存令牌桶配置，注册自定义的PostgreSQL令牌桶：

// 注册Npgsql数据源
builder.Services.AddNpgsqlDataSource(builder.Configuration.GetConnectionString("PostgreSQL"));

// 注册限流策略
builder.Services.AddRateLimiter(limiterOptions =>
{
    limiterOptions.AddPolicy(userPolicyName, context =>
    {
        string username = Utils.getUsername(context);
        var options = new TokenBucketRateLimiterOptions
        {
            TokenLimit = myOptions.TokenLimit,
            QueueProcessingOrder = QueueProcessingOrder.OldestFirst,
            QueueLimit = 0,
            ReplenishmentPeriod = TimeSpan.FromDays(myOptions.ReplenishmentPeriod),
            TokensPerPeriod = myOptions.ReplenishmentPeriod,
            AutoReplenishment = myOptions.AutoReplenishment
        };

        var dataSource = context.RequestServices.GetRequiredService<NpgsqlDataSource>();
        return RateLimitPartition.Get<PostgreSqlTokenBucketRateLimiter>(username, _ => 
            new PostgreSqlTokenBucketRateLimiter(username, options, dataSource));
    });
});

方案二：分布式缓存间接持久化 #

若不想完全自定义令牌桶，可通过IDistributedCache将内存令牌桶状态序列化后存储到PostgreSQL（依赖Npgsql分布式缓存实现），需处理并发更新冲突：

1. 注册Npgsql分布式缓存 #

builder.Services.AddDistributedNpgsqlCache(options =>
{
    options.ConnectionString = builder.Configuration.GetConnectionString("PostgreSQL");
    options.SchemaName = "public";
    options.TableName = "distributed_cache";
});

2. 包装内存令牌桶同步状态 #

创建包装类，每次操作令牌桶后将状态序列化存入缓存，初始化时从缓存读取状态，需通过乐观锁等机制处理并发冲突。

注意事项 #

• 并发控制：所有数据库操作必须包裹在事务中，避免多请求同时更新导致数据不一致。
• 性能优化：可添加本地内存缓存层，定时同步状态到PostgreSQL，减少数据库访问频次。
• 配置一致性：令牌桶核心配置（TokenLimit、ReplenishmentPeriod等）建议统一存放在配置中心，避免代码与数据库配置不一致。

内容的提问来源于stack exchange，提问作者pkExec