Python 비동기 작업에서 메모리 누수 추적하기

문제 상황

운영 중인 FastAPI 서비스가 24시간 이상 구동되면 메모리 사용량이 4GB를 넘어서는 현상이 발생했다. 초기 메모리는 500MB 정도였는데, 시간이 지날수록 선형적으로 증가했다.

원인 추적

tracemalloc을 활용해 메모리 스냅샷을 비교했다.

import tracemalloc
import asyncio

tracemalloc.start()

# 작업 실행 후
snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics('lineno')

for stat in top_stats[:10]:
    print(stat)

결과를 보니 WebSocket 연결 처리 코드에서 생성한 asyncio 태스크들이 완료되지 않은 채 누적되고 있었다. asyncio.create_task()로 생성한 태스크의 참조를 제대로 관리하지 않아서 발생한 문제였다.

해결 방법

태스크 생성 시 weakref를 사용한 컨테이너로 관리하도록 수정했다.

import weakref
from typing import Set

class TaskManager:
    def __init__(self):
        self._tasks: Set[asyncio.Task] = weakref.WeakSet()
    
    def create_task(self, coro):
        task = asyncio.create_task(coro)
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)
        return task
    
    async def cancel_all(self):
        tasks = list(self._tasks)
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)

태스크가 완료되면 자동으로 컨테이너에서 제거되고, 서비스 종료 시 남은 태스크를 명시적으로 취소할 수 있게 되었다.

결과

패치 배포 후 72시간 동안 모니터링한 결과, 메모리 사용량이 600MB 내외로 안정화되었다. 비동기 프로그래밍에서는 태스크 생명주기 관리가 중요하다는 것을 다시 한번 느꼈다.