MNTSQ Techブログ

「MNTSQ(モンテスキュー)」のTechブログです。

Celery + SQS で、処理中タスクの Visibility Timeout を ハートビートで動的に延長する

MNTSQ Tech Blog TOP > 記事一覧 > Celery + SQS で、処理中タスクの Visibility Timeout を ハートビートで動的に延長する

はじめに

SREの寺島です。

Amazon SQS をタスクキューのブローカーとして使うとき、実行時間の短いジョブと長いジョブが混在していると、処理に時間のかかるタスクが途中で再配信されてしまう問題に当たります。可視性タイムアウト(Visibility Timeout)を長くすれば防げますが、長くするほど障害時の復旧が遅くなるというトレードオフがあります。

これを解決するのが、可視性タイムアウトを短く設定し、処理中にハートビートで動的に延長する方法です。本記事では Celery を使った実装例を紹介します。

SQS + Worker の非同期処理の構成

まず前提として、Celery + SQS でよくある非同期処理の構成を示します。

  1. クライアントからジョブのリクエストが送られる
  2. API がリクエストを受け付けてキュー(SQS)に enqueue し、クライアントに即座にAcceptedレスポンスを返す
  3. Worker が SQS に入ったメッセージを拾って処理を実行する

可視性タイムアウト設計の難しさ

可視性タイムアウトとは、Workerがメッセージを受信してから、そのメッセージが他のWorker から見えなくなる(再受信されなくなる)時間のことです。Worker はこの時間内に処理を終えてメッセージを削除します。削除されないままタイムアウトが切れると、メッセージは再び可視化され、別の Worker に再配信されます。

ここで問題になるのが、実行時間にばらつきのあるジョブです。例として、LLM を用いた処理を考えます。こうした処理は入力の大きさや複雑さ、呼び出し先 API のレイテンシによって実行時間が大きく変動し、多くは数秒〜10 秒程度で終わるものの、入力によっては数分、最悪で数十分かかることもあります。

可視性タイムアウトは最も遅い処理時間に合わせるのが定石ですが、長すぎると障害時の復旧が遅れます。この例で可視性タイムアウトを数十分に設定すると、Worker がクラッシュしてもメッセージは「処理中」として扱われ続けるため、その数十分間どの Worker も再処理しません。

逆に短すぎると、重複処理の無駄が発生します。処理が終わる前にメッセージが再可視化され、別の Worker が同じジョブを処理してしまうためです。

冪等性を担保すれば、重複処理されても最終的な状態は変わりません。ただし冪等性が保証するのは「結果が二重に反映されないこと」だけで、重い処理(外部API呼び出し、DB書き込み、計算など)が二重に走るという無駄までは消えません(実装で回避できるケースはあります)。特に LLM 呼び出しのようにコストのかかる処理では、重複処理が頻発するのは避けたいところです。

なお、実行時間の長短が別々の業務に対応しているのであれば、キューを分けてそれぞれに適切な可視性タイムアウトを設定するのが素直です。短時間ジョブ用・長時間ジョブ用にキューを分ければ、可視性タイムアウトの設計はキューごとに完結します。

しかし今回のように、同一のジョブの中で実行時間が大きく変動する場合は、同じキューに速いものと遅いものが混在するため、キュー分離では解決できません。

解決策

長すぎても短すぎても問題があるとなると、どちらかに固定するのは難しそうです。そこで最初は短く設定しておき、処理が続いている間だけ可視性タイムアウトを動的に延長する、というアプローチを取ります。

こうすれば、通常は短いタイムアウトのおかげで障害時に素早く再処理でき、長時間かかるジョブは延長によって再配信を防げます。固定値では避けられなかった長すぎ・短すぎのトレードオフの問題が解消されるわけです。

これは AWS のドキュメントでも紹介されている方法です。

タイムアウトの設定と調整。 まず、アプリケーションがメッセージを処理して削除するのに通常必要な最大時間に合わせて可視性タイムアウトを設定します。正確な処理時間について不明な場合は、短いタイムアウト (2 分など) で開始し、必要に応じて延長します。ハートビートメカニズムを実装して可視性タイムアウトを定期的に延長し、処理が完了するまでメッセージを非表示にします。これにより、未処理メッセージの再処理の遅れを最小限に抑えるとともに、再表示が早すぎないようにします。

https://docs.aws.amazon.com/ja_jp/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-visibility-timeout.html

具体的な動作

  • ベースの可視性タイムアウトを短めに設定します(例: 30 秒)。
  • Worker はバックグラウンドで定期的に(例: 15 秒ごと)ChangeMessageVisibility を呼び出し、可視性タイムアウトを延長し続けます。

