Azure Functionsでの大量データ処理とグレースフルシャットダウン（後編）

技術本部Sansan Engieering Unit Data Hubグループの藤原です。前回はAzure Functions好きにしか刺さらないとがった内容を書いてしまいました。反省しているので、今回は間口を広げて.NETの標準クラスライブラリ好きにも刺さる内容になっています。

前回、グレースフルシャットダウン対応のバグを修正したパッチを適用したが、実はうまくいっていなかった……というところまで書きました。今回はさらに深く入り込み、その問題も直した話になっています。一言でいうと、Event Hubトリガーを使っている場合、SDKのバージョンによってはメッセージが処理されないことがあります（Microsoft.Azure.WebJobs.Extensions.EventHubs の v6.1.0 以上を使用する必要があります）。また、 LinkedCancellationTokenSource を使う場合、リンク先のトークンとの競合状態があるので注意が必要です。

なお、このブログでは事象を理解するための前提知識としてAzure Functionsの内部構造に触れています。詳しくない方は前回をご参照ください。また、今回は .NET の内部実装にも踏み込んでいます。

さらなるグレースフルシャットダウン処理のバグ
CancellationToken、CancellationTokenSource、LinkedCancellationTokenSource
グレースフルシャットダウンかどうかの判定
調査と修正
仲間を募集中です！

さらなるグレースフルシャットダウン処理のバグ

前回説明したグレースフルシャットダウン対応のパッチを適用した後、Data Hubの運用メンバーからメッセージロストが再発したと報告されました。

この場合、可能性としては以下のいずれかです。

他にもメッセージロストを引き起こす問題がある。
前回適用したパッチでは完全に問題が修正されていない。

手始めに、「上記のパッチでは完全に問題が修正されていない」可能性を考慮することにしました。あらゆる可能性を模索するには時間がかかるためです。幸運にも、パッチ適用の際にSDKのソースコードを読み込んでいたため、グレースフルシャットダウン時の挙動を再現する方法の勘所が付いていました。そのため、SDKのコードにデバッグコードを仕込みつつ、グレースフルシャットダウンの状態を再現させることができました。

実行してみると、グレースフルシャットダウン時にチェックポイントは進まないようでした。しかし、繰り返し実行してみると、ときどきチェックポイントが進んでしまうことがありました。なぜでしょうか。今回も、Azure Functionsと、さらには.NETの標準クラスライブラリの実装に踏み込みながら、何が起こっていたのかを説明します。

CancellationToken、CancellationTokenSource、LinkedCancellationTokenSource

ここで、.NETの標準クラスライブラリにあまり詳しくない方のためにちょっと補足します。.NET標準ライブラリに慣れている方は図の後まで読み飛ばしていただいて大丈夫です。

CancellationTokenとCancellationTokenSource

.NETの CancellationToken は、キャンセル状態になっているかどうかをアプリケーションに伝える役割を持つオブジェクト（構造体）です。CancellationToken は以下の機能を持ちます。

キャンセル状態になっていることを示す IsCancellationRequested プロパティ
キャンセル状態になったときに実行されるコールバックを登録する Register メソッド（一部のチェックを省略した UnsafeRegister もあります）
キャンセル状態になっていた時にキャンセル済み例外をスローする ThrowIfCancellationRequested メソッド

ここで、注目すべきポイントが2つあります。まず、プロパティ名が IsCancelled ではなく、IsCancellationRequested になっています。これは、CancellationToken はあくまでキャンセルが要求されていることだけを伝え、実際に処理をキャンセルするのは CancellationToken を利用する側の役割だからです。次に、CancellationToken にはキャンセル状態にする（キャンセル要求をする）機能がありません。これは、 CancellationTokenSource という別のオブジェクト（クラス）に実装されています。

実は、CancellationToken は、CancellationTokenSource の Token プロパティから返されるオブジェクトです。取得元の CancellationTokenSource の Cancel メソッドが呼び出されると、その CancellationTokenSource から取得された全ての CancellationToken がキャンセル状態になります。実際のところ、キャンセル状態や後述するコールバックの管理は CancellationTokenSource に実装されており、CancellationToken はキャンセル状態の問い合わせとコールバック登録用のインターフェースであると言えるでしょう。

