История о недостающих метриках: странности с замыканиями в Rust

Как один нюанс, связанный с обработкой замыканий в Rust, привёл к потере метрик в нашей компании, и какие уроки мы из этого извлекли

В потоковой базе данных у вас должны быть в распоряжении точные метрики – это не просто плюс, а важнейшая предпосылка для отслеживания производительности, выявления узких мест и обеспечения ровного функционирования системы в режиме реального времени. Метрики можно сравнить с циферблатами на приборном щитке автомобиля. Именно по ним вы получаете оперативную обратную связь и узнаёте, что творится «под капотом». Один из ключевых показателей, который необходимо отслеживать — это скорость извлечения данных с того уровня, где они хранятся. В системе реального времени такая функция неоценима.

Нам довелось столкнуться с загадочной ситуацией: одна из наших метрик, характеризующая хранение данных, якобы оставалась на нуле, даже когда нам было точно известно, что через систему прокачиваются данные. Всё равно, что вести машину со сломанным спидометром — мы же едем, но не знаем, с какой скоростью. В результате нам пришлось устроить настоящее расследование и углубиться в наш код на Rust. В результате мы узнали много нового о том, как в Rust раньше обрабатывались замыкания.

Что такое RAII и почему это важно

В Rust существует концепция под названием RAII. Эта аббревиатура означает «Приобретение ресурсов равноценно их инициализации». Представьте себе обслуживающую команду, которая автоматически прибирает за вами после вечеринки.

Когда некоторые данные (ресурс) больше не нужны, механизм RAII отвечает за их правильную утилизацию. Это важно для предотвращения утечек памяти и, естественно, для сообщения самих метрик.

Допустим, у вас есть счётчик, который нужно увеличивать на единицу всякий раз, как только произойдёт определённое событие. Используя RAII, можно настроить счётчик так, чтобы он автоматически обновлялся, а затем, по окончании использования сообщал своё значение. Тогда вам не пришлось бы писать код, который бы каждый раз выполнял эту операцию. В таком случае вы не пропустите обновление, даже если ваш код усложнится.

Но даже с таким мощным инструментом мы видели в наших метриках нули – это какой-то нонсенс.

Наша конфигурация: MonitoredStateStoreIterStats и try_stream

Чтобы не сомневаться, что все данные учитываются, мы разработали специальный инструмент, MonitoredStateStoreIterStats. Он отслеживает метрики — например, сколько элементов данных было обработано. Он похож на скрупулезного счетовода, который каталогизирует все события, происходящие в потоке данных.

Всякий раз, когда фрагмент данных проходит через нашу систему, этот счётчик его добросовестно записывает. Ниже в упрощённом виде показано, как выглядит эта структура:

struct MonitoredStateStoreIterStats {
    total_items: usize,
    total_size: usize,
    storage_metrics: Arc<MonitoredStorageMetrics>,
}

Также мы воспользовались удобной возможностью Rust под названием try_stream — она упрощает работу с потоками данных. Она похожа на конвейер, эффективно транспортирующий данные. За этим конвейером расположена наша MonitoredStateStoreIterStats, ведущая счёт проезжающих по ленте элементов.

Вот как мы организовали поток:

pub struct MonitoredStateStoreIter<S> {
    inner: S,
    stats: MonitoredStateStoreIterStats,
}

impl<S: StateStoreIterItemStream> MonitoredStateStoreIter<S> {
    #[try_stream(ok = StateStoreIterItem, error = StorageError)]
    async fn into_stream_inner(mut self) {
        let inner = self.inner;
        futures::pin_mut!(inner);
        while let Some((key, value)) = inner.try_next().await? {
            self.stats.total_items += 1;
            self.stats.total_size += key.encoded_len() + value.len();
            yield (key, value);
        }
    }
}

Закончив работу, MonitoredStateStoreIterStats автоматически сообщает все собранные метрики в нашу систему мониторинга. Это делается в реализации Drop:

impl Drop for MonitoredStateStoreIterStats {
    fn drop(&mut self) {
        self.storage_metrics.iter_item.observe(self.total_items as f64);
        self.storage_metrics.iter_size.observe(self.total_size as f64);
    }
}

Расследование: углубляемся в код

