beautifulsoups.dev
Published on

async contextmanager로 커넥션 누수 막는 법

Authors

비동기 서버를 운영하다 보면 어느 순간부터 DB 풀의 available=0 상태가 길어지고, 지연이 튀거나 타임아웃이 연쇄적으로 발생하는 경우가 있습니다. 원인이 쿼리 성능이 아니라 커넥션 누수(connection leak) 인 경우가 꽤 많습니다. 특히 async 코드에서는 예외, return, break, 그리고 CancelledError 같은 취소 경로가 다양해서 “반납을 깜빡한” 코드가 숨어들기 쉽습니다.

이 글에서는 파이썬의 contextlib.asynccontextmanager를 사용해 커넥션 획득과 반납을 한 군데로 모으고, 예외·취소 상황에서도 항상 정리되도록 만드는 실전 패턴을 정리합니다.

운영 관점에서 비슷한 유형의 장애 대응 글로는 K8s CrashLoopBackOff·OOMKilled 원인과 해결, DB 유지보수 측면에서는 PostgreSQL VACUUM 안 하면 폭증하는 3가지 증상도 함께 참고하면 좋습니다.

비동기에서 커넥션 누수가 더 자주 터지는 이유

동기 코드에서는 보통 try/finally로 닫는 습관이 잘 자리잡아 있습니다. 하지만 비동기에서는 다음이 겹치면서 누수 확률이 올라갑니다.

  • 조기 반환 경로가 많다: 조건 검사 후 return이 여러 군데 존재
  • 예외가 더 흔하다: 네트워크 I/O, 타임아웃, 재시도 로직
  • 취소가 개입한다: 요청 타임아웃, 서버 셧다운, 상위 태스크 취소로 asyncio.CancelledError가 발생
  • 풀 기반 자원: 닫지 않으면 “프로세스가 끝날 때까지”가 아니라 “풀이 고갈될 때까지” 버틴 뒤 장애로 번진다

핵심은 “닫는다”를 각 함수에서 매번 신경쓰게 만들면 언젠가 새는 지점이 생긴다는 점입니다. 그래서 획득/반납을 구조로 강제하는 것이 가장 안전합니다.

안티패턴: await pool.acquire() 후 반납 누락

다음 코드는 얼핏 정상처럼 보이지만, 예외나 조기 반환에서 커넥션이 반납되지 않을 수 있습니다.

import asyncpg

async def get_user(pool: asyncpg.Pool, user_id: int):
    conn = await pool.acquire()

    # 조건에 따라 조기 반환
    if user_id <= 0:
        return None  # conn 반납 누락

    row = await conn.fetchrow(
        "select id, email from users where id=$1",
        user_id,
    )

    # 예외가 여기서 터지면 반납 누락 가능
    if row is None:
        raise ValueError("user not found")

    await pool.release(conn)
    return dict(row)

물론 try/finally로 감싸면 해결되지만, 이런 코드가 프로젝트 곳곳에 흩어지면 유지보수 중에 다시 누락되기 쉽습니다.

기본 해법: asynccontextmanager로 획득/반납을 표준화

contextlib.asynccontextmanager를 쓰면 async with 블록이 끝날 때 항상 정리 코드를 실행하도록 강제할 수 있습니다.

1) 풀 커넥션을 안전하게 빌려주는 컨텍스트 매니저

from contextlib import asynccontextmanager
import asyncpg

@asynccontextmanager
async def acquire_conn(pool: asyncpg.Pool):
    conn = await pool.acquire()
    try:
        yield conn
    finally:
        await pool.release(conn)

사용 측 코드는 다음처럼 단순해집니다.

async def get_user(pool: asyncpg.Pool, user_id: int):
    async with acquire_conn(pool) as conn:
        if user_id <= 0:
            return None

        row = await conn.fetchrow(
            "select id, email from users where id=$1",
            user_id,
        )
        if row is None:
            raise ValueError("user not found")

        return dict(row)

