記事を書いたきっかけ

ITエンジニア部川上貴士です。電子薬歴システムである Solamichi というプロダクトの開発に携わっております。

グラフデータベースである Amazon Neptune を利用した機能の開発をしましたが、オンライン更新処理のスループットが思ったより出ず、スループット改善に多くの時間を費やしました。本記事を残すことで、Amazon Neptune を利用しようと考えている方、すでに利用されている方の一助になれば幸いです。
※本修正を行うことで、更新スループットがおよそ8倍に向上しました。

今回スループットで問題となった処理の概要

Solamichiでは患者基礎情報と呼ばれるデータを管理しており、このデータを Amazon Neptune に保持し、ユーザーに対して患者基礎情報を更新・参照するAPIエンドポイントを公開しています。クエリはGremlinを利用しています。ユーザーはリアルタイムで患者基礎情報を更新・参照することができます。今回スループットの問題が発生した箇所は、この更新APIになります。

患者基礎情報は患者を起点として患者基礎情報を時系列でEdgeでつなぎ、その日に入力のあった情報に対してEdgeでつなぐ構造になっています。 ※当日に入力がない場合は直近の日付の情報に対してEdgeをつないています。

※「〇」はVertex、→は「Edge」を指しています。

患者基礎情報の更新処理は、患者以下の1日分のVertex、および、Edgeの更新を行う必要があり、更新は患者単位で悲観ロックを獲得した後に行います。

Amazon Neptuneのクライアントとなるアプリケーションの実行環境、および、Amazon Neptuneのコンピュートリソースを使い切る前に更新リクエストが滞留する事象が発生しました。
アプリケーション側のトレースを確認すると、クエリ発行から60秒経過後にAmazon Neptuneからエラーが返却される挙動となっており、60秒もの間コネクションが占有されてしまっていました。

Amazon Neptuneは、複数のトランザクションによるロックで競合が発生した場合、最大60秒処理を待機する仕様となっており、更新負荷が高まった際にこれが多発していることがわかりました。
また、このロック競合の解決による待ちは、クライアント側でクエリタイムアウトを短く設定していたにもかかわらず60秒待ってしまう挙動となり、アプリケーション側のスレッドを使い尽くしていることがわかりました。

更新スループットを改善するために行ったこと

1) Upsertクエリの記述修正

GremlinでUpsertを行う際、開発当初はmergeV句がサポートされていなかったため、以下のように記述していました。

g.V('vertex-1')
 .fold()
 .coalesce(unfold(),
           addV('L-1').property(id, 'vertex-1'))
 .property(single, 'prop1', 'prop-value-1')

※ ID : vertex-1、Label : L-1、property : { 'prop1' = 'prop-value-1' } の vertex を Insert / Update するクエリ

上記のクエリは意図通りに動作はしますが、指定のVertexが存在しない状態でInsertとして動作する場合に全てのVertexのロックを獲得してしまうことが分かりました。
これは絶対に回避すべきです。幸いにも、クエリの記述方法により容易に回避可能です。
回避方法はmutationクエリの最後にid句を追加することです。

g.V('vertex-1')
 .fold()
 .coalesce(unfold(),
           addV('L-1').property(id, 'vertex-1'))
 .property(single, 'prop1', 'prop-value-1')
 .id()

※ このようなクエリが実装に混入しないよう、Gremlinのclientのクエリ発行イベントに対してInterceptorを実装し、id句で完結しないmutationクエリの発行を検出するようにしておくと安心かもしれません。

2) ロックの方式の見直し

同一リソースに対する同時更新が行われた場合、ロックの競合による最大60秒の待機が発生します。
このような長期の待機イベントは通常発生しないようにする必要があります。

例えば、以下のように5月16日分のデータを新規登録するケースを想定します。

このような更新を実施する場合、更新起点となる患者単位でクエリを直列化する必要があります。
当初、一連の更新処理は明示的なトランザクション（tx句を利用したbegin - commit）を利用し、トランザクション開始直後に患者ロック用Vertex（患者IDをVertexのIDの一部に利用）を登録することで排他制御を行っていました。

トランザクション開始
患者ロック用Vertexを登録
更新前のデータを取得
更新前のデータと、更新後のデータを比較し、更新クエリの実行
患者ロック用Vertexを削除
コミットしてトランザクション終了

※例外をフックした場合はロールバックしてトランザクション終了

しかし、この処理フローでは、同一患者に対する更新が同時期に行われた場合、患者ロック用Vertexの登録で競合が発生し、後続トランザクションは待機状態になります。
この待機状態のトランザクションが積み重なると、ロック解放待ちの待機列が伸びていき、コネクションプールが枯渇、アプリケーションがAmazon Neptuneへ新規のクエリを発行できなくなります。

