Tower Layer 직접 만들기: 커스텀 미들웨어 작성 가이드

Tower 미들웨어 실전에서 본 빌트인 미들웨어로 대부분의 자리는 커버되지만, 가끔은 직접 만들어야 하는 순간이 옵니다. 사내 인증 토큰 검증, 도메인 특화 메트릭 수집, 비표준 헤더 처리 같은 자리죠.

이번 글에서는 Tower의 Service와 Layer를 직접 구현해서 미들웨어를 작성하는 방법을 정리해보겠습니다. 비동기 future를 직접 만드는 패턴까지 포함해서요. 입문 개념은 Tower 입문에서 다뤘으니, 이 글은 그 위에서 한 단계 깊이 들어가는 내용입니다.

만약 Future의 poll, Waker, Pin 같은 단어가 아직 낯설다면 Rust Future란 무엇인가를 먼저 훑어보시면 이 글의 중반부가 훨씬 편하게 읽힙니다.

Service 트레이트 다시 보기

Service 트레이트의 전체 모양은 다음과 같습니다.

trait Service<Request> {
    type Response;
    type Error;
    type Future: Future<Output = Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
    fn call(&mut self, req: Request) -> Self::Future;
}

직접 구현할 때 핵심은 두 메서드입니다.

poll_ready는 “요청 받을 준비됐어?”에 답하는데요. 내부 자원이 가득 찼으면 Pending을, 사용 가능하면 Ready(Ok(()))를 돌려줍니다. 미들웨어 입장에서는 보통 안쪽 서비스의 poll_ready를 그대로 위로 전달해주면 됩니다.

call은 실제로 요청을 처리해서 응답 future를 돌려주는데요. 호출자는 반드시 직전에 poll_ready가 Ready(Ok)를 돌려줬을 때만 call을 호출해야 한다는 계약이 있습니다. 이 계약을 어기면 패닉이 날 수도 있죠.

가장 단순한 Layer: 로깅 미들웨어

요청이 들어올 때마다 로그를 찍는 미들웨어를 만들어보겠습니다. 비동기 future를 새로 만들지 않고 안쪽 서비스의 future를 그대로 돌려주는 가장 단순한 형태입니다.

먼저 미들웨어 구조체를 정의하고 Service를 구현합니다.

use tower::{Service, Layer};
use std::task::{Context, Poll};

#[derive(Clone)]
struct LogService<S> {
    inner: S,
}

impl<S, Request> Service<Request> for LogService<S>
where
    S: Service<Request>,
    Request: std::fmt::Debug,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request) -> Self::Future {
        println!("요청: {req:?}");
        self.inner.call(req)
    }
}

call에서 로그를 한 줄 찍고, 나머지는 안쪽 서비스에 그대로 위임합니다. Future 타입을 S::Future로 선언했기 때문에 변환 없이 그대로 돌려줄 수 있죠.

이 서비스를 ServiceBuilder에서 쓸 수 있게 하려면 대응되는 Layer를 정의해야 합니다.

#[derive(Clone)]
struct LogLayer;

impl<S> Layer<S> for LogLayer {
    type Service = LogService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        LogService { inner }
    }
}

Layer::layer는 안쪽 서비스를 받아 우리 미들웨어로 감싼 서비스를 돌려줍니다. 이 패턴이 모든 Tower 미들웨어의 뼈대입니다.

이제 사용 측에서는 깔끔하게 끼울 수 있습니다.

use tower::{ServiceBuilder, ServiceExt, BoxError, service_fn};

#[tokio::main]
async fn main() -> Result<(), BoxError> {
    let service = ServiceBuilder::new()
        .layer(LogLayer)
        .service_fn(|req: String| async move {
            Ok::<_, BoxError>(format!("ok: {req}"))
        });

    service.oneshot("hello".into()).await?;
    Ok(())
}

응답을 변환하는 미들웨어

응답을 변형하려면 안쪽 서비스의 future를 그대로 돌려주는 대신, 그 future를 감싼 새 future를 만들어야 합니다. 여기서부터 pin-project가 등장하죠.

응답 시간을 측정해서 로그에 남기는 미들웨어를 예로 들어보겠습니다.

use pin_project::pin_project;
use std::pin::Pin;
use std::future::Future;
use std::time::Instant;

#[pin_project]
struct TimedFuture<F> {
    #[pin]
    inner: F,
    start: Instant,
}

