Rust에서 인체공학적인 Async 트레이트 객체

Async Rust는 다루기 까다로울 수 있고, 라이브러리 유지관리자 입장에서는 사용자가 비동기 코드의 골칫거리를 떠안지 않도록 가능한 한 많은 무거운 일을 대신 처리해주는 것이 중요합니다. 우리는 최근 cot에 동적 캐시 백엔드 지원을 추가했는데, 그 과정에서 작은 도전 과제 하나가 API 설계와 사용성(ergonomics)을 제대로 맞추는 일이었습니다. 여기서 사용하는 사용 사례와 코드는 코드베이스에서 그대로 가져온 것입니다. Cot은 Django에서 영감을 받은 Rust 웹 프레임워크입니다.

문제 배경

우리는 cot 웹 프레임워크에 캐싱 메커니즘을 추가해야 했습니다. 사용자는 단순한 캐시 인터페이스를 사용하되, 객체가 캐시에 어떻게 저장되고 조회되는지는 추상화되어야 합니다. 또한 여러 캐시 백엔드를 지원하고, 원한다면 사용자가 자신만의 커스텀 캐시 백엔드를 구현할 수 있게 하고 싶었습니다. 특히 async 코드를 다루는 상황에서, Rust로 이런 API를 어떻게 설계할까요?

동기(synchronous) 버전

먼저 단순한 동기 구현부터 시작해봅시다. 아래 코드를 보겠습니다:

rust
use std::collections::HashMap;
use serde_json::Value;
use serde::{Serialize, de::DeserializeOwned};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum CacheError {
    #[error("Cache operation failed: {0}")]
    Backend(String),
    #[error("Serde error: {0}")]
    SerdeJson(#[from] serde_json::Error),
}

pub type CacheResult<T> = std::result::Result<T, CacheError>;

pub trait CacheStore: Send + Sync + 'static {
    fn get(&self, key: &str) -> CacheResult<Option<Value>>;
    fn insert(&mut self, key: &str, value: Value) -> CacheResult<()>;
    fn remove(&mut self, key: &str) -> CacheResult<()>;
}

#[derive(Debug, Clone)]
pub struct Memory {
    map: HashMap<String, Value>,
}

impl Memory {
    pub fn new() -> Self {
        Self {
            map: HashMap::new(),
        }
    }
}

impl CacheStore for Memory {
    fn get(&self, key: &str) -> CacheResult<Option<Value>> {
        Ok(self.map.get(key).cloned())
    }

    fn insert(&mut self, key: &str, value: Value) -> CacheResult<()> {
        self.map.insert(key.to_string(), value);
        Ok(())
    }

    fn remove(&mut self, key: &str) -> CacheResult<()> {
        self.map.remove(key);
        Ok(())
    }
}

pub struct Cache {
    store: Box<dyn CacheStore>,
}

impl Cache {
    pub fn new(store: impl CacheStore) -> Self {
        let store = Box::new(store);
        Self { store }
    }

    pub fn get<K, V>(&self, key: K) -> CacheResult<Option<V>>
    where
        K: AsRef<str>,
        V: DeserializeOwned,
    {
        self.store
            .get(key.as_ref())?
            .map(serde_json::from_value)
            .transpose()
            .map_err(CacheError::from)
    }

    pub fn insert<K, V>(&mut self, key: K, value: V) -> CacheResult<()>
    where
        K: AsRef<str>,
        V: Serialize,
    {
        self.store.insert(key.as_ref(), serde_json::to_value(value)?)
    }

    pub fn remove<K>(&mut self, key: K) -> CacheResult<()>
    where
        K: AsRef<str>,
    {
        self.store.remove(key.as_ref())
    }
}

이 코드에서는 여러 캐시 저장소를 지원하려고 공통 인터페이스로 CacheStore 트레이트를 정의합니다. 캐시 데이터를 관리하기 위한 몇 가지 필수 연관 메서드(associative methods)만 요구하도록 최소한으로 유지해봅시다.

