Amazon S3で分散ロックを実装する

先日Amazon S3でconditional write機能がリリースされました。本記事では、この機能を用いた分散ロックについて検討します。

aws.amazon.com

分散ロックとは

分散ロック (distributed lock) とは、分散環境で排他制御を実現するために必要な機構です。実現できることはロックですが、分散環境から利用できることが特徴です *1。

実装はRedisを利用したものが有名ですが、AWSネイティブな実装としてはDynamoDBを利用することも多いでしょう。(実装例: DynamoDBLockClient, Powertools for Lambda)

分散ロックは強い整合性を持つ条件付き書き込みが可能なストレージがあれば、実現することが出来ます。

// 分散ロックの擬似コード
結果 = 条件を満たしたら書き込み(共通のキー)
if (結果 == 成功) {
  // ロックが取得できたのでメインの処理を実行
  メイン処理

  // メイン処理が終わったらロックを解放する
  ロックの解放
} else {
  // ロックを取得できなかった。再試行や終了などする
}

S3のconditional writeも強い整合性を持つため、分散ロックを実装できます。「条件を満たしたら」の条件は、「同じキーのオブジェクトが存在しなければ」という条件になります。

AWS SDK for JavaScriptによる実装

それでは、S3による分散ロック実装例をTypeScriptで見てみましょう。以下は100個のタスクがロックを取り合う例です:

import { S3 } from '@aws-sdk/client-s3';
import { setTimeout } from 'timers/promises';

const s3 = new S3();
const key = '.lock';
const bucket = process.env.BUCKET;

const task = async (id: number) => {
  while (true) {
    // 各タスクでタイミングをバラつかせる
    await setTimeout(Math.random() * 500 + 500);
    try {
      // ロックの取得を試みる
      await s3.putObject({
        Bucket: bucket,
        Key: key,
        IfNoneMatch: '*',
        Body: '\n',
      });
    } catch (e) {
      // ロックの取得に失敗。再試行する
      continue;
    }

    // ロック取得に成功
    console.log(`acquired lock ${id}`);
    // メイン処理 (ここでは仮にsleepするだけ)
    await setTimeout(2000);

    // ロックを解放する
    console.log(`releasing lock ${id}`);
    await s3.deleteObject({
      Bucket: bucket,
      Key: key,
    });
  }
};

// 上記タスクを100個起動
new Array(100).fill(0).forEach((_, i) => {
  task(i);
});

すべてのタスクは同じオブジェクトキー (ここでは .lock) をロックのオブジェクトとして利用します。これにより、全体でひとつのロックを取り合う形となります。

putObject に IfNoneMatch: '*' を指定することで、オブジェクトが存在しない場合は作成、存在すればエラーとなります。強い整合性を持つ書き込みのため、同時にリクエストが発生した場合、ただ1つのリクエストだけが成功することが保証されています。

ロックを取得できたタスクは .lock という空オブジェクトをS3バケット上に作成し、メイン処理を実行後、そのオブジェクトをバケットから削除してロックを解放します。

実行すると

実際に実行すると、各タスクがロックを取り合いつつ、排他制御ができている様子が観察できます。

acquired lock 3
releasing lock 3
acquired lock 8
releasing lock 8
acquired lock 65
releasing lock 65
acquired lock 54
releasing lock 54
acquired lock 38
releasing lock 38
acquired lock 77
releasing lock 77
...

ちなみに、ロックの取得に失敗した場合は下記のエラーが得られるようです:

PreconditionFailed: At least one of the pre-conditions you specified did not hold
...
{
  '$fault': 'client',
  '$metadata': {
    httpStatusCode: 412,
    requestId: 'REDACTED',
    extendedRequestId: 'REDACTED,
    cfId: undefined,
    attempts: 1,
    totalRetryDelay: 0
  },
  Code: 'PreconditionFailed',
  Condition: 'If-None-Match',
  RequestId: 'REDACTED',
  HostId: 'REDACTED'
}

この例では各タスクがすべて同じプロセス内にいるため分散ロックの必要すらないわけですが、概観をつかむことはできますね。

実用性を考える

ここまでで、S3を利用した分散ロックを実装できることがわかりました。追加の観点から、実用性を考えてみましょう。なお私は分散処理の専門家ではないため、間違っていたら教えて下さい🙇

ロックの期限は？

多くの分散ロックの実装では、ロックに期限 (expiry) を設定できます。これにより、ロックを取得した処理が何らかの原因でロックの解放に失敗したときも、設定した期限以降は再び他の処理がロックを取得することができます。

例えばDynamoDBでは、conditional writeの条件に不等号などを利用できるため、ロックの期限を実装可能です。

S3の場合、書き込み時に利用できる条件は現状「オブジェクトがすでに存在するかどうか」のみのため、単純に期限を実装するのは難しそうです。

一案として、S3のライフサイクルルールを使えば、「オブジェクトが作成されてからN日後にオブジェクトを自動で削除する」ことができます。オブジェクトの削除はロックの解放と同義のため、これを使えばロックの期限を実装できると思われます。また、ロックを取得した処理が定期的にPutObjectし直すことで、ハートビートも実装できそうです。しかしながら、期限の設定単位は日毎になる (最短でも1日後) になるので、ユースケースは限られてくるでしょう。

あるいはロックのオブジェクトを削除するワーカーを別途用意し、「オブジェクトの作成日時を見て、期限を超えていたら削除する」という方法も可能かもしれません (要はライフサイクルルールのセルフ実装)。ただし、DeleteObjectのAPI は条件付きの削除などは現状できないため、削除とハートビートのリクエストの競合を完全に回避することは難しいでしょう。実用上は、期限よりも十分短い間隔でハートビートすれば問題にはなりづらいと思われます。(これはライフサイクルルールを使う場合もそう)

さらに別解として、AWS Step Functions (SFn)を使う方法も考えられます。ロック取得・メイン処理・ロック解放を別々のタスクとしてもつSFnステートマシンとして実装し、SFnが正常に動作している限りは必ずロックが解放される前提を置く (期限に頼らない) という方法です。万が一S3やSFnの障害などでロックが解放されなかったときは、手作業などで復旧を行います。

少なくともDynamoDBやRedisよりは期限に関する実装の選択肢が減るので、ここは重要な考慮点となりそうですね。

コストは？

気になるコストも確認しましょう。ドキュメントを読む限りはconditional writeでコストが変わるわけでもないようなので、通常のPUTと同一コストがかかるとします。この場合は、1000リクエストあたり0.005USD です (us-east-1)。ロック取得失敗したときのコストも同様に掛かるようです (こちらのドキュメントに課金されないエラーコードがまとめられていますが、conditional writeの失敗は含まれないように読めます。)

DynamoDBと比べてどうでしょうか？オンデマンドキャパシティの料金と比べると、強整合の書き込みは2つのWRUを使うので 1000リクエストあたり0.0025USD です (取得失敗時も同様)。S3のちょうど半額となります。

コスト面ではDynamoDBが有利ですが、そもそもが安いので、リクエスト量次第でいずれにせよ許容できるコストに収まることもあるでしょう。