Чтобы выяснить, что же идёт не так, мы переписали наш код в упрощённом виде. Как будто поставили небольшой лабораторный эксперимент.

Этот тест мы опробовали в Rust Playground. Вот код, при помощи которого мы воспроизвели проблему:

struct Stat {
    count: usize,
    vec: Vec<u8>,
}

impl Drop for Stat {
    fn drop(&mut self) {
        println!("count: {}", self.count);
    }
}

fn main() {
    let mut stat = Stat {
        count: 0,
        vec: Vec::new(),
    };
let mut f = move || {
        stat.count += 1;
        1
    };
    println!("num: {}", f());
}

В этом коде есть структура Stat, похожая на нашу MonitoredStateStoreIterStats. В ней есть поле count, в котором отслеживается, сколько раз произошло некоторое событие, а также поле vec, которым мы воспользуемся позже. Реализация Drop выводит на экран значение счётчика, актуальное на момент сброса Stat.

Главная функция создаёт экземпляр Stat и замыкание f. Предполагается, что это замыкание должно увеличивать count на 1 всякий раз, когда её вызывают.

Но, запустив этот код, получаем следующий вывод:

num: 1
count: 0

num: 1 демонстрирует, что было вызвано замыкание, и возвращает 1. Но count: 0 свидетельствует, что инкремент в поле count нашей структуры Stat не произошёл, хотя и должен был. Это было удивительно и в точности напоминало проблему, с которой мы столкнулись в продакшене: метрики не обновлялись.

Всё дело в замыканиях Rust

Оказалось, проблема связана с тем, как в Rust обрабатываются замыкания. Замыкания подобны мини-функциям, встроенным в код. Как правило, когда вы используете переменную внутри замыкания, эта переменная «захватывается». То есть, в зависимости от варианта использования, переменная либо попадает во владение, либо заимствуется…

Всё становится ещё интереснее с ключевым словом move. Если move используется до замыкания, то замыкание вынуждено брать во владение все те переменные, которые использует. Можно было бы подумать, что таким образом проблема решается, но на самом деле всё только усложняется.

Аналогия: разборчивый покупатель

Допустим, у вас есть список покупок (ваши данные) и вы посылаете кого-то (замыкание) с этим списком в магазин. В старых версиях Rust покупателю приходилось брать с собой весь список, хотя, купить из него требовалось только что-то одно.

Но, начиная с версии Rust 2021, покупатель стал работать эффективнее. Он берёт с собой только часть списка, ту, которая прямо требуется для решения поставленной задачи.

Обычно это плюс: экономится память и вообще всё работает быстрее. Но именно в нашем случае счётчик не получал той информации, которая требовалась ему для обновления метрик. Всё равно, словно покупатель видит в списке слово «яблоки», но игнорирует, сколько яблок нужно купить.

Вникаем в детали

Чтобы понять, почему так происходит, давайте рассмотрим несколько сценариев.

Перемещение структуры целиком

Сначала мы попытались модифицировать структуру так, чтобы в неё явно перемещалось stat целиком:

let mut f = move || {
    let mut stat = stat; // Явно переносим stat в замыкание
    stat.count += 1;
    1
};

После такого изменения получаем вывод:

num: 1
count: 1

Теперь count корректно увеличивается на 1. Как видите, если явно переместить stat в замыкание, то замыкание ею полностью завладевает. Любые изменения, вносимые в замыкание, затрагивают исходную stat.

Использование других полей

Далее мы попытались воспользоваться другим поле vec в структуре stat внутри замыкания:

let mut f = move || {
    let _ = stat.vec.len(); // Используем поле vec
    stat.count += 1;
    1
};

Удивительно, но вывод остался корректным:

num: 1
count: 1

Таким образом, если мы воспользуемся таким полем как vec, которое не реализует Copy, замыканию ничего не остаётся, кроме как завладеть всей stat целиком. Оно не сможет просто скопировать поле count.

Проблема с типами Copy

Такая проблема возникает, когда замыкание использует только поля, реализующие Copy. Таково, например, наше поле count, имеющее размер usize. В данном случае, даже с move, замыкание не завладевает всей структурой. Оно лишь копирует поля Copy.