LinkedCancellationTokenSource

さて、複数の CancellationToken のいずれかがキャンセル状態になったときに処理をキャンセルしたい場合があるかもしれません。あるいは、受け取った CancellationToken とは別に、独自のキャンセル要求を行いたい場合があるかもしれません。たとえば、一定時間経過した後に処理をタイムアウトさせたい場合が考えられます。これらのケースをサポートするために、1つ以上の CancellationToken に「リンクした」CancellationTokenSource、すなわち LinkedCancellationTokenSource を作成できます。LinkedCancellationTokenSource から取得した CancellationToken は、リンクされた CancellationToken のいずれか（の取得元 CancellationTokenSource）、または LinkedCancellationTokenSource 自身がキャンセルされたときにキャンセル状態になります。

LinkedCancellationTokenSourceの実装詳細

LinkedCancellationTokenSource の実装はどうなっているのでしょうか？　.NETのコードを見てみると、リンクする CancellationToken の Register にコールバックを登録し、そのコールバックで LinkedCancellationTokenSource 自身のキャンセル状態を更新しています。その理由は明記されていませんが、推測する限り、次のような理由だと思われます。仮に、単にリンクされた CancellationToken の IsCancellationRequested の論理積をとる実装にした場合、複数の CancellationToken をリンクする場合に、処理のオーダーは O(N) になります。キャンセル状態はプロパティで確認されるので、その性能特性は O(1) であることが期待されます。プロパティの参照は低コストであることを想定し、ループ内で繰り返し参照されることが一般的だからです。コールバックによって LinkedCancellationTokenSource のキャンセル状態を直接書き替えるようにすると、キャンセル状態の取得と更新のいずれも O(1) で済むはずです。

さらに、このコールバックの実行やキャンセル状態はロックフリーで実行されており、高い並行性を確保しています。コールバックの実行中に LinkedCancellationTokenSource （から取得した CancellationToken）の IsCancellationRequested を参照しても、ロックの解放待ちで待たされるようなことはありません。ここでも、プロパティの参照が高コストにならないよう考慮されているようです。さらに、コールバックは後入れ先出し（LIFO）で処理されます。これは、あるメソッドの呼び出し元と呼び出し先の両方でコールバックを登録したとき、呼び出し先のコールバックが先に呼び出されるようにしていると推察されます。

グレースフルシャットダウンかどうかの判定

先ほど説明したように、グレースフルシャットダウンの場合にチェックポイントを進めないようにするのはEvent Hubリスナーの役割です。

Event Hubリスナーは、Azure Functionsランタイムからグレースフルシャットダウンが通知されると、関数に渡している CancellationToken をキャンセル状態にします。そのため、この CancellationToken がキャンセル状態になっていれば、チェックポイント処理を行わないという単純な判定をしています。

さらに、追加の処理として、Event Hubリスナーではオーナーシップを失った場合の考慮も行っています。オーナーシップは、スケールアウトの結果、現在のインスタンスが担当していたEvent Hubパーティションを別のインスタンスが処理するようになった場合に失われます。オーナーシップを失ったケースをフォローするためには、オーナーシップを失ったときキャンセル状態になる別の CancellationToken を併せて使用する必要があります。そのため、Event Hubリスナーはこれらの CancellationToken にリンクされた LinkedCancellationTokenSource を使用しています。先ほど述べたように、グレースフルシャットダウンの判定処理も、この LinkedCancellationTokenSource から取得した CancellationToken で判定しています。

なお、関数実行にはFunctionsランタイムのキャンセル状態を判定するCancellationToken だけを渡すようになっていました。これは上記のオーナーシップを失ったときの対応を入れたこの PR で実装され、SDKのv5.5.0としてリリースされたものですが、issueコメントを見る限り単なるうっかりミスだったようです。実は、これが今回の問題の原因でした。

調査と修正

さて、アプリケーション実行がキャンセルされているのにチェックポイント処理が進んでしまう現象が発生したとき、実際には何が起こっていたのでしょうか？

キャンセルされているはずなのにされていない

