Rust không nhanh như bạn nghĩ!
Tôi đã viết lại một thư viện từ Python sang Rust vì muốn tối ưu hiệu năng. Kết quả chương trình của tôi chạy nhanh bằng 1/15 lần so với thư viện Python gốc.
Vấn đề
Mình có xây dựng một hệ thống xử lý bất đồng bộ đơn giản gồm 2 services như hình:
Producer được viết bằng Python, liên tục tạo ra các messages bằng định dạng protobuf trên một Inter-Process Communication (IPC) channel.
Consumer được viết bằng ngôn ngữ khác Python, cũng liên tục lấy các messages trên IPC này để xử lý.
Vấn đề xuất hiện khi Consumer của mình xử lý messages quá nhanh, trong khi tốc độ tạo messages của Producer lại chậm hơn nhiều. Điều này dẫn đến mình không thể tận dụng tối đa tài nguyên hệ thống vì Consumer luôn có khoảng thời gian chờ messages mới từ Producer.
Mình đã quyết định viết lại Producer bằng Rust để giải quyết bottleneck. Producer của mình sử dụng thư viện Domato do team security của Google viết để tạo messages, viết lại Producer đồng nghĩa với việc mình phải viết lại thư viện này sang Rust.
Sau khi viết lại bằng Rust xong, mình nhanh trí dùng hyperfine để benchmark xem chương trình mình chạy nhanh hơn bao nhiêu và kết quả:
Thật ngạc nhiên, code Python trung bình chạy chỉ tốn 255.4 milliseconds, trong khi Rust trung bình chạy mất 4.016 seconds. Nghĩa là code Python cũ chạy nhanh hơn gấp 15 lần code Rust mới (╯‵□′)╯︵┻━┻
Giải quyết vấn đề
Python chạy nhanh hơn Rust gấp 15 lần là một điều phi lý, chắc chắn code mình viết có vấn đề. Nhưng code mình viết gần như là sao chép lại code Python cũ, mình chỉ sửa lại cho phù hợp cú pháp cũng như rebuild một số structures và enums.
Để tìm ra được code mình có vấn đề ở đâu, mình dùng flamegraph-rs để profiling xem CPU tốn thời gian xử lý nhất ở những hàm nào.
Như hình ở trên, mình có khoanh vùng xanh lá, hàm regex_automata::meta::strategy::new chiếm tận 48% thời gian thực thi của CPU. Sau khi đọc được một bình luận ở Rust Forum, mình đã hiểu nguyên nhân do hàm Regex::new được gọi nhiều lần, mỗi lần gọi như vậy thì regex sẽ thực hiện compile regex expression lại. Thực ra mình cũng đã nhận ra vấn đề này nên mình đã không đưa Regex::new vào vòng lặp, mà khởi tạo nó trước vòng lặp như này:
Nhưng mình lại sót trường hợp hàm include_from_string lại được gọi nhiều lần và không ngờ cái giá phải trả cho CPU quá đắt 🥲
Mình đã sử dụng crate once_cell theo hướng dẫn và regex macro để sửa lại việc gọi regex. Code trông như thế này:
Chạy lại benchmark sau khi đã tối ưu phần regex, mình nhận được:
Tin vui là tốc độ đã cải thiện gấp 6 lần so với lần trước, tin buồn thì code Python vẫn chạy nhanh hơn code của mình 2.5 lần (╯‵□′)╯︵┻━┻
Lại tiếp tục phân tích flamegraph xem mình code lởm ở chỗ nào nữa.
Để ý 2 vùng khoanh màu tím than, mỗi vùng chiếm khoảng 12%, tổng là 24% CPU time đều gọi Option<&T>::cloned từ hàm select_creator
Clone/cloned mặc định của Rust gần giống với deep copy ở Python, nên nếu mình bỏ Option<&T>::cloned thì sẽ cải thiện được kha khá tốc độ. Ngoài ra, ở dòng 30 và 35, mình đều gọi .clone() rồi, tức lời gọi .cloned() ở dòng 17 và 22 là dư thừa, mình không nhớ tại sao mình lại gọi dư như vậy nữa. Và clippy cũng không phát hiện ra được pattern này luôn.
Xóa bỏ .cloned() trong hàm select_creator, mình chạy benchmark một lần nữa:
Lần này thì khả quan hơn rồi, tốc độ đều xấp xỉ nhau. Nhưnggggg, bỏ công sức ra để port một thư viện từ Python sang Rust với lý do hiệu năng, mà lại kết thúc ở mức tốc độ xấp xỉ nhau ư? Tất nhiên là không rồi.
Mình lại tiếp tục công cuộc vạch lá tìm sâu cùng với flamegraph:
Flamegraph trông có vẻ bình thường, nhưng với nhiều năm kinh nghiệm bắt sâu, kết hợp với bài học .cloned()/.clone() khá là đắt, mình dùng tính năng search và highlight của flamegraph để chắt lọc lại dữ liệu.
Bằng việc search từ khóa ::clone, kết quả ở dưới góc phải có đến 40.3% matches. Và các matches chủ yếu là:
<domato::grammar::Grammar as core::clone::Clone>::clone
<domato::grammar::Rule as core::clone::Clone>::clone
<domato::grammar::Context as core::clone::Clone>::clone
<domato::grammar::Tag as core::clone::Clone>::clone
Mình có kiểm tra xem mình có gọi .clone() thừa cho những structures này, nhưng không có chút dư thừa nào, mọi lời gọi .clone() đều phù hợp để đảm bảo borrow checker của Rust trả về là hợp lệ.
Phân tích kỹ hơn một ví dụ <domato::grammar::Rule as core::clone::Clone>::clone xem nó thực thi những gì:
Bên trong Rule::clone tiếp tục gọi ::clone của những thuộc tính bên trong, phần chiếm nhiều CPU nhất là alloc::slice::hack::to_vec. Mình hiểu ngay nếu vector càng lớn, thì việc clone nó sẽ càng tốn nhiều thời gian bởi số lượng item cần được clone trong vector.
Kiểm tra lại cấu trúc của các structures Grammar, Rule, Context, Tag.
Những cấu trúc này đều dùng Vector để chứa lồng lẫn nhau. Biết được nguyên nhân rồi, nhưng bây giờ phải làm thế nào để giảm tối thiểu việc clone như deep copy?
Giải pháp là sử dụng std::rc::Rc (Reference Counted), nếu từng đọc qua RustBook Chapter 15.4 hoặc sử dụng các ngôn ngữ dùng Garbage Collector để quản lý bộ nhớ như Python, Java, JavaScript hay Golang thì sẽ quen thuộc với thuật ngữ này.
Trong tài liệu của Rc https://doc.rust-lang.org/std/rc/struct.Rc.html#method.clone giải thích khá rõ về việc cách Rc thực hiện clone:
fn clone(&self) -> Rc<T, A>
Makes a clone of the Rc pointer.
This creates another pointer to the same allocation, increasing the strong reference count.Hiểu đơn giản thì Rc::clone sẽ chỉ thực hiện self.reference_counter += 1 thay vì self.inner_object.clone(). Như vậy clone sẽ không còn là deep copy nữa, đúng thứ mà mình cần. Giờ chỉ cần sửa lại một chút code và cấu trúc các structures.
Chạy lại benchmark để hóng kết quả:
Lần này chương trình viết bằng Rust của mình cuối cùng cũng chạy nhanh gấp 1.8 lần so với code cũ Python. Mình cũng tạm hài lòng với kết quả này nên không tìm cách tối ưu hơn nữa. Nhưng mình đoán là mình vẫn chưa thực sự làm nó tối ưu tốt nhất.
Bài học
Toàn bộ những thay đổi về code để tối ưu tốc độ nằm ở commit https://github.com/ByteDeflect/domato-rs/commit/065dc427e94fb30f71df103af7a74e08c5a0a1b1.
“Blazing fast” đúng là một cái meme của Rust 🫠
Nên sử dụng regex::new cùng với once_cell, điều này đã được khuyến cáo ở regex repo: Usage: Avoid compiling the same regex in a loop.
Không nên sử dụng clone/cloned bừa bãi.
Nếu bạn muốn viết lại một công cụ, thư viện từ ngôn ngữ khác sang Rust để chạy trong production, hãy benchmark cẩn thận.
Rust không phải là ngôn ngữ dễ sử dụng, đặc biệt sử dụng với lý do hiệu năng. Bạn cần nắm một số kiến thức cơ bản của Khoa học máy tính thì mới thực sự tận dụng được hiệu năng của Rust.