이제 return이 어디에 있든, 예외가 어디서 나든 finally에서 반납이 보장됩니다.

2) 트랜잭션까지 함께 묶기

실무에서는 “커넥션 반납”만큼 “트랜잭션 종료”도 누락되기 쉽습니다. asyncpg는 conn.transaction() 컨텍스트를 제공하지만, 프로젝트 표준으로 묶어두면 더 안전합니다.

from contextlib import asynccontextmanager
import asyncpg

@asynccontextmanager
async def transaction(pool: asyncpg.Pool):
    async with acquire_conn(pool) as conn:
        tr = conn.transaction()
        await tr.start()
        try:
            yield conn
        except Exception:
            await tr.rollback()
            raise
        else:
            await tr.commit()

사용 예:

async def transfer(pool: asyncpg.Pool, from_id: int, to_id: int, amount: int):
    async with transaction(pool) as conn:
        await conn.execute(
            "update accounts set balance = balance - $1 where id = $2",
            amount,
            from_id,
        )
        await conn.execute(
            "update accounts set balance = balance + $1 where id = $2",
            amount,
            to_id,
        )

이 패턴은 DDD에서 애그리거트 경계가 흔들리며 트랜잭션 범위가 비대해질 때 특히 중요합니다. 트랜잭션이 길어질수록 커넥션 점유 시간이 늘고, 누수가 터졌을 때 영향 범위가 커집니다. 관련해서는 DDD 애그리거트 경계 오류로 트랜잭션 폭발할 때도 같이 보면 좋습니다.

취소(CancelledError)에서도 안전한가

async with는 블록 탈출 시점에 __aexit__가 호출되므로, 일반적으로 취소가 걸려도 finally가 실행됩니다. 다만 실무에서는 다음을 주의해야 합니다.

  • finally 안에서 또 await를 하기 때문에, 이벤트 루프 상태에 따라 취소가 전파될 수 있습니다.
  • 파이썬은 취소가 들어오면 다음 await에서 CancelledError가 발생할 수 있습니다.

그래서 “반납/정리” 같은 중요한 작업은 finally에서 실행하되, 필요하면 asyncio.shield로 감싸 취소 전파를 막는 선택지도 있습니다.

import asyncio
from contextlib import asynccontextmanager

@asynccontextmanager
async def acquire_conn_shielded(pool):
    conn = await pool.acquire()
    try:
        yield conn
    finally:
        # 취소가 와도 반납은 끝까지 시도
        await asyncio.shield(pool.release(conn))

주의: shield는 만능이 아닙니다. 서버가 강제 종료되거나 프로세스가 죽으면 어떤 정리도 보장할 수 없습니다. 하지만 “요청 취소” 수준에서는 누수 방지에 도움이 됩니다.

HTTP 클라이언트도 똑같이 새기 쉽다

DB뿐 아니라 aiohttp.ClientSession 같은 HTTP 세션도 누수의 단골입니다. 세션을 매 요청마다 만들거나, 예외 시 close를 놓치면 소켓이 쌓입니다.

세션을 컨텍스트로 감싸기

from contextlib import asynccontextmanager
import aiohttp

@asynccontextmanager
async def http_session(timeout_s: float = 10.0):
    timeout = aiohttp.ClientTimeout(total=timeout_s)
    session = aiohttp.ClientSession(timeout=timeout)
    try:
        yield session
    finally:
        await session.close()

async def fetch_json(url: str):
    async with http_session() as session:
        async with session.get(url) as resp:
            resp.raise_for_status()
            return await resp.json()

핵심은 동일합니다. “열었으면 닫아야 한다”를 개발자 기억력에 맡기지 말고, 컨텍스트로 강제합니다.

FastAPI 의존성(Dependency)과 함께 쓰는 패턴

