種類を分けて流量制限をする(System.Threading.RateLimiting)

始めに

System.Threading.RateLimitingというものが.NET 7より追加された。
詳しくは別記事を見て欲しいが、例えばリソースの同時アクセス制限機能で、URLのパス部分等で区切って流量制限をしたい場合がある。
そのような時に、同パッケージでは包括的に管理可能な仕組みを提供しており、それが今回紹介するPartitionedRateLimiterである。

これは短く言うと、 IDictionary<TResource, RateLimiter>をうまいことやってくれる仕組み である。

RateLimiterの基本的な概念や使い方については、System.Threading.RateLimitingの記事を別途書いているので、そちらを参考にしてほしい。

基本的な使い方

まず、単純な使い方は以下のようになる。細かいクラスの説明は省く。

// using System.Threading.RateLimiting;
void PartitionDemo()
{
    // 第一型引数(TResource)が、外部で指定するキーの型で、第二引数(TKey)が内部で使うキーの型
    // トークン取得要求(AttemptAquire)するごとに中の処理が実行される。
    // この例ではTResourceにDateTimeを選んでいるが、実際はHttpContext等が入ることが多いと思われる。
    using var plimiter = PartitionedRateLimiter.Create<DateTime, string>(idx =>
    {
        // 区分けのキー+中身のRateLimiterの生成処理を返す
        // 内部でキャッシュするため、TKeyが同一の場合はRateLimiterの生成処理は通らない
        return RateLimitPartition.GetConcurrencyLimiter(DateTime.ToString("dd"), key =>
        {
            return new ConcurrencyLimiterOptions() { PermitLimit = 10, };
        });
    });
    // 第一引数がTResourceで、第二引数がtokenの要求数
    using(var lease = plimiter.AttemptAquire(DateTime.Now, 1))
    {
        // 後は通常のRateLimiterの処理と一緒
    }
}

PartitionedRateLimiter.Create<TResource> でRateLimiterを作り、
それを運用していく形になる。
通常のRateLimiterと同じく、PartitionedRateLimiterもシングルトンで運用するのが望ましい。
中で生成するRateLimiterの生存期間等は、PartitionedRateLimiterの内部で管理することになる。

型の説明

主に覚えておく必要があるのは下記

• PartitionedRateLimiter<TResource>
• PartitionedRateLimiter
• RateLimitPartition<TKey>
• RateLimitPartition

PartitionedRateLimiter<TResource>

実際に流量制限を行う型。
この型を生成して保持することとなる。
あまり頻繁に生成するようなことは想定していないので、シングルトンで管理するのが望ましい。

この型自体はabstract classで、デフォルト実装はpublicではなく、PartitionedRateLimiterクラスのCreateを使って生成する。
abstract classなので独自実装も可能だが、大体の場合はライブラリ側で提供されるもので用は足りると思う。

生成後は、以下二つのトークン取得メソッドを使ってleaseを取得して、通常のRateLimiterと同じように運用していく。

• ValueTask<RateLimitLease> AquireAsync(TResource resource, int requiredToken = 1, CancellationToken cancellationToken = default)
• RateLimitLease AttemptAquire(TResource resource, int requiredToken = 1)

どちらも、TResource resourceで指定された値に基づきRateLimiterを選択し、Aquireの結果を返す。
返ってくるRateLimitLeaseは、内部で使用しているRateLimiterに準ずる。
使い終わったら、DisposeあるいはDisposeAsyncを行う事。

独自実装する場合

独自実装する場合は最低限下記APIを実装する

API	内容
public abstract RateLimiterStatistics? GetStatistics(TResource resource)	統計情報(現在の取得可能トークン数等)を返さなければいけない
protected abstract RateLimitLease AttemptAcquireCore(TResource resource, int permitCount)	トークン取得時に内部で呼ばれる(同期版)
protected abstract ValueTask<RateLimitLease> AcquireAsyncCore(TResource resource, int permitCount, CancellationToken cancellationToken)	トークン取得時に内部で呼ばれる(非同期版)

これに加えて、必要ならばDisposeあるいはDisposeAsyncを実装すること。
デフォルトでは何もしてないので、十中八九実装することになると思われるが。

TResourceを別の型に変換して使う(WithTranslatedKey)

TResourceで指定した型ではなく、別の型で受け取りたい場合は、下記のWithTranslatedKeyを使う。
PartitionedRateLimiter<TOuter> PartitionedRateLimiter<TResource>.WithTranslatedKey<TOuter>(Func<TOuter, TResource> keyAdapter, bool leaveOpen)

keyAdapterはTOuterからTResourceへの変換処理を指定する。leaveOpenは、新しく生成したPartitionedRateLimiterをDisposeする時、元のPartitionedRateLimiterをDisposeするかどうかを指定する(falseで一緒にDispose、trueでそのまま)。