또한 CacheStore가 제공하는 메서드를 구현한 간단한 인메모리 저장소를 제공합니다. file이나 redis 같은 다른 저장소를 추가하는 것은 CacheStore의 메서드들을 구현하기만 하면 됩니다. 그리고 사용자가 직접 상호작용할 수 있는 공개 API인 Cache 구조체를 제공합니다. 이는 저장소와 캐시 관련 설정들을 함께 감싸고, 대부분의 작업을 저장소로 전달하는 캐시 관리 메서드를 제공합니다.

아래 코드로 테스트할 수 있습니다:

rust
fn main() {
    let store = Memory::new();
    let cache = Cache::new(store);
    
    cache.insert("foo", "bar");
    let expected: Option<String> = cache.get("foo");
    assert_eq!(expected, Some("foo".to_string()));
}

이는 잘 동작합니다. 하지만 코드는 동기 방식이며, 특히 redis나 I/O를 수행하는 file 같은 다른 저장소를 추가하면 확장성이 떨어집니다. 따라서 설계를 비동기 방식으로 바꿔야 합니다.

순진한(naive) async 버전

Rust는 async 트레이트를 지원하므로, (생략된 라이프타임을 보여주지 않는다면) 업데이트된 코드는 대략 아래처럼 보일 것입니다:

rust
pub trait CacheStore: Send + Sync + 'static{
    async fn get(&self, key: &str) -> CacheResult<Option<Value>>;
    async fn insert(&self, key: &str, value: Value) -> CacheResult<()> ;
    async fn remove(&self, key: &str) -> CacheResult<()>;
}

#[derive(Debug, Clone)]
pub struct Memory {
    map: Arc<Mutex<HashMap<String, Value>>>,
}

impl CacheStore for Memory {
    async fn get(&self, key: &str) -> CacheResult<Option<Value>> {
        // implementation
    }

    ...
}

...
impl CacheStore {

    ...
    pub async fn get<K, V>(&self, key: K) -> CacheResult<Option<V>>
    where
        K: AsRef<str>,
        V: DeserializeOwned,
    {
        self.store.get(key.as_ref()).await?.map(serde_json::from_value).transpose().map_err(CacheError::from)
    }

    pub async fn insert<K, V>(&self, key: K, value: V) -> CacheResult<Option<V>>
    where
        K: AsRef<str>,
        V: Serialize,
    {
        self.store.insert(key.as_ref(), serde_json::to_value(serde_json::to_value(value)?)).await
    }

    ...

}

위 코드에서는 인메모리 저장소를 Mutex 뒤에 두고, 트레이트 메서드에서 mut 요구사항을 제거합니다. 또한 메인 함수를 tokio의 async 런타임을 사용하도록 바꿉니다:

rust
#[tokio::main]
fn main() -> Result<(), String> {
    let store = Memory::new();
    let cache = Cache::new(store);
    
    cache.insert("foo", "bar");
    let expected: Option<String> = cache.get("foo").await?;
    assert_eq!(expected, Some("foo".to_string()));
}

하지만 불행히도, 이는 기대대로 동작하지 않습니다. 보통 아래와 같은 에러를 만나게 됩니다:

text
error[E0038]: the trait `CacheStore` is not dyn compatible
  --> src/bin/asyn_t.rs:52:16
   |
52 |         Self { store }
   |                ^^^^^ `CacheStore` is not dyn compatible
   |
note: for a trait to be dyn compatible it needs to allow building a vtable
      for more information, visit <https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility>
  --> src/bin/asyn_t.rs:11:14

게다가 트레이트 객체와 함께 공개 트레이트 메서드에서 async fn을 쓰려고 시도하면 컴파일러가 여러분에게 고함을 지를 것입니다.

Dyn 호환성