웹 API에서는 요청 단위로 커넥션을 빌려주고, 핸들러가 끝나면 반납하는 구조가 가장 자연스럽습니다. FastAPI라면 yield 기반 의존성을 쓰는 것이 사실상 컨텍스트 매니저와 같은 개념입니다.

from contextlib import asynccontextmanager
from fastapi import Depends, FastAPI
import asyncpg

app = FastAPI()

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.pool = await asyncpg.create_pool(dsn="postgresql://...")
    try:
        yield
    finally:
        await app.state.pool.close()

app.router.lifespan_context = lifespan

async def get_conn():
    pool = app.state.pool
    async with acquire_conn(pool) as conn:
        yield conn

@app.get("/users/{user_id}")
async def read_user(user_id: int, conn=Depends(get_conn)):
    row = await conn.fetchrow("select id, email from users where id=$1", user_id)
    return dict(row) if row else None

이렇게 하면 요청 처리 중 예외가 나도 커넥션이 반납됩니다.

누수가 의심될 때 운영에서 확인할 것

async contextmanager로 구조를 잡아도, 이미 누수가 있거나 풀 설정이 부적절하면 문제가 계속 보일 수 있습니다. 다음을 함께 점검하세요.

  • 풀 메트릭: 사용 중 커넥션 수, 대기열 길이, 획득 대기 시간
  • 타임아웃: 풀 획득 타임아웃, 쿼리 타임아웃
  • 트랜잭션 체류 시간: 긴 트랜잭션이 커넥션을 오래 점유하는지
  • DB 상태: autovacuum 지연, bloat로 인한 쿼리 지연이 커넥션 점유를 늘리는지

특히 PostgreSQL은 VACUUM이 밀리면 쿼리가 느려지고, 그 결과 커넥션 점유 시간이 증가해 “누수처럼 보이는” 풀 고갈이 발생할 수 있습니다. 관련 증상은 PostgreSQL VACUUM 안 하면 폭증하는 3가지 증상에서 체크리스트로 정리되어 있습니다.

팀 규칙으로 굳히는 팁

  • pool.acquire() 직접 호출을 금지하고, 반드시 acquire_conn() 또는 transaction()만 사용하게 한다
  • 코드리뷰 체크리스트에 “자원 획득 후 반납 경로가 단일화되었는가”를 추가한다
  • 테스트에서 의도적으로 예외/취소를 발생시켜도 풀의 in_use가 원복되는지 확인한다

간단한 예로, 취소를 시뮬레이션하는 테스트는 다음처럼 작성할 수 있습니다.

import asyncio
import pytest

@pytest.mark.asyncio
async def test_cancel_does_not_leak(pool):
    async def task_body():
        async with acquire_conn_shielded(pool) as conn:
            await conn.execute("select pg_sleep(5)")

    t = asyncio.create_task(task_body())
    await asyncio.sleep(0.1)
    t.cancel()

    with pytest.raises(asyncio.CancelledError):
        await t

    # 여기서 풀 상태/메트릭을 확인하거나, 추가 쿼리가 즉시 수행되는지 확인
    async with acquire_conn(pool) as conn:
        await conn.execute("select 1")

정리

async contextmanager의 가치는 문법 편의가 아니라 누수를 불가능하게 만드는 구조화에 있습니다.

  • 커넥션 획득/반납을 async with로 강제
  • 트랜잭션 시작/커밋/롤백을 한 곳에 모아 실수 제거
  • 취소까지 고려하면 asyncio.shield 같은 방어막을 선택적으로 적용

비동기 시스템에서 커넥션 누수는 “가끔 발생하는 버그”가 아니라 “언젠가 장애로 바뀌는 부채”에 가깝습니다. 초기에 컨텍스트 매니저로 패턴을 표준화해두면, 운영 단계에서 풀 고갈과 연쇄 타임아웃을 겪을 확률이 크게 줄어듭니다.

Discuss on Twitter • View on GitHub