如何在Cosmos DB查询中限制RU/s吞吐量并添加退避延迟？

问题背景 #

现有一个异步查询Cosmos DB集合的方法，会占用最大吞吐量，可能导致数据库无响应。需要新增maxThroughputInRUsPerSecond参数，限制低优先级查询的RU消耗，已知可通过response.RequestCharge获取单次请求的RU消耗，需要解决两个问题：如何估算当前RU/s，以及如何添加退避延迟来近似达到限制值。

原方法代码：

public async IAsyncEnumerable<TObject> Query(Expression<Func<TObject, bool>> filter)
{
    FeedIterator<TObject> feedIterator = this.CosmosClient
        .GetContainer(this.DatabaseName, this.CollectionName).GetItemLinqQueryable<TObject>()
        .Where(filter)
        .ToFeedIterator();

    while (feedIterator.HasMoreResults)
    {
        FeedResponse<TObject> response = await feedIterator.ReadNextAsync();
        foreach (TObject r in response.Resource)
        {
            yield return r;
        }
    }
}

目标方法签名：

public async IAsyncEnumerable<TObject> Query(Expression<Func<TObject, bool>> filter, 
    double maxThroughputInRUsPerSecond)

一、RU/s 的估算方式 #

要计算当前的RU消耗速率，核心是关联单次请求的RU消耗和请求耗时：

1. 记录每次ReadNextAsync请求的开始时间和结束时间，计算请求耗时（单位：秒）。
2. 用单次请求的response.RequestCharge除以耗时，得到本次请求的瞬时RU/s。
3. 如果需要更稳定的速率判断，可以维护一个滑动窗口（比如最近3-5次请求的总RU和总耗时），计算窗口内的平均RU/s，避免单次请求的波动干扰。

二、退避延迟的实现逻辑 #

通过延迟后续请求的执行，降低整体的RU消耗速率，具体步骤：

1. 计算本次请求的实际RU/s，如果超过设定的maxThroughputInRUsPerSecond，说明请求速度过快。
2. 计算理想耗时：消耗本次请求的RU数，按照限制速率应该花费的时间（理想耗时 = 本次RU消耗 / 限制RU/s）。
3. 计算需要额外延迟的时间：延迟时间 = 理想耗时 - 实际耗时。如果结果为正，就执行对应时长的延迟；如果为负，说明实际耗时已经符合要求，无需延迟。

修改后的完整代码 #

public async IAsyncEnumerable<TObject> Query(Expression<Func<TObject, bool>> filter, 
    double maxThroughputInRUsPerSecond)
{
    FeedIterator<TObject> feedIterator = this.CosmosClient
        .GetContainer(this.DatabaseName, this.CollectionName).GetItemLinqQueryable<TObject>()
        .Where(filter)
        .ToFeedIterator();

    while (feedIterator.HasMoreResults)
    {
        // 记录请求开始时间（用UTC时间避免时区问题）
        var requestStartTime = DateTimeOffset.UtcNow;
        
        var response = await feedIterator.ReadNextAsync();
        
        // 计算本次请求耗时
        var requestEndTime = DateTimeOffset.UtcNow;
        double elapsedSeconds = (requestEndTime - requestStartTime).TotalSeconds;

        // 先返回查询结果
        foreach (var item in response.Resource)
        {
            yield return item;
        }

        // 避免除以0的异常情况
        if (elapsedSeconds <= 0)
        {
            continue;
        }

        // 计算本次请求的实际RU/s
        double actualRUsPerSecond = response.RequestCharge / elapsedSeconds;

        // 当实际速率超过限制时，计算并执行延迟
        if (actualRUsPerSecond > maxThroughputInRUsPerSecond)
        {
            double idealSeconds = response.RequestCharge / maxThroughputInRUsPerSecond;
            double delaySeconds = idealSeconds - elapsedSeconds;

            if (delaySeconds > 0)
            {
                await Task.Delay(TimeSpan.FromSeconds(delaySeconds));
            }
        }
    }
}

补充说明 #

• 这是一种近似控制方案，因为Cosmos DB的请求耗时、RU消耗会受数据分布、查询复杂度等因素影响，无法做到绝对精准，但能有效降低峰值吞吐量，避免挤占核心业务资源。
• 如果需要更平滑的速率控制，可以扩展滑动窗口逻辑：维护最近N次请求的总RU和总耗时，用平均速率计算延迟，减少单次请求波动的影响。
• 避免设置过低的maxThroughputInRUsPerSecond，否则会导致查询耗时急剧增加，甚至触发超时。

内容的提问来源于stack exchange，提问作者Mo B.