なお、ベースの可視性タイムアウトはハートビートの延長間隔より長くしておく必要があります。延長間隔の方が長いと、最初のハートビートが届く前にベースの可視性タイムアウトが切れて再配信されてしまうためです(同様に、1 回あたりの延長幅も延長間隔より長くしておきます)。

長時間実行されるジョブでも、延長リクエストを出し続けている間は再配信されません。処理が完了したらハートビートを止め、メッセージを削除します。

sequenceDiagram
  participant W as 処理スレッド
  participant H as ハートビートスレッド
  participant Q as SQS

  Q->>W: メッセージ受信 (初期 30s)
  W->>H: ハートビート開始
  activate H
  activate W
  Note over W: 重い処理を実行中

  loop 15秒ごと
    H->>Q: ChangeMessageVisibility(30s)
    Note over Q: 期限を30sに上書き
  end

  W->>W: 処理完了
  deactivate W
  W->>H: stop.set() で停止
  deactivate H
  W->>Q: DeleteMessage
  Note over Q: メッセージ削除

一方、Worker がクラッシュした場合は延長リクエストが途絶えます。すると最後の延長から(最大でも延長幅の時間で)メッセージが再可視化され、別の Worker に再処理されます。短時間で終わるジョブがクラッシュした場合も同様に、ベースの可視性タイムアウト経過後に再処理されます。

sequenceDiagram
  participant W as 処理スレッド (Worker A)
  participant H as ハートビートスレッド
  participant Q as SQS
  participant W2 as 別ワーカー (Worker B)

  Q->>W: メッセージ受信 (初期 30s)
  W->>H: ハートビート開始
  activate H
  activate W
  H->>Q: ChangeMessageVisibility(30s)
  Note over Q: 期限 = 今から30s

  W->>W: クラッシュ / 強制kill
  deactivate W
  Note over H: スレッドも道連れで停止
  deactivate H

  Note over Q: 延長が来ない → 30s経過で期限切れ
  Q->>W2: 再配信(別ワーカーが受信)
  activate W2
  Note over W2: 処理を最初からやり直す<br/>※冪等性で二重実行を吸収
  deactivate W2

実装サンプル

Celery/kombu は SQS の可視性タイムアウトを自動延長しないため、ハートビートは自前で実装します。

ここでは要点を絞って簡略化しています。そのまま動かせる完全版は GitHub に置いてあるので、こちらを参照してください。

github.com

ReceiptHandle / Queue URL の取得

可視性タイムアウトを延長するには、対象メッセージの ReceiptHandle とキューの URL が必要です。kombu の SQS transport は、受信したメッセージの ReceiptHandle と Queue URL を task.request.delivery_info に載せてくれるので、タスク側から取得できます。

def extract_sqs_receipt(request):
    """Celery タスクの request から (ReceiptHandle, QueueURL) を取り出す。"""
    for src in (
        getattr(request, "delivery_info", None),
        (getattr(request, "properties", None) or {}).get("delivery_info"),
    ):
        if isinstance(src, dict):
            msg = src.get("sqs_message")
            if msg and msg.get("ReceiptHandle"):
                return msg["ReceiptHandle"], src.get("sqs_queue")
    return None, None

取得方法の出典: celery/celery Discussion #7388

ReceiptHandle さえ取れれば、あとは boto3 の change_message_visibility を定期的に叩くだけです。

ハートビート本体

再利用しやすいよう context manager にします。タスク実行中は daemon thread でハートビートを回し、処理本体はメインスレッドで実行、ハートビートスレッドが定期的に ChangeMessageVisibility を叩く構成です。

import contextlib
import threading

import boto3


@contextlib.contextmanager
def sqs_visibility_heartbeat(task, interval=15, extend_by=30):
    receipt_handle, queue_url = extract_sqs_receipt(task.request)
    if not (receipt_handle and queue_url):
        yield                      # SQS 以外 / handle 取得不可なら何もしない
        return

    sqs = boto3.client("sqs")
    stop = threading.Event()

    def _beat():
        # stop されるか interval 経過のたびにループ
        while not stop.wait(interval):
            try:
                sqs.change_message_visibility(
                    QueueUrl=queue_url,
                    ReceiptHandle=receipt_handle,
                    VisibilityTimeout=extend_by,   # 呼んだ時点から extend_by 秒に設定し直す
                )
            except Exception:
                break              # 期限切れ等は延長を諦めて停止

    t = threading.Thread(target=_beat, daemon=True)
    t.start()
    try:
        yield
    finally:
        stop.set()                 # 必ずハートビートを止める
        t.join(timeout=2)