왜 이런 일이 발생하는지 이해하려면 Rust의 dyn 호환성(dyn compatibility) 개념을 알아야 합니다. Rust에서 트레이트는 dyn Trait 형태의 트레이트 객체로 사용할 수 있는데, 오직 트레이트가 dyn compatible일 때만 가능합니다. 트레이트가 dyn compatible이 되려면 두 가지 주요 조건을 만족해야 합니다:

1. 모든 트레이트 메서드는 제네릭 타입 파라미터를 가지면 안 된다.
2. 모든 메서드는 수신자(receiver)가 &self, &mut self, 또는 Box<Self>여야 한다.

트레이트의 async 함수 문제는, 이것들이 impl Future를 반환하는 함수로 디슈가(desugar)된다는 점입니다. 즉, 우리의 예시는 내부적으로 대략 이런 형태가 됩니다:

rust
pub trait CacheStore: Send + Sync + 'static{
    fn get(&self, key: &str) -> impl Future<Output = Result<Value, String>>;
    ...
}

이는 메서드가 이제 제네릭 반환 타입을 갖게 된다는 뜻이며, dyn 호환성의 첫 번째 조건을 위반합니다.

async-trait 크레이트

이 제한을 우회하는 한 가지 방법은 async-trait 크레이트를 사용하는 것입니다. 이 크레이트는 dyn 호환성을 유지하면서 트레이트에 async 함수를 정의할 수 있게 해줍니다. 코드를 async-trait를 사용하도록 바꾸면 다음과 같습니다:

rust
use async_trait::async_trait;

#[async_trait]
pub trait CacheStore: Send + Sync + 'static{
    async fn get(&self, key: &str) -> Result<Option<Value>, String>;
    ...
}

...
#[async_trait]
impl CacheStore for Memory {
    async fn get(&self, key: &str) -> CacheResult<Option<Value>> {
        // implementation
    }
    ...
}

트레이트와 구현에 #[async_trait]를 붙이면, 이 크레이트가 async 메서드를 dyn 호환적으로 만들기 위한 필요한 변환을 처리해줍니다. 덕분에 dyn 호환성 문제 없이 dyn CacheStore를 사용할 수 있습니다.

하지만 우리는 async-trait 사용을 제한하기로 했습니다. 때때로 디버깅하기 어려운 에러 메시지를 생성하기 때문입니다. 또한 async-trait를 사용하면 트레이트 메서드 문서에 불필요한 디테일이 잔뜩 포함되어 가독성이 떨어집니다. 예를 들어 cargo doc --open을 실행하면 CacheStore 트레이트가 다음과 같이 표시됩니다:

rust
pub trait CacheStore:
Send
+ Sync
+ 'static {
    // Required methods
    fn get<'life0, 'life1, 'async_trait>(
        &'life0 self,
        key: &'life1 str,
    ) -> Pin<Box<dyn Future<Output = CacheResult<Option<Value>>> + Send + 'async_trait>>
    where Self: 'async_trait,
          'life0: 'async_trait,
          'life1: 'async_trait;
    fn insert<'life0, 'life1, 'async_trait>(
        &'life0 self,
        key: &'life1 str,
        value: Value,
    ) -> Pin<Box<dyn Future<Output = CacheResult<()>> + Send + 'async_trait>>
    where Self: 'async_trait,
          'life0: 'async_trait,
          'life1: 'async_trait;
    fn remove<'life0, 'life1, 'async_trait>(
        &'life0 self,
        key: &'life1 str,
    ) -> Pin<Box<dyn Future<Output = CacheResult<()>> + Send + 'async_trait>>
    where Self: 'async_trait,
          'life0: 'async_trait,
          'life1: 'async_trait;
}

그러면 이런 질문이 생깁니다: async-trait 없이도 사용성, 성능, 문서 품질을 어떻게 균형 있게 맞출 수 있을까요? 답을 찾으려면 async-trait가 내부에서 하는 일을, 수동으로 비슷하게 해줘야 합니다.

수동 Pin 및 박싱(Boxed) Future

