Published on

contextlib로 async 컨텍스트 누수·예외 잡기

Authors

서버에서 async with 블록을 쓴다고 해서 리소스가 자동으로 안전해지는 건 아닙니다. 특히 데이터베이스 커넥션, HTTP 세션, 스트리밍 응답, 임시 파일 같은 자원은 예외/취소가 끼어드는 순간 정리 코드가 건너뛰어지기 쉽고, 그 결과가 “가끔 커넥션이 안 돌아옴”, “파일 디스크립터가 쌓임”, “백그라운드 태스크가 계속 살아있음” 같은 형태로 나타납니다.

이 글은 contextlib를 중심으로 async Context Manager의 예외·누수 패턴을 재현하고, 이를 구조적으로 봉인하는 방법을 다룹니다. 핵심은 다음 세 가지입니다.

  • @contextlib.asynccontextmanager로 “정리 로직이 반드시 실행되는” 형태를 강제하기
  • 여러 자원을 동적으로 쌓을 때 contextlib.AsyncExitStack으로 누수 방지하기
  • 비동기 제너레이터/스트리밍 자원에 contextlib.aclosing을 적용해 종료 보장하기

운영 환경에서 “취소”는 예외의 한 종류로 취급해야 합니다. 이 관점은 취소/역압으로 버그가 터지는 케이스(예: 스트림 결합, 파이프라인 중단)와도 연결됩니다. 비슷한 맥락의 문제로는 Kotlin Flow combine 최신 문법과 역압·취소 버그도 참고할 만합니다.

왜 async Context Manager에서 누수가 생기나

1) __aenter__ 성공 후, __aexit__가 호출되지 않는 케이스는 거의 없다

정상적으로 async with에 들어갔다면, 블록을 빠져나올 때 __aexit__는 호출됩니다. 문제는 보통 다음에서 생깁니다.

  • 자원을 컨텍스트 밖에서 열어놓고 async with로 감싸지 않음
  • 여러 자원을 조건에 따라 열다가 중간 예외가 나서, 일부만 정리됨
  • 비동기 제너레이터/스트리밍을 “끝까지 소비하지 않고” 중간에 탈출함
  • 취소(asyncio.CancelledError)가 들어왔을 때 정리 순서가 꼬이거나, 실수로 취소를 삼켜서 태스크가 이상 상태가 됨

2) “부분 성공”이 가장 위험하다

예를 들어 커넥션을 하나 열고, 그 다음 임시 파일을 만들고, 그 다음 원격 스트림을 연결하는데 두 번째 단계에서 예외가 나면 첫 번째 커넥션은 닫혀야 합니다. 이런 “단계적 자원 획득”은 사람이 try/finally로 관리하면 빠뜨리기 쉽습니다.

이때 contextlib의 스택 기반 정리 도구가 빛을 봅니다.

@asynccontextmanager로 예외/취소에도 정리 보장하기

contextlib.asynccontextmanageryield 앞을 “획득”, yield 뒤를 “정리”로 강제하는 패턴입니다. 코드 리뷰 때도 의도가 명확해지고, try/finally를 표준화할 수 있습니다.

예시: HTTP 세션을 안전하게 제공하는 컨텍스트

import contextlib
import aiohttp

@contextlib.asynccontextmanager
async def http_session():
    session = aiohttp.ClientSession()
    try:
        yield session
    finally:
        await session.close()

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

포인트는 finally입니다.

  • 블록 내부 예외
  • 타임아웃 예외
  • 취소(CancelledError)

모두 finally를 통과하며 close가 호출됩니다.

예외를 로깅하되, 삼키지 않는 패턴

정리 단계에서 예외 정보를 로깅하고 싶을 때가 많습니다. 하지만 주의할 점은 원래 예외를 덮어쓰지 않는 것입니다.

import contextlib
import logging

log = logging.getLogger(__name__)

@contextlib.asynccontextmanager
async def traced(name: str):
    try:
        yield
    except Exception:
        log.exception("context %s failed", name)
        raise
    finally:
        log.info("context %s exit", name)
  • except에서 로깅 후 raise로 재전파
  • finally는 항상 실행

취소도 다루고 싶다면 Exception 대신 BaseException을 고려할 수 있지만, 운영 코드에서는 취소를 삼키는 실수를 막기 위해 보통은 취소는 그대로 전파하는 쪽이 안전합니다.

여러 자원을 조건부로 열 때: AsyncExitStack이 정답

AsyncExitStack은 “정리 함수들을 스택에 쌓아두고, 마지막에 역순으로 실행”하는 도구입니다. 자원 획득이 여러 단계이고, 중간에 예외가 나도 이미 쌓인 정리 작업은 실행됩니다.

문제 상황: 단계적으로 연결되는 파이프라인

# 안티패턴: 중간 예외 시 일부 자원이 남을 수 있음
async def build_pipeline(cfg):
    session = aiohttp.ClientSession()
    conn = await open_db(cfg.dsn)
    stream = await open_stream(cfg.url)  # 여기서 실패하면 session/conn 정리 누락 가능
    return session, conn, stream

해결: AsyncExitStack으로 획득과 정리를 한 곳에 모으기

import contextlib
import aiohttp

@contextlib.asynccontextmanager
async def pipeline(cfg):
    async with contextlib.AsyncExitStack() as stack:
        session = aiohttp.ClientSession()
        stack.push_async_callback(session.close)

        conn = await open_db(cfg.dsn)
        stack.push_async_callback(conn.close)

        stream = await open_stream(cfg.url)
        stack.push_async_callback(stream.aclose)

        # 여기까지 성공한 자원만 yield
        yield session, conn, stream
        # 블록 종료 시 stack이 역순으로 정리

async def run(cfg):
    async with pipeline(cfg) as (session, conn, stream):
        ...

