Zenn
👌

FastAPI: プロセスを起動し、その出力を WebSocket でブラウザに流す

2023/08/18に公開
Python
FastAPI
#
WebSocket
#
asyncio
tech

FastAPI で、

  1. 何かリクエストを受けたら、
  2. プロセスを起動して、
  3. そのプロセスの出力を WebSocket で流して
  4. ブラウザに表示する

というのをやってみました。

FastAPI と、起動したプロセスの間の通信は標準入出力でやりとりします。

実装

import asyncio
import atexit

from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from websockets.exceptions import ConnectionClosed

app = FastAPI()
background_tasks = set()


html = """
<!DOCTYPE html>
<html>
    <head>
        <style>
            pre { white-space: pre-wrap; }
        </style>
    </head>
    <body>
        <div>
            <p>tail -f</p>
            <pre id="log"></pre>
        </div>
        <script>
            var ws = new WebSocket("ws://localhost:8000/ws");
            ws.onmessage = function(event) {
                const log_element = document.getElementById('log')
                log_element.insertAdjacentText('beforeend', event.data)
            };
        </script>
    </body>
</html>
"""


@app.get("/", response_class=HTMLResponse)
def root():
    return HTMLResponse(content=html)


async def tail_f(filename: str) -> tuple[asyncio.Task, asyncio.StreamReader]:
    """
    tail -f をバックグランドで実行する

    usage:
        tail_task, reader = await tail_f("path/to/somefile")
        await some_coroutine()
        tail_task.cancel()  # tail -f の終了
    """
    proc = await asyncio.create_subprocess_exec(
        "tail", "-f", filename, stdout=asyncio.subprocess.PIPE
    )
    task = asyncio.create_task(proc.wait())
    task.add_done_callback(lambda _: proc.terminate())

    background_tasks.add(task)
    task.add_done_callback(background_tasks.discard)

    return task, proc.stdout


# 終了時に実行中のプロセスも終了させる
def exit_handler():
    for task in background_tasks:
        task.cancel()


atexit.register(exit_handler)


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()

    tail_task, reader = await tail_f("log.txt")
    try:
        async for line in reader:
            await websocket.send_bytes(line.decode("utf-8"))
    except ConnectionClosed:
        tail_task.cancel()
  • / にアクセスすると WebSocket でサーバにアクセスする HTML を返します（サンプルなので js は埋め込み）
  • /ws が WebSocket のエンドポイントで、ここではサンプルとして tail -f を起動しています
  • python終了時に tail -f を終了したいので、タスクを background_tasks に保存してゴニョゴニョしています

動作確認

  • サーバの起動
    python3 -m venv venv
source venv/bin/activate
pip install 'fastapi[all]'

touch log.txt
uvicorn main:app --reload
  • ブラウザから http://localhost:8000 を開く
  • log.txt にデータを書き込と、ブラウザに反映されるはず
    echo hoge >> log.txt
echo fuga >> log.txt

Discussion