Astral의 Rust 기반 파이썬 패키지 매니저 uv가 중복 작업을 제거하기 위해 사용하는 경량 동시 메모이제이션 프리미티브 OnceMap의 설계와 레이스 컨디션 수정 과정을 살펴본다.
uv는 Astral에서 만든 Rust 기반 파이썬 패키지 매니저로, pip보다 10~100배 빠르다. 이 속도 향상의 일환으로 uv는 의존성 해석, 휠 다운로드, 소스 빌드를 공격적으로 병렬화한다.
uv pip install scipy pandas를 실행하면, 두 패키지 모두 numpy에 의존하므로 두 스레드가 각각 독립적으로 numpy 메타데이터를 가져오려 할 수 있다. 조율이 없다면 중복 요청이 발생한다. OnceMap은 첫 번째 호출자만 데이터를 가져오게 하고, 나머지는 그 결과를 기다리게 함으로써 이를 방지한다.
병렬화는 조율 문제를 만든다. 여러 태스크가 같은 리소스가 필요할 때, 어떻게 그 일이 정확히 한 번만 일어나도록 보장할까? uv의 해법은 OnceMap이다. 이는 리졸버와 인스톨러 전반에서 중복 제거를 떠받치는, 가벼운 동시 메모이제이션(concurrent memoization) 프리미티브다.
OnceMap 구조체는 동시 해시맵인 DashMap을 감싼다. 단일 락을 쓰는 표준 HashMap과 달리, DashMap은 데이터를 여러 세그먼트로 샤딩(sharding)하고 각 세그먼트마다 별도 락을 둔다. 덕분에 락 경합 없이 병렬 접근이 가능하다.
Value enum은 핵심 요소다. 엔트리는 단순히 “존재/비존재”만이 아니라 진행 중(in progress), 완료(completed), 시작 전(not started) 같은 상태를 추적한다. Waiting 변형(variant)은 Tokio의 Notify를 담고 있는데, 이는 태스크가 효율적으로 잠들고 깨어날 수 있게 해주는 가벼운 동기화 프리미티브다.
OnceMap은 register, done, wait의 3-메서드 프로토콜을 사용한다.
register는 불리언 계약을 반환한다. true라면 당신은 반드시 작업을 수행하고 done을 호출해야 한다. false라면 다른 태스크가 이미 그 작업을 처리 중이다.
done은 Waiting 센티널을 결과로 교체하고, notify_waiters()로 대기 중인 모든 태스크를 깨운다.
이 설계는 개발 중 발견된 레이스 컨디션을 해결한다(Issue #3724).
원래 구현에는 결함이 있었다.
register()를 호출하고 작업을 시작한다.register()를 호출하고, 진행 중임을 확인한 뒤 wait()를 호출한다.Notify 핸들을 가져온다.done()을 호출해 대기자에게 신호를 보낸다.notify.notified().await를 호출한다.Notify의 신호는 큐잉되지 않으므로, 태스크 B는 알림을 놓치고 영원히 멈춘다.
해결책은 위 wait() 구현처럼, 맵을 다시 확인하기 전에 먼저 대기자로 등록하는 것이다.
pin!(notify.notified())는 태스크 B가 즉시 신호를 수신하도록 등록한다. 두 번째 get(key) 확인 도중에 태스크 A가 notify_waiters()를 호출하더라도, 태스크 B는 이를 놓치지 않는다. pin! 매크로는 future를 고정된 메모리 주소에 붙잡아 둔다. Rust의 future는 기본적으로 이동 가능(movable) 값이다. Notified는 Tokio의 대기 큐에 자기 자신의 주소를 등록하므로, 이후에 이동해버리면 댕글링 포인터가 된다. pin!은 future가 그 자리에 머무르도록 보장한다.
레이스 컨디션 수정은 전적으로 wait() 안에만 들어가 있다는 점에 유의하자.
OnceMap은 휠 다운로드와 메타데이터 가져오기를 조율한다. 다음은 인스톨러가 이를 사용하는 방식이다.
astral-sh/uv:preparer.rs#L109-L211
한 태스크가 작업을 수행하는 동안, 동시에 들어온 요청들은 공유 결과를 기다린다.
왜 표준 캐시가 아닌가? 캐시는 결과를 저장하지만, 진행 중(in-flight)인 작업을 추적하지는 못하므로 중복 요청을 막을 수 없다.
왜 DashMap + Notify 인가? DashMap은 세밀한 락으로 경합을 최소화하고, Notify는 비동기 동기화를 효율적으로 처리한다.
왜 가져올 때 clone을 하나? 참조를 반환하려면 async 경계 너머로 락을 잡은 채로 유지해야 하며, 이는 데드락 위험을 만든다. Arc<V>를 쓰면 clone 비용이 작다.
PR #544 OnceMap으로 진행 중인(un-flight) unzip을 추적 by @konstin (konsti)
PR #3627 리졸버 병렬화 by @ibraheemdev (Ibraheem Ahmed)
PR #3987 OnceMap에서 레이스 컨디션 회피 by @ibraheemdev (Ibraheem Ahmed)