今回は実際にSDKのコードに手を入れ、デバッグコードを入れて確認することにしました。そうすると、関数実行に渡している CancellationToken の IsCancellationRequested が true になっているのに、チェックポイント処理で判定に使用している LinkedCancellationTokenSource 由来の CancellationToken の IsCancellationRequested は false のままだったのです。いったい何が起こっているのでしょうか。

先ほど説明したように、LinkedCancellationTokenSource のキャンセル状態は、リンクされた CancellationToken に登録したコールバックで変更されます。コールバックは順番に（LIFOで）実行されるため、リンクされた CancellationToken がキャンセル状態になった後、 LinkedCancellationTokenSource がキャンセル状態になるまでには若干のタイムラグが生じます。さらに、このコールバック実行と IsCancellationRequested はロックフリー実装なので、このタイムラグの間も、LinkedCancellationTokenSource 由来の CancellationToken の IsCancellationRequested の状態はブロックされることなく確認でき、false を返します。

LinkedCancellationTokenSourceのキャンセル状態の更新と競合状態

競合状態の解消

この競合状態を解消する方法は簡単です。 LinkedCancellationTokenSource 由来の CancellationToken を関数に渡すか、コールバック判定処理でリンクされた CancellationToken（2つしかありません）それぞれの状態を見るかのどちらかを行えばよいのです。先ほど説明したように、ドレイン対応PR が入る前は関数実行に LinkedCancellationTokenSource 由来の CancellationToken を渡しているのですから、ドレインモードを考える必要がない限りこのバージョンの SDK を使うように修正さればよさそうです。幸いにも当該処理ではドレインモードを考慮する必要がなかったので、このグレースフルシャットダウンのバグのみが直っているバージョンに戻した結果、今のところ問題なく動いています。

なお、これだけでは別の問題のためのパッチが出たときに現象が再発して困ってしまうので、SDK チームにフィードバックする必要もあります。そのため、この問題はissueとして報告しており、既に修正パッチがマージされています。ちなみに、結局、 LinkedCancellationTokenSource 由来の CancellationToken を関数に渡すように修正されました。

本当に？

ところで、ここで疑問が生じます。コールバックが実行されてから LinkedCancellationTokenSource の状態が変わるまでにタイムラグがあることはわかりましたが、このタイムラグの間にチェックポイント判定処理まで進んでしまう可能性はどの程度あるのでしょうか？　コードの行数を考えても、コールバックの完了までのCPU命令数と、関数実行がキャンセルされてリトライループを抜けてチェックポイント判定処理までのCPU命令数にはかなりの差があるように思えます。そこまで低い確率を引き当てられるなら、まず宝くじを連番で3枚買うべきでしょう。他にも要因がありそうです。

試しに、グレースフルシャットダウンの処理でコールスタックを取ってみました。すると、以下のように、グレースフルシャットダウンの通知による CancellationTokenSource の Cancel() メソッドの呼び出しがコールスタックに含まれていたのです。

これは、以下の理由が合わさって発生したと考えられます。

Task.Delay のキャンセル処理は CancellationToken のコールバックとして実装されている。
このコールバックはLIFOで実行されるため、LinkedCancellationTokenSource の状態変更よりも、この Task.Delay のキャンセル処理が実行される。
キャンセルが行われたとき、.NETの非同期処理の実装として、そこまでの await 式の後続の処理が、継続（continuation）として、可能な限り効率的に実行される。具体的には、同じスレッドで連続して継続が実行される。
Event Hubリスナーは関数呼び出しを await し、その後でチェックポイント判定処理を行っている。

コールバックで継続が実行されたために競合状態でチェックポイント判定処理が実行された

つまり、グレースフルシャットダウンが行われたとき、LinkedCancellationTokenSource がキャンセル状態になる前に、チェックポイント判定処理が行われているということです。

実際のアプリケーション実行において、グレースフルシャットダウンの発生時に、非同期での通信のキャンセル対応やタイムアウト処理のために CancellationToken のコールバックが使用されている可能性は高いと考えられます。そのため、この継続が同期的に実行されるという事象によりチェックポイント処理判定が不正になる可能性は高いと言えそうです。