外部API呼び出しを含む処理を、advisory lockではなく自作のハートビートで排他制御する

問題

時間のかかる後処理を、複数のリクエストが同時に始めてしまうことがある。
PostgreSQLのadvisory lockでは対応しきれないことがある。

具体例

ユーザーの操作をきっかけに、後処理を始める。
後処理の中身は「外部APIへの問い合わせ」と「DBへの保存」の2段階。
外部APIの応答には数秒から数分かかることがある。

同じユーザーが、別のタブや別の操作から同時にリクエストが送られ、後処理が同時に2つ動いてしまう可能性がある。

解決方法

DBに、状態を書くための行を1つだけ用意する。
最初に気づいたリクエストが、その行に自分の担当開始時刻を記録する。
これを「担当権を取る」と呼ぶ。

担当権を持っている間、一定間隔（例として30秒）で、その同じ行に「まだ動いている」という時刻を書き込み続ける。
これがハートビート。

別のリクエストが同じ処理を始めようとしたときは、まずこの行を見に行く。
直近のハートビートが一定時間（例として2分）以内なら、担当中とみなして何もしない。
一定時間ハートビートが止まっていたら、担当が死んだとみなして、自分が担当権を奪う。

advisory lockとの使い分け

PostgreSQLにはadvisory lockがある。

• advisory lockは、短い処理用
  DBの読み書きだけ

• ハートビートは、長い処理用
  外部へのAPI問い合わせなど

実装例（疑似コード）

担当権テーブルには、対象のID・使い捨ての合言葉（lease_token）・更新時刻の3列を用意する。

create table job_lease (
  target_id text primary key,
  lease_token uuid not null,
  updated_at timestamptz not null
);

担当権を取るたびに、新しいlease_tokenを発行する。
このトークンは、以降のハートビート・保存・解放、すべてで使い回す。

async function acquireLease(targetId: string, timeoutMs: number) {
  const leaseToken = crypto.randomUUID();

  const result = await db.execute(sql`
    insert into job_lease (target_id, lease_token, updated_at)
    values (${targetId}, ${leaseToken}, now())
    on conflict (target_id) do update
      set lease_token = ${leaseToken}, updated_at = now()
      where job_lease.updated_at < now() - interval '${timeoutMs} milliseconds'
    returning lease_token
  `);

  // 1行返ってきたら、担当権を取れたということ。このtokenを持ち運ぶ
  return result.rowCount > 0 ? leaseToken : null;
}

ハートビートは、target_idだけでなく、lease_tokenも条件に含めて更新する。

function startHeartbeat(targetId: string, leaseToken: string, intervalMs: number) {
  return setInterval(async () => {
    const result = await db.execute(sql`
      update job_lease set updated_at = now()
      where target_id = ${targetId} and lease_token = ${leaseToken}
    `);

    // 0行なら、担当権はすでに他のリクエストに奪われている
    if (result.rowCount === 0) {
      // ここでタイマーを止め、処理そのものも中断する
    }
  }, intervalMs);
}

lease_tokenが一致しなければ、更新は0行のまま失敗する。
担当権を奪われた旧担当は、ここで自分がもう担当ではないと気づける。

保存する直前の確認だけでは、確認と保存の間に隙間ができる。
確認した直後に担当権を奪われ、その後で保存してしまうと、二重保存を防げない。

これを防ぐには、確認と保存を1つのトランザクションにまとめる。
確認の時点で担当権の行をロックし、保存が終わるまでロックを持ち続ける。
ロックしている間、他のリクエストはその行を更新できないので、確認から保存の間に横から奪われることがなくなる。

async function saveIfStillOwner(targetId: string, leaseToken: string, resultData: unknown) {
  return db.transaction(async (tx) => {
    const lease = await tx.execute(sql`
      select lease_token from job_lease
      where target_id = ${targetId}
      for update
    `);

    if (lease.rows[0]?.lease_token !== leaseToken) {
      // 担当権はもう自分のものではないので、結果は保存せず捨てる
      return false;
    }

    await tx.execute(sql`
      insert into job_result (target_id, data)
      values (${targetId}, ${resultData})
      on conflict (target_id) do update set data = excluded.data
    `);

    return true;
  });
}

処理が終わったら、同じlease_tokenを条件にして行を消す。

async function releaseLease(targetId: string, leaseToken: string) {
  await db.execute(sql`
    delete from job_lease
    where target_id = ${targetId} and lease_token = ${leaseToken}
  `);
}

target_idだけで削除すると、後から担当権を取った別のリクエストの行まで消してしまう。
lease_tokenも条件に含めることで、自分が取った担当権だけを消せる。

気をつけた点

タイムアウトの値は、ハートビートの間隔より、十分に長くする。
間隔とタイムアウトが近いと、通信の遅延だけで、担当権を誤って奪ってしまうことがある。

担当権を奪った直後に、前の担当がまだ生きていて、処理を続けていることもある。
そうなると、2つの処理が同時に、結果を保存しようとしてしまう。

これを防ぐため、target_idだけでなく、lease_tokenも条件に含める。
lease_tokenが変わっていれば、それは誰かが先に担当権を奪ったということ。
その場合は、自分の結果を保存せずに、捨てる。

ただし、確認と保存を別々の処理にすると、確認した直後に奪われる隙間が残ってしまう。
確認と保存を1つのトランザクションにまとめ、確認の時点で行をロックすることで、この隙間をなくしている。

つまり、担当権の仕組みだけに頼らない。
ハートビート・保存・解放、すべてのタイミングでlease_tokenを確認し、二重に保存されないようにしている。

まとめ

外部APIへの問い合わせを含む、時間の読めない処理を、複数のリクエストが同時に始めてしまう問題があった。
advisory lockはDB接続を握り続ける前提の仕組みなので、この用途には向いていなかった。

DBの1行に時刻を書くだけの、ハートビート方式の担当権で解決した。
仕組みはシンプルだが、担当権を外す主体をPostgresからアプリ自身に移す、という発想の転換が肝になっている。