クラッシュ/強制終了時はこの daemon thread もプロセスごと止まるため延長が途絶え、ベースの可視性タイムアウト経過後に SQS が自動で再配信してくれます。

タスク側

タスク本体は with で囲むだけです。

@celery_app.task(bind=True, name="long_task")
def long_task(self, seconds):
    with sqs_visibility_heartbeat(self):
        do_heavy_work()

なお、この方式が機能する前提として task_acks_late = True が必要です。これが無いと kombu は受信直後にメッセージを削除してしまい、延長対象が消える上にクラッシュ時も再配信されなくなります。

注意点

本実装のハートビートはスレッドで動くため、Python の GIL の影響を受けます。とはいえ、I/O 待ちや通常の Python 処理では GIL は定期的に解放されるため、ほとんどのケースでは問題になりません。注意が必要なのは、GIL を解放しない C 拡張が 1 回の呼び出しで長時間 GIL を握り続けるようなケースです。この間はハートビートスレッドが動けず延長が止まり、可視性タイムアウト経過後に再配信されてしまう可能性があります。そのようなワークロードでは、別プロセスでの延長など別の方式を検討する必要があります。

動作確認

ローカルの SQS 互換サーバ(ElasticMQ)で動作を確認できます。確認しやすいように base は短め(10 秒)にし、ハートビートは 4 秒ごと・延長幅 10 秒で動かします。

以下の手順は前掲の GitHub リポジトリ をクローンすればそのまま実行できます。まず ElasticMQ を起動し、各ターミナルでは下記の環境変数を設定しておきます。

docker compose up -d   # ElasticMQ(jobs キュー: base 10 秒)

# 以降の各ターミナルで設定
export SQS_ENDPOINT_URL=http://localhost:9324 \
  AWS_ACCESS_KEY_ID=x AWS_SECRET_ACCESS_KEY=x AWS_REGION=us-east-1 \
  HEARTBEAT_INTERVAL=4 HEARTBEAT_EXTEND_BY=10

通常時:長時間タスクでも再配信されない

Worker を起動し、base を超える 25 秒のタスクを投入します。

# Worker(別ターミナル)
uv run celery -A celery_app worker -Q jobs --concurrency=4 --loglevel=INFO

# 投入(別ターミナル)
uv run python enqueue.py --seconds 25

すると Worker のログは次のようになります。

long_task start: 25s (task=01cb2cbb-...)
heartbeat #1: +10s (task=01cb2cbb-...)
heartbeat #2: +10s (task=01cb2cbb-...)
...
heartbeat #6: +10s (task=01cb2cbb-...)
long_task done (task=01cb2cbb-...)

base(10 秒) を超える 25 秒タスクでも、4 秒ごとのハートビートで延長され続けるため再配信されず、1 回だけで完走しますlong_task start が 1 回しか出ないことで確認できます)。

クラッシュ時:Worker A が落ちると Worker B が引き継ぐ

図のように「処理中の Worker が落ちて、別の Worker が引き継ぐ」様子を再現します。Worker を 2 つ(A・B)起動し、長めのタスクを投入してから、処理中の Worker を強制終了します。

# ターミナル1: Worker A
uv run celery -A celery_app worker -Q jobs -n workerA@%h --concurrency=1 --loglevel=INFO

# ターミナル2: Worker B
uv run celery -A celery_app worker -Q jobs -n workerB@%h --concurrency=1 --loglevel=INFO

# ターミナル3: 60 秒のタスクを投入
uv run python enqueue.py --seconds 60

タスクはどちらか一方(仮に Worker A)が拾い、A 側に long_task startheartbeat #N が出ます。処理中に A を強制終了します。

pkill -9 -f workerA

A が止まるとハートビートが途絶えるため、base(10 秒) 経過後にメッセージが再可視化されます。するとまだ生きている Worker B がそのメッセージを拾い、最初から処理をやり直します(B 側に long_task start が出ます)。タスクの長さ(60 秒)を待つことなく、base 程度の時間で別の Worker に引き継がれることが確認できます。

最後に

実行時間にばらつきのある非同期処理ワークロードに対して、可視性タイムアウトの動的延長は非常に有効です。ジョブの実行時間が予測できない場合や、データ量の増加によってジョブ時間が伸びる可能性がある場合にも役立つと思います。

Celery + SQS の構成では可視性タイムアウトの動的延長がサポートされておらず、自前で実装する必要がありました。対応内容はシンプルですが、Webで実装例をなかなか見つけられなかったので、本記事でサンプル実装を紹介しました。同じ悩みを抱えている方の参考になれば幸いです。