Rust: async의 늪

actix 웹 프레임워크에서 POST를 받으면 load_courses라는 async 함수를 부르는 상황이다.

 

CSV 파일이 커서 tokio를 이용해 async read를 했다.

use csv::ReaderBuilder;
use tokio::fs::File;
use tokio::io::AsyncReadExt;

async fn load_csv_data(file_path: &str) -> Result<Vec<Course>> {
    let mut buf = Vec::new();
    let mut file = File::open(file_path).await?;
    file.read_to_end(&mut buf).await?;

    let mut reader = ReaderBuilder::new()
        .delimiter(b',')
        .flexible(true)
        .from_reader(buf.as_slice());

    let mut courses: Vec<Course> = vec![];
    for result in reader.deserialize() {
        let record: Course = result?;
        courses.push(record);
    }

    Ok(courses)
}

 

csv 파일을 한 번만 부르고 프로그램을 실행 후 인덱싱을 한 번만 하고 싶었다.

include_str처럼 바이너리에 넣으려고 했으나 너무 커서 안 되는 것 같았다.

use lazy_static::lazy_static;
use std::sync::{Arc, Mutex, Once};

lazy_static! {
    static ref INIT: Once = Once::new();

    static ref SUB_ENG_NAME_INDEX: Mutex<Arc<BTreeMap<String, Vec<Course>>>> = Mutex::new(Arc::new(BTreeMap::new()));
    static ref SUB_KOR_NAME_INDEX: Mutex<Arc<BTreeMap<String, Vec<Course>>>> = Mutex::new(Arc::new(BTreeMap::new()));
    static ref COURSE_PROF_NAME_INDEX: Mutex<Arc<BTreeMap<String, Vec<Course>>>> = Mutex::new(Arc::new(BTreeMap::new()));
}

 

BTreeMap을 이용해 키가 정렬된 상태도 유지되고 바이너리 서치를 해 O(n) 보다 빠르게 해 보았다.

그리고 std::sync의 Once를 이용해서 매번 함수가 부를 때마다 인덱스를 만들었는지 확인하는 것보다

low-level에서 지원해 주는 방법을 썼다. 하지만 Rust는 아직 async closure를 지원하지 않기에 load_csv를 사용할 수 없었다.

그래서 OnceCell이란 crate로 바꿔서 set/get을 지원하고 async tokio::Mutex를 지원해서 이걸로 바꿔봤다.

use once_cell::sync::OnceCell;

static INIT: OnceCell<tokio::sync::Mutex<()>> = OnceCell::new();

pub async fn load_courses(query: &str) -> Result<Vec<Course>> {
    let init_lock = INIT.get_or_init(|| tokio::sync::Mutex::new(()));
    {
        let _guard = init_lock.lock().await;
        if SUB_ENG_NAME_INDEX.lock().unwrap().is_empty() {
            let courses = load_csv_data("course.csv").await?;
            index_from_courses(&courses);
        }
    }

    let mut matching_courses = search(query);
    matching_courses.truncate(10);
    Ok(matching_courses)
}

그럼 위같이 바꿀 수 있는데 왜 각각 인덱스를 만들어 따로따로 저장했는지 떠올라서 하나의 인덱스로 바꿨다.

use std::collections::{BTreeMap, HashSet};
use std::sync::Arc;

lazy_static! {
    static ref GLOBAL_INDEX: Mutex<Arc<BTreeMap<String, Vec<Course>>>> =
        Mutex::new(Arc::new(BTreeMap::new()));
}

fn index_from_courses(courses: &[Course]) {
    let mut index = GLOBAL_INDEX.lock().unwrap();
    let mut index_inner = Arc::make_mut(&mut index);

    for course in courses {
        let keys = vec![
            course.subject_english_name.to_lowercase(),
            course.subject_korean_name.to_lowercase(),
            course.main_lecturer_name.to_lowercase(),
        ];

        for key in keys {
            index_inner
                .entry(key)
                .or_insert_with(Vec::new)
                .push(course.clone());
        }
    }
}

fn search(query: &str) -> Vec<Course> {
    let query = query.to_lowercase();
    let index = GLOBAL_INDEX.lock().unwrap();
    let index_clone = Arc::clone(&index);

    let mut results: Vec<Course> = Vec::new();
    let mut added_ids: HashSet<String> = HashSet::new();

    for (key, courses_list) in index_clone.iter() {
        if key.contains(&query) {
            for course in courses_list {
                if !added_ids.contains(&course.id) {
                    results.push(course.clone());
                    added_ids.insert(course.id.clone());
                }
            }
        }
    }

    results
}