ユーザー側が同一患者に対する同時更新を行うことは、業務上ほぼないと想定されていたため、悲観ロックを獲得する方式とし、ロック獲得待ちで待機させないようにしました。
ただし、Amazon NeptuneではSELECT - FOR UPDATE NOWAITのような悲観ロックを簡単に実現することができなかったため、以下のようなロック用Vertexを用いた排他処理を実装しました。

① 更新フローの修正

患者ロック用Vertexの登録・削除をトランザクションスコープの外で実行するようにし、ロックが獲得できなかった場合は待機せずにエラーになるようにします。

患者ロック用Vertexを登録
トランザクション開始
更新前のデータを取得
更新前のデータと、更新後のデータを比較し、更新クエリの実行
コミットしてトランザクション終了
患者ロック用Vertexを削除

② 患者ロック用Vertexの取り扱い修正

一連の更新処理がなんらかの理由で完了できなかった場合、患者ロック用Vertexが削除されずに残留する場合が想定されます。
患者ロック用Vertexにプロパティ：ロック開始タイムスタンプを追加し、このタイムスタンプ値を使ったロックタイムアウト判定を行うことで、永続ロックに陥ることを回避します。

具体的には、患者ロック用Vertexを登録する際に以下のようにロック獲得成功・失敗を判定します。

登録する患者ロック用Vertexと同一IDのVertexが存在しなかった ⇒ ロック獲得成功
登録する患者ロック用Vertexと同一IDのVertexが存在したが、ロック獲得したタイムスタンプを確認したところロックタイムアウト時間を超過していた ⇒ ロック獲得成功
登録する患者ロック用Vertexと同一IDのVertexが存在し、ロック獲得したタイムスタンプを確認したところロックタイムアウト時間を超過していなかった ⇒ ロック獲得失敗

上記の判定はトランザクショナルに行う必要があり、Gremlinはそれを実現するクエリが記述可能です。
以下、ID : 20240610abcdeに対するロック獲得クエリの例

g.V('Lock-20240610abcde')                                -- 患者ロックVertexのID
  .fold()
  .coalesce(
    __.unfold()                                         -- 同一IDの患者ロックVertexが既に存在する場合、このunfold句によって、既に存在するVertexが返却される。
      .choose(
        __.has('lock_timestamp', lt(1717999050)),               -- ロックタイムアウト判定。trueの場合はタイムアウトしている。
        __.property(single,'lock_timestamp', 1718001038217),        -- ロックタイムアウト判定がtrueの時に実行される。タイムスタンプを更新してロック獲得成功
        __.fail('lock-failure')),                               -- ロック獲得失敗のため、failステップによりクエリを失敗させる。
    __.addV('Lock')                                     -- unfold句の結果が空だった場合（同一IDの患者ロックVertexが既に存在しなかった場合）に実行される。ロック獲得成功のため、患者ロックVertexを登録して終了。
        .property(T.id, 'Lock-20240610abcde')
        .property(single, 'lock_timestamp', 1718001038217))
  .id()                                             -- id句は全体ロックにしないために必ずつける

3) トランザクション処理の修正

更新処理は明示的なトランザクション（tx句を利用したbegin - commit）を利用していました。
JavaアプリケーションからAmazon Neptuneへの接続は、org.apache.tinkerpop.gremlin-driver:3.6.2を利用していましたが、明示的なトランザクションを利用した場合、コネクションプールの接続を利用せずに、毎回新規コネクションを生成していることが分かりました。

※本件は将来的にライブラリの修正によって解消するかもしれません。

一方、Amazon Neptuneは明示的なトランザクションを利用せずとも、1クエリはトランザクショナルに動作する仕様があり、Gremlinは複数の更新を1クエリにまとめて記述することができます。
これによりコネクションプールを利用して更新クエリを行うことが可能です。
以下、EdgeとVertexの登録、更新、削除を一括で行うクエリの例を記載します。

g.E('Edge-1', 'Edge-2', …, 'Edge-N')               -- 削除対象のEdge IDを列挙（EdgeのID指定での抽出は、gの直後のみ記述可能）
  .fold()
  .aggrigate('drop-resource')                   -- drop句はクエリを終端してしまう（他の更新クエリや、VertexのID指定の抽出ができない）ため、必ず最後に記述する必要がある。そのため、エイリアス「drop-resource」に保持しておく
  .V('Vertex-1', 'Vertex-2', …, 'Vertex-N')           -- 削除対象のVertex IDを列挙
  .fold()
  .aggrigate('drop-resource')                   -- エイリアス「drop-resource」には削除対象のVertex ID、Edge IDの両方が格納される

  .addV(...)                                -- 以下、任意のVertex、Edgeの登録・更新クエリを列挙する
  .addV(...)
  .addV(...)
  .addE(...)
  .addE(...)

  .select('drop-resource').unfold().unfold().drop() -- 最後に、エイリアス「drop-resource」のリソースを削除。
  .id()

上記の修正を行う場合、更新フローは以下のようになります。