앞서 언급했듯이, 트레이트의 async 함수는 impl Future를 반환하는 함수로 디슈가되며 이는 dyn 호환이 아닙니다. 하지만 반환된 future를 박싱(Box<dyn Future + Send>)하면 구체 타입(concrete type)이 되므로 dyn 호환이 가능합니다.

다만 이것만으로는 충분하지 않습니다. future는 자기참조(self-referential)일 수 있기 때문입니다(자기 내부에 자기 데이터에 대한 참조를 포함). 따라서 Pin 래퍼 타입으로 메모리에 고정(pin)해야 합니다. 이는 future가 메모리 내에서 이동하는 것을 막아 내부 참조가 무효화되는 것을 방지합니다.

그 결과 트레이트 메서드 시그니처는 다음과 같이 됩니다:

rust
use std::pin::Pin;
use std::future::Future;

pub trait CacheStore: Send + Sync + 'static{
    fn get(&self, key: &str) -> Pin<Box<dyn Future<Output = CacheResult<Option<Value>>> + Send>>;
    ...
    // Other trait methods (insert and remove) should have similar signature.
}

이제 Memory에 대한 CacheStore 구현을 아래처럼 업데이트할 수 있습니다:

rust
use std::pin::Pin;
use std::future::Future;
use std::boxed::Box;
use std::collections::HashMap;
use serde_json::Value;
use std::sync::{Arc, Mutex};
use futures::future;

#[derive(Debug, Clone)]
pub struct Memory {
    map: Arc<Mutex<HashMap<String, Value>>>,
}

...

impl CacheStore for Memory {
    fn get(&self, key: &str) -> Pin<Box<dyn Future<Output = CacheResult<Option<Value>>> + Send>> {
        let map = self.map.clone();
        let key = key.to_string();
        Box::pin(async move {
            let map = map.lock().unwrap();
            Ok(map.get(&key).cloned())
        })
    }

   ...
}

공개 API인 Cache는 이전과 동일하게 유지할 수 있고, 이를 실행하면 기대대로 동작합니다. 좋은 점은 Cache API를 사용하는 사용자들은 async 트레이트 복잡성을 전혀 다룰 필요가 없고, 그대로 사용하면 된다는 것입니다.

하지만 트레이트 메서드 구현의 시그니처는 매우 보기 흉하고 장황합니다:

rust
fn get(&self, key: &str) -> Pin<Box<dyn Future<Output = CacheResult<Option<Value>>> + Send>>

예를 들어 다운스트림 사용자나 서드파티 라이브러리 작성자가 자신만의 캐시 저장소를 구현하려면 이런 장황함을 감당해야 하는데, 사용성 측면에서 이상적이지 않습니다. 이들을 위해서도 더 편하게 만들어줄 수는 없을까요?

블랭킷(blanket) async 트레이트 구현

CacheStore 트레이트가 고정된 박싱 future를 직접 반환하는 대신, 모든 더러운 작업을 수행하는 헬퍼 트레이트에 이를 위임하고, CacheStore 트레이트는 평소처럼 async 메서드를 정의하도록 만들 수 있습니다. 실제로 어떻게 되는지 봅시다:

rust
use std::pin::Pin;
use std::future::Future;
use serde_json::Value;

pub(crate) trait BoxedCacheStore: Send + Sync + 'static {
    fn get<'a>(&'a self, key: &'a str) -> Pin<Box<dyn Future<Output = CacheResult<Option<Value>>> + Send + 'a>>;
    ...
}

pub trait CacheStore: Send + Sync + 'static {
    fn get(&self, key: &str) -> impl Future<Output = CacheResult<Option<Value>>> + Send;
    ...
}

impl<T: CacheStore> BoxedCacheStore for T {
    fn get<'a>(&'a self, key: &'a str) -> Pin<Box<dyn Future<Output = CacheResult<Option<Value>>> + Send + 'a>> {
        Box::pin(async move {
            T::get(self, key).await
        })
    }