Мы предположили: что ж, пусть stat реализует Drop и потому не поддаётся частичному перемещению. Но, если внутри замыкания move используются только поля типа Copy, относящиеся к stat, эти поля и будем копировать в замыкания. Любые изменения, вносимые в эти поля в пределах замыкания, затрагивают только скопированные значения, а оригинальные поля остаются нетронутыми.

Подтверждение нашей теории

Чтобы убедиться, что наша теория верна, мы решили заглянуть под капот и изучить код, сгенерированный компилятором Rust. Мы воспользовались инструментом под названием MIR (среднеуровневое промежуточное представление), чтобы посмотреть, чем компилятор занимается за кулисами.

По поводу проблемного кода MIR показал, что замыкание захватывает только поле usize (stat.count), а не всю структуру stat:

_3 = {closure@src/main.rs:19:17: 19:24} { stat: (_1.0: usize) };

Правда, в изменённом коде, где мы явно перемещали stat, MIR показывал, что stat целиком перемещается в замыкание:

_3 = {closure@src/main.rs:19:17: 19:24} { stat: move _1 };

Тем самым наша гипотеза подтвердилась. В проблематичном коде замыкание не завладевает stat, а просто копирует поля Copy.

Решение

Оказалось, что исправить проблему удивительно просто. Всего лишь требовалось явно сообщить замыканию, что оно должно завладевать всем объектом stats в нашей функции into_stream_inner.

Вот так исходно выглядел код, с которым возникала проблема:

#[try_stream(ok = StateStoreIterItem, error = StorageError)]
async fn into_stream_inner(mut self) {
    let inner = self.inner;
    ...
    self.stats.total_items += 1;
    self.stats.total_size += key.encoded_len() + value.len();
    ...
}

А вот исправленный код:

#[try_stream(ok = StateStoreIterItem, error = StorageError)]
async fn into_stream_inner(self) {
    let inner = self.inner;
    let mut stats = self.stats; // Take ownership of self.stats
    ...
    stats.total_items += 1;
    stats.total_size += key.encoded_len() + value.len();
    ...
}

Добавив let mut stats = self.stats;, мы вынуждаем замыкание завладеть объектом stats. Теперь при инкременте stats.total_items и stats.total_size мы изменяем сам объект stats, а не его копию.

Более широкий контекст: что это значит для Rust-разработчиков

Наше расследование показало одну тонкую деталь, которая ранее касалась поведения замыканий move в Rust: они захватывали минимально необходимую часть структуры. Притом, насколько это эффективно, такая практика может приводить к неожиданным результатам, особенно, если не знать всей её специфики. Это наиболее касается работы с типами Copy внутри структур, реализующих Drop.

В нашем случае из-за такого поведения терялись метрики. Дополнительно усложняет ситуацию инкапсулированный макрос try_stream, внутри которого применяются замыкания. Из-за него возникает ситуация, в которой, несмотря на наше намерение передать владение параметрами в поток, такая передача на самом деле не происходит.

Мы выявили две основные области, которые можно было бы доработать:

1. Инкапсуляция макросов: макрос try_stream должен более явно описывать, как в нём реализуется владение. Это помогло бы предотвратить схожие проблемы в случаях, когда разработчик предполагает, что перемещение в поток происходит move — а на самом деле нет.
2. Поведение замыканий в Rust: Притом, что поведение Rust технически корректно, в нём легко запутаться. Мы написали об этой проблеме в сообщество Rust, чтобы обсудить возможные улучшения, помочь прояснить код или документацию.

Заключение: сделанные выводы

Этот опыт лишний раз подчеркнул, как важно понимать нюансы владения в Rust и поведения замыканий в этом языке. Также стало понятнее, насколько важно тщ0ательно прорабатывать API макросов во избежание неоднозначности.

Для нас это был не просто очередной исправленный баг. Мы узнали много нового, что позволило нам повысить надёжность нашей системы. Мы не просто поправили в ней механизм сбора метрик, но и подробнее разобрались в Rust, что в будущем помогло нам писать более надёжный код.

Надеемся, что, поделившись этим опытом, мы принесли пользу и всему сообществу Rust. Хорошо, если другие не попадут в такие ситуации при работе с замыканиями, семантикой перемещения и инкапсуляцией макросов. Если вам интересно подробнее узнать, как мы разрабатываем на Rust потоковую базу данных нового поколения, загляните в наш репозиторий на GitHub!