핵심 이점:

  • 어떤 단계에서 실패하든, 그 시점까지 성공한 자원은 모두 정리됨
  • 정리 순서가 “획득의 역순”으로 고정됨
  • 코드가 커져도 try/finally 중첩 지옥을 피함

enter_async_context로 “진짜 컨텍스트”를 스택에 올리기

이미 async with로 써야 하는 객체라면 enter_async_context가 더 깔끔합니다.

import contextlib
import aiohttp

@contextlib.asynccontextmanager
async def session_and_response(url: str):
    async with contextlib.AsyncExitStack() as stack:
        session = await stack.enter_async_context(aiohttp.ClientSession())
        resp = await stack.enter_async_context(session.get(url))
        yield resp  # resp와 session은 자동 정리

스트리밍/비동기 제너레이터 누수: aclosing으로 마무리 강제

비동기 제너레이터는 “끝까지 소비”하지 않으면 내부 정리 코드가 늦게 실행되거나(가비지 컬렉션 시점), 아예 실행되지 않은 것처럼 보일 수 있습니다. 스트리밍 응답을 중간에 끊는 경우가 대표적입니다.

contextlib.aclosingasync with 블록이 끝날 때 aclose를 호출해줍니다.

예시: 비동기 제너레이터를 안전하게 소비

import contextlib

async def lines(stream):
    async for chunk in stream:
        yield chunk.decode("utf-8")

async def consume_some(stream):
    async with contextlib.aclosing(lines(stream)) as gen:
        async for line in gen:
            if "STOP" in line:
                break  # 중간 탈출해도 gen.aclose() 호출

이 패턴은 다음에 특히 중요합니다.

  • 서버에서 SSE, chunked response, 로그 tailing 같은 무한 스트림
  • 일부만 읽고 끊을 수 있는 프로토콜
  • 백프레셔/취소가 자주 걸리는 파이프라인

취소(CancelledError)와 정리 코드: “정리 중 예외”를 조심하기

실무에서 흔한 사고는 다음입니다.

  • 취소가 들어와서 정리 중인데, 정리 코드에서 또 예외가 발생
  • 그 예외가 원래 취소/원래 예외를 덮어써서 원인 분석이 어려워짐

정리 단계에서는 가능하면 “최선의 정리”를 하되, 실패해도 원래 예외를 보존하는 전략이 좋습니다.

정리 예외는 로깅하고 진행하는 패턴

import contextlib
import logging

log = logging.getLogger(__name__)

@contextlib.asynccontextmanager
async def safe_close(resource, *, name: str):
    try:
        yield resource
    finally:
        try:
            await resource.close()
        except Exception:
            log.exception("failed to close %s", name)

주의: 여기서도 취소를 삼키지 않도록, close에서 CancelledError가 난다면 그대로 전파되게 두는 설계가 일반적으로 낫습니다. 정말로 취소를 무시하고 닫아야 한다면(드물지만 파일 flush 같은 경우), 별도 정책으로 분리하세요.

디버깅: 누수는 어떻게 “증명”하나

1) 테스트에서 종료 여부를 강제 검증

자원 객체에 closed 플래그를 두거나, close가 호출되었는지 스파이로 검증합니다.

import pytest
import contextlib

class Dummy:
    def __init__(self):
        self.closed = False
    async def close(self):
        self.closed = True

@contextlib.asynccontextmanager
async def dummy_ctx(d: Dummy):
    try:
        yield d
    finally:
        await d.close()

@pytest.mark.asyncio
async def test_close_on_exception():
    d = Dummy()
    with pytest.raises(RuntimeError):
        async with dummy_ctx(d):
            raise RuntimeError("boom")
    assert d.closed is True

2) 운영에서는 “증상 기반”으로 먼저 잡는다

누수는 종종 다음 증상으로 시작합니다.

  • 커넥션 풀 고갈
  • 파일 디스크립터 증가
  • 메모리 증가(버퍼/스트림)
  • 타임아웃 증가

이런 류의 장애는 “원인 추적이 어려운 운영 이슈”로 이어지기 쉽습니다. 예를 들어 인프라 레벨에서 원인 가설을 빠르게 좁히는 접근은 K8s ImagePullBackOff - registry auth·CA 10분 진단처럼 체크리스트 기반 진단이 효과적입니다. 애플리케이션 레벨에서도 동일하게, 컨텍스트 종료 여부를 로깅/메트릭으로 관측 가능하게 만드는 게 좋습니다.

실전 체크리스트

  • async with로 감쌀 수 있는 자원은 무조건 감싼다
  • 단계적으로 자원을 획득한다면 AsyncExitStack을 기본값으로 쓴다
  • 비동기 제너레이터/스트리밍은 aclosing으로 중간 탈출에도 종료를 보장한다
  • 정리 코드에서 예외가 나도 원래 예외를 덮지 않게 설계한다
  • 취소는 “정상적인 운영 이벤트”로 보고, 취소 경로에서도 정리가 되는지 테스트한다

마무리

contextlib는 “파이썬스러운 편의 기능” 정도로 보이지만, async 환경에서는 예외/취소/부분 성공이 얽힌 누수 문제를 구조적으로 제거하는 핵심 도구입니다. @asynccontextmanager로 패턴을 표준화하고, AsyncExitStack으로 자원 획득 단계를 안전하게 만들고, aclosing으로 스트리밍 종료를 강제하면, 애매한 누수/헛도는 태스크 문제를 큰 폭으로 줄일 수 있습니다.

다음 단계로는, 컨텍스트 종료를 메트릭으로 계측하고(예: 열린 커넥션 수, 스트림 수), 장애 시나리오에서 취소를 인위적으로 주입해 회귀 테스트를 만드는 것을 추천합니다.