例えば既にPartitionedRateLimiter<string>があるが、それを更に包んで最上位の処理ではPartitionedRateLimiter<HttpContext>で使いたい、という時等に使う。

AttemptAquireで返ってくるRateLimitLeaseは、元のParitionedRateLimiterが持つRateLimiterに準ずる。

PartitionedRateLimiter

PartitionedRateLimiter<TResource> に関するデフォルト実装のインスタンス生成やその他ユーティリティ的な機能を提供する静的クラス。
ユーザーはここを起点に各種PartitionedRateLimiterを生成することとなる。

PartitionedRateLimiter<TResource> Create<TResource, TKey>(Func<TResource, RateLimitPartition<TKey>> partitioner, IEqualityComparer<TKey> comparer = null)

RateLimitPartitionについては後述する。

PartitionedRateLimiter<TResource, TKey>を生成する。
partitionerはトークン取得要求の時に常に呼ばれる。

PartitionedRateLimiter<TResource, TKey>は内部でTKeyの値で各RateLimiterを管理しており、
partitionerで返されるRateLimitPartition<TKey>.PartitionKeyが既に内部リストに存在している場合、
生成処理は呼ばず、同じRateLimiterを使いまわす。
comparerが指定されていた場合、内部リスト確認時の比較関数として使われる。

PartitionedRateLimiter<TResource> CreateChained<TResource>(params PartitionedRateLimiter<TResource>[] limiters)

各種PartitionedRateLimiterをまとめたParititionedRateLimiterを作る。
これで作ったPartitionedRateLimiterは、引数で指定されたものを次々実行していき、全てのAquireが成功したら成功と判断する。

AttemptAcquireで返ってくるRateLimitLeaseは、複数のRateLimitLeaseを内包しているので必ずDisposeを行う。
TryGetMetadataすると、内部で保持しているRateLimitLeaseの中で最初にTryGetMetadataが成功したものの情報を返す。

RateLimitPartition<TKey>

デフォルト実装で必要な、パーティションのキーと生成処理を渡すための構造体。
以下のプロパティを持っている。

• TKey PartitionKey
  • RateLimiterに関連付けられたキー
• Func<TKey, RateLimiter> Factory
  • RateLimiter生成処理

RateLimiterの生成処理と、RateLimiterに関連付けられたキーを保持する型。
これはトークン取得処理の都度生成される想定。
トークン取得の時にPartitionedRateLimiterの中でTKeyの値を見て、RateLimiterを生成するかどうかが決定される。
なので、Factoryで指定された処理というのは毎回実行されるわけではない。
RateLimiterが生成される時は、デフォルト実装ではFactoryにPartitionKeyが渡される。

生成は自分でnewすることも可能だが、後述するRateLimitPartitionで生成することも可能。

RateLimitPartition

RateLimitPartition<TKey> を生成するためのユーティリティを集めた静的クラス。

RateLimitPartition<TKey> Get<TKey>(TKey key, Func<TKey partitionKey, RateLimiter> factory)

任意のRateLimiterを生成する。単純にPartitionKeyとFactoryに値をセットしているだけ。

RateLimitPartition<TKey> Get[各RateLimiter]<TKey>(TKey partitionKey, Func<TKey, [オプション型]> factory)

上記Getを更に各実装に細分化したもの。内部的にやっていることは上記Getとほぼ同じ。

PartitionedRateLimiter<TResource>とRateLimitPartition<TKey>の関係

ここで筆者が個人的に疑問に思ったこととして、"RateLimitPartitionってなんでいるの?"というのがあった。
自分なりに考えたところ、PartitionedRateLimiterで渡したい型でstring等の直接的な型を指定すると、
実行ユーザー側で型変換の処理が毎回入り、上位層に区分けのロジックが入ってしまうことになり、
うまく共通化できてないと、漏れが発生してしまうことある。

そこで、区分けのロジックは可能な限り上位層に出さないようにすれば、上位層では大雑把な型(例えばHttpContext)さえ渡してしまえば
後は共通のロジックで区分けをしてくれるので、漏れも少なくなるということになる。

また、区分けの条件を変えたいとき等は、生成時の処理で変えてしまえばいいので、より網羅的に修正を行うことが可能になる。

終わりに

RateLimiterの区分けを考える場合、IDictionary<TResource, RateLimiter>で自前で管理するというのが真っ先に思いつくところだが、
今回紹介したPartitionedRateLimiterの仕組みを使えば、ライブラリの方に生成破棄等のロジックの大部分を任せることができ、
RateLimiterを利用する際の手間も減ると思われるので、あてはまる場合があれば試してみるのもいいと思う。