Zenn
🕌

FastAPIでタイムアウト処理を入れる

2021/12/12に公開
Python
FastAPI
tech

Sansan Advent Calendar 2021 12日目の記事です。

FastAPI は Python のWebフレームワークです。ECS Fargate 上で FastAPI アプリケーションを動かしていたときに、たまたま長時間処理し続けていた場合に意図せぬ挙動をしていたので調べたり対策してみました。

事象

ALB + ECS Fargate + FastAPI（Gunicorn + Uvicorn）で構成していました。ALB の Idle Timeout は60秒に設定しているので、60秒かかったリクエストには Gateway Timeout（504）を返します。

しかし、リクエストログを見ると ALB 側で Gateway Timeout を返した後も FastAPI アプリケーション内部は処理し続けていることがわかりました。Gunicorn のタイムアウトはデフォルトで30秒に設定されていることから正常に動作しないことがわかりました。

https://docs.gunicorn.org/en/stable/settings.html#timeout

検証

Sleepを入れただけの簡単な FastAPI アプリケーションを実行して確認してみました。

import time
import asyncio
from fastapi import FastAPI

app = FastAPI()

@app.get("/sleep/{wait_time}")
def wait_for(wait_time: int):
    time.sleep(wait_time)
    return {"Wait time(sec)": wait_time}

@app.get("/async_sleep/{wait_time}")
async def async_wait_for(wait_time: int):
    await asyncio.sleep(wait_time)
    return {"Wait time(sec)": wait_time}

/sleep/300 のようにある程度長めに Sleep をいれると、クライアント側でリクエストを打ち切ったとしても FastAPI 側は止まらず待ち続けていることがわかりました。

[2021-11-26 13:16:47 +0900] [15903] [INFO] Starting gunicorn 20.1.0
[2021-11-26 13:16:47 +0900] [15903] [INFO] Listening at: http://0.0.0.0:8000 (15903)
[2021-11-26 13:16:47 +0900] [15903] [INFO] Using worker: uvicorn.workers.UvicornWorker
[2021-11-26 13:16:47 +0900] [15948] [INFO] Booting worker with pid: 15948
[2021-11-26 13:16:47 +0900] [15949] [INFO] Booting worker with pid: 15949
[2021-11-26 13:16:47 +0900] [15950] [INFO] Booting worker with pid: 15950
[2021-11-26 13:16:47 +0900] [15948] [INFO] Started server process [15948]
[2021-11-26 13:16:47 +0900] [15948] [INFO] Waiting for application startup.
[2021-11-26 13:16:47 +0900] [15948] [INFO] Application startup complete.
[2021-11-26 13:16:47 +0900] [15949] [INFO] Started server process [15949]
[2021-11-26 13:16:47 +0900] [15949] [INFO] Waiting for application startup.
[2021-11-26 13:16:47 +0900] [15949] [INFO] Application startup complete.
[2021-11-26 13:16:47 +0900] [15950] [INFO] Started server process [15950]
[2021-11-26 13:16:47 +0900] [15950] [INFO] Waiting for application startup.
[2021-11-26 13:16:47 +0900] [15950] [INFO] Application startup complete.
# ここでSleepが入る

# クライアント側でリクエストを打ち切るがSleepしたまま

# ここでgunicornを切る
[2021-11-26 13:17:19 +0900] [15903] [INFO] Handling signal: int
[2021-11-26 13:17:19 +0900] [15903] [WARNING] Worker with pid 15948 was terminated due to signal 3
[2021-11-26 13:17:19 +0900] [15903] [WARNING] Worker with pid 15949 was terminated due to signal 3
[2021-11-26 13:17:19 +0900] [15903] [WARNING] Worker with pid 15950 was terminated due to signal 3
[2021-11-26 13:17:19 +0900] [15903] [INFO] Shutting down: Master

対策

クライアント側でタイムアウトを設定したとしても、FastAPI 側で打ち切らないことがわかったので、FastAPI 内部でタイムアウト処理を入れる必要があります。

調べたところ、以下の Issue を発見しました。FastAPI の Middleware という機能を利用して asyncio.wait_for() でタイムアウトを設定するとできそうです。

[QUESTION] How does one setup a global timeout to all requests? · Issue #1752 · tiangolo/fastapi

実装

フルバージョンは以下にあげています。

Add timeout middleware by ohsawa0515 · Pull Request #3 · ohsawa0515/fastapi_logging_demo

Middleware ではなく、Custom Router で実装するバージョンもあります。

[not merge] Add timeout by ohsawa0515 · Pull Request #4 · ohsawa0515/fastapi_logging_demo

# main.py
import os
import time
import asyncio
from fastapi import FastAPI
from timeout_middleware import TimeoutMiddleware

# タイムアウト値を環境変数で設定できるように
REQUEST_TIMEOUT = os.getenv("REQUEST_TIMEOUT", 5)

app = FastAPI()
app.add_middleware(TimeoutMiddleware, timeout=REQUEST_TIMEOUT)

@app.get("/sleep/{wait_time}")
def wait_for(wait_time: int):
    time.sleep(wait_time)
    return {"Wait time(sec)": wait_time}

@app.get("/async_sleep/{wait_time}")
async def async_wait_for(wait_time: int):
    await asyncio.sleep(wait_time)
    return {"Wait time(sec)": wait_time}

# timeout_middleware.py
import asyncio
from typing import Callable
from fastapi import Request, Response
from fastapi.responses import JSONResponse
from starlette.types import ASGIApp
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.status import HTTP_504_GATEWAY_TIMEOUT

class TimeoutMiddleware(BaseHTTPMiddleware):
    """
    Return gateway timeout error (504)
    if the request processing time is above a certain threshold
    """

    def __init__(self, app: ASGIApp, timeout: int = 10):
        super().__init__(app)
        self.timeout = int(timeout)

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        try:
            return await asyncio.wait_for(call_next(request), timeout=self.timeout)
        except asyncio.TimeoutError:
            return JSONResponse(
                content={'detail': 'Request processing time excedeed limit'}, status_code=HTTP_504_GATEWAY_TIMEOUT)

Discussion