impl<F, Res, E> Future for TimedFuture<F>
where
    F: Future<Output = Result<Res, E>>,
{
    type Output = Result<Res, E>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>)
        -> Poll<Self::Output>
    {
        let this = self.project();
        match this.inner.poll(cx) {
            Poll::Ready(result) => {
                let elapsed = this.start.elapsed();
                println!("응답 시간: {elapsed:?}");
                Poll::Ready(result)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

#[pin_project] 매크로가 Pin<&mut Self>에서 안전하게 필드를 꺼낼 수 있게 해줍니다. 비동기 future는 자기 위치를 옮기면 안 되는데(!Unpin), 안에 들어 있는 future도 마찬가지로 고정해줘야 하거든요. #[pin] 어트리뷰트가 그 약속을 코드로 표현한 거고, this.inner가 Pin<&mut F>로 안전하게 꺼내집니다.

이제 이 future를 사용하는 서비스를 정의합니다.

struct TimedService<S> { inner: S }

impl<S, Request> Service<Request> for TimedService<S>
where
    S: Service<Request>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = TimedFuture<S::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request) -> Self::Future {
        TimedFuture {
            inner: self.inner.call(req),
            start: Instant::now(),
        }
    }
}

대응되는 TimedLayer는 위 LogLayer와 같은 패턴으로 만들면 됩니다.

빌트인 미들웨어 들여다보기: Timeout

직접 만든다는 게 어떤 일인지 감을 잡기 위해 Tower의 Timeout 구현을 살펴보면 좋습니다. 핵심만 추리면 다음과 같은 모양인데요.

#[pin_project]
struct ResponseFuture<F> {
    #[pin]
    response_future: F,
    #[pin]
    sleep: tokio::time::Sleep,
}

impl<F, Response, Error> Future for ResponseFuture<F>
where
    F: Future<Output = Result<Response, Error>>,
    Error: Into<BoxError>,
{
    type Output = Result<Response, BoxError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();

        match this.response_future.poll(cx) {
            Poll::Ready(result) => return Poll::Ready(result.map_err(Into::into)),
            Poll::Pending => {}
        }

        match this.sleep.poll(cx) {
            Poll::Ready(()) => Poll::Ready(Err(Box::new(TimeoutError))),
            Poll::Pending => Poll::Pending,
        }
    }
}

응답 future와 sleep future를 둘 다 갖고 있다가, 어느 쪽이 먼저 끝나는지를 보고 결과를 결정합니다. 응답이 먼저 오면 그대로 통과시키고, sleep이 먼저 끝나면 타임아웃 에러로 빠지는 거죠.

이 패턴, 즉 “두 future를 동시에 폴링하다가 먼저 끝나는 쪽을 따른다”는 발상은 직접 미들웨어를 만들 때 자주 등장합니다. 재시도, 회로 차단기(circuit breaker), 헤지(hedging) 같은 고급 패턴이 모두 이 위에 서 있습니다.

합성할 때 주의할 점

미들웨어를 직접 만들고 다른 빌트인 미들웨어와 조합할 때 흔히 마주치는 함정이 있습니다.

첫째로 Service는 &mut self로 호출됩니다. 그래서 안쪽 서비스가 Clone을 구현하지 않으면 한 번에 한 곳에서만 들고 있을 수 있어 동시 호출이 어렵습니다. 보통 미들웨어 자체와 안쪽 서비스 모두 Clone을 요구하도록 설계하면 안전한데요. 빌트인 미들웨어들도 모두 Clone을 요구합니다.

둘째로 에러 타입을 통일해줄 필요가 있습니다. 미들웨어를 합성하다 보면 안쪽 서비스의 에러 타입과 미들웨어가 만들어내는 에러 타입(예: 타임아웃 에러)이 달라지는데요. Tower는 보통 BoxError = Box<dyn std::error::Error + Send + Sync>를 공통 에러 타입으로 쓰고, 각 미들웨어는 S::Error: Into<BoxError> 같은 바운드를 두어 자동 변환되도록 설계합니다. 여기서 BoxError는 표준 라이브러리에 있는 타입이 아니라 tower 크레이트가 제공하는 단축 별칭인데, Box<dyn ...>를 매번 길게 쓰지 않게 해주는 정도의 역할입니다. 실제 변환에는 특별한 장치가 없습니다. 앞서 본 Timeout 코드에서 응답을 통과시키던 result.map_err(Into::into)가 바로 안쪽 에러를 BoxError로 박싱하는 지점이죠. ? 연산자도 같은 Into 변환을 자동으로 끼워주기 때문에, 호출하는 쪽에서는 에러 타입이 서로 다르다는 걸 거의 의식하지 않고 합성할 수 있습니다.

셋째로 poll_ready의 계약을 지켜야 합니다. poll_ready가 Ready(Ok)를 돌려주면 다음 call까지 그 슬롯을 예약한 것과 같은데, 만약 poll_ready만 호출하고 call을 안 하면 자원이 낭비될 수 있습니다. 직접 호출하지 말고 ServiceExt::oneshot이나 ready_and() 같은 헬퍼를 쓰는 편이 안전한 이유입니다.

마치며

Tower 미들웨어를 직접 만드는 일은 처음에는 트레이트 바운드가 무서워 보이지만, 패턴은 단순합니다. Service로 처리 단계를 정의하고, 응답 변환이 필요하면 pin-project로 future를 만들고, 대응되는 Layer로 빌더에 끼울 수 있게 만들어주는 흐름이죠.

이 글에서 본 패턴은 Tower 빌트인 미들웨어도 똑같이 따르고 있으니, 빌트인 코드를 한 번 들여다보시면 직접 만든 미들웨어를 어떻게 다듬을지 감이 잡힙니다. Axum 서버에서 직접 만든 미들웨어를 끼우는 실전 예제는 Tower 미들웨어 실전에 함께 정리되어 있습니다.

더 자세한 내용은 Tower 공식 가이드의 미들웨어 작성 문서를 참고하세요.