    ...
}

이제 무슨 일이 일어나는지 풀어봅시다. 우리는 내부 헬퍼 트레이트 BoxedCacheStore를 제공하는데, 이는 이전 CacheStore와 유사한 시그니처의 메서드들을 정의합니다. 그리고 CacheStore 트레이트를 구현하는 어떤 타입 T에 대해서도 이 트레이트를 블랭킷 구현합니다. 이 구현은 CacheStore의 실제(구체) 구현을 호출하고, 반환된 future를 박싱하는 역할만 합니다.

또한 BoxedCacheStore 메서드에 명시적 라이프타임을 제공한 점을 볼 수 있습니다. 이는 Rust의 트레이트 객체에 적용되는 라이프타임 생략 규칙(elision rules) 때문에 매우 중요합니다. 명시적 라이프타임이 없다면 컴파일러는 트레이트에 대해 'static을 가정합니다:

rust
fn get(&self, key: &str) -> Pin<Box<dyn Future<Output = ...> + Send>>;

// is equivalent to

fn get(&self, key: &str) -> Pin<Box<dyn Future<Output = ...> + Send + 'static>>;

하지만 우리의 future는 self와 key 파라미터로부터 빌림(borrow)을 하고 있기 때문에 'static일 수 없습니다. 따라서 명시적 라이프타임을 생략하는 것은 불법이 됩니다. 명시적으로 라이프타임을 제공함으로써, 우리는 박싱된 future의 라이프타임이 캡처한 참조들의 라이프타임에 묶여 있음을 컴파일러에게 알려줍니다.

이제 Memory 구현을 아래처럼 업데이트합니다:

rust
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use serde_json::Value;

#[derive(Debug, Clone)]
pub struct Memory {
    map: Arc<Mutex<HashMap<String, Value>>>,
}

impl Memory {
    pub fn new() -> Self {
        Self {
            map: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

impl CacheStore for Memory {
    async fn get(&self, key: &str) -> CacheResult<Option<Value>> {
        let map = self.map.lock().unwrap();
        Ok(map.get(key).cloned())
    }

    ...
}

훨씬 깔끔합니다! 고정된 박스도 없고, future를 수동으로 박싱할 필요도 없고, 그저 평범한 async 함수입니다.

이제 Cache 구조체를 업데이트하여 공개 API에서는 CacheStore를 노출하되 내부적으로는 BoxedCacheStore를 사용하도록 할 수 있습니다:

rust
use std::sync::Arc;
use serde::{Serialize, de::DeserializeOwned};

#[derive(Clone)]
pub struct Cache {
    store: Arc<dyn BoxedCacheStore>,
}

impl Cache {
    pub fn new(store: impl CacheStore) -> Self {
        let store = Arc::new(store) as Arc<dyn BoxedCacheStore>;
        Self { store }
    }

    pub async fn get<K, V>(&self, key: K) -> CacheResult<Option<V>>
    where
        K: AsRef<str>,
        V: DeserializeOwned,
    {
        self.store
            .get(key.as_ref())
            .await?
            .map(serde_json::from_value)
            .transpose()
            .map_err(CacheError::from)
    }

    ...
}

보시다시피 이제 store 필드는 트레이트 객체로 BoxedCacheStore를 사용하지만, 공개 API는 CacheStore를 구현하는 어떤 타입이든 받아들입니다. 그리고 Cache 인스턴스를 만들 때 제공된 저장소를 BoxedCacheStore로 캐스팅합니다. 개인적으로 훨씬 보기 좋다고 생각합니다.

이 접근의 단점은 보일러플레이트가 늘어나고 유지보수할 코드가 더 많아진다는 점입니다. 하지만 보일러플레이트는 최소한이고 대부분 기계적인 코드입니다. 헬퍼 트레이트는 단지 실제 구현으로 위임(delegation)하기만 합니다. 더 나은 사용성과 문서 품질을 위한 작은 대가라고 할 수 있습니다.