Портируем декодер AV1 с С на Rust для повышения быстродействия и безопасности

AV1 становится всё более значимым видеоформатом, которому требуется безопасный и производительный декодер. Исходя из этой идеи, мы в тандеме с командой из Immutant создали rav1d, портировав на Rust написанный на С декодер dav1d. Перед вами первая из двух статей, посвящённых решению этой задачи.

— Джош Аас, глава проекта Prossimo организации ISRG

В современном мире программного обеспечения парсинг сложных данных в плане безопасности является одной из наиболее критических операций. Браузерам нужно декодировать недоверенный ввод аудио/видео, кодируемый в реальном времени в крайне сложные форматы. В итоге при декодировании баги, связанные с безопасностью памяти, становятся не только чудовищными, но и частыми. Например, исследователи реализаций декодера H.264, продемонстрировали, что они являются опасными источниками багов. AV1 тоже является сложным и широко используемым форматом. И чтобы избежать уязвимостей в активно используемых программах вроде браузеров, нам нужна безопасная по памяти и производительная реализация парсинга этого формата.

Для создания этого быстрого и безопасного декодера AV1, который получил название rav1d, мы портировали существующую библиотеку декодирования dav1d с её нативного языка C на Rust. Наша реализация полностью совместима с API dav1d. Парсеры форматов, которые были написаны на небезопасном С, теперь стали безопасными по памяти реализациями на Rust.

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

В отношении реализации rav1d мы поставили перед собой такие цели:

• Перенести код dav1d на Rust для повышения безопасности памяти.
• Добиться полной совместимости с API языка C.
• Добиться производительности, сравнимой с реализацией на С.
• Повторно использовать код из dav1d, чтобы упростить его частую синхронизацию.
• Обеспечить поддержку X86-64 и ARM64.

Наш подход к миграции

Написать высокопроизводительный и полноценный декодер AV1 с нуля сложно. Для этого нужно хорошо разбираться в стандарте AV1 и предметно понимать, как лучше реализовать формат кодека, чтобы он получился и быстрым, и совместимым.

Реализация dav1d, развиваемая сообществами VideoLAN и FFmpeg при спонсорской поддержке Alliance for Open Media, разрабатывается уже 6 лет. Она содержит около 50 К строк С и 250 К строк ассемблера. Это зрелая, быстрая и широко распространённая реализация. И вместо того, чтобы пытаться заново создать аналогичный декодер на Rust, мы предпочли перенести на этот язык существующий код dav1d.

Мы хотим предоставить совместимый с С API, чтобы сделать портирование кода в новую библиотеку на Rust более плавным. Мы также хотим повторно использовать код ассемблера dav1d для сохранения быстродействия. Эти ограничения в плане совместимости означают, что мы должны частично сохранить внутреннюю и внешнюю структуру данных dav1d. Кроме того, нам нужно вызывать процедуры ассемблера в то же время, когда это делает dav1d. С функциональной точки зрения это требует реализации декодирования практически тем же образом, каким оно реализовано в dav1d.

Мы могли бы вручную переписать dav1d на Rust по одной функции или модулю за раз. Однако, учитывая необходимость сохранения совместимости, это бы стало очень утомительной задачей. Чтобы достичь точки, где мы могли бы вносить во внутренние структуры данных сквозные изменения, необходимые для повышения безопасности памяти, потребуется переписать значительную часть базы кода. Вместо этого мы решили использовать c2rust, чтобы изначально транспилировать код С в равнозначный, но небезопасный код Rust. Это позволило нам начать переписывание с полностью рабочей и легко совместимой базы кода Rust, не внося в процессе новые логические баги. В итоге основная часть работы заключалась в ручном рефакторинге и переписывании небезопасного кода Rust в его безопасную, идеоматичную форму.

Транспиляция в небезопасный Rust с последующим переписыванием обеспечила для нашего проекта два важных преимущества: 1) начало с полностью рабочей реализации на Rust позволило нам тщательно протестировать функциональность декодирования, производя поэтапный рефакторинг; 2) транспиляция сложной логики декодирования сократила потребность в предметном знании спецификации AV1.

В результате оказалось очень полезным проводить полноценное тестирование непрерывной интеграции с самого начала, параллельно переписывая и дорабатывая код Rust. Мы могли вносить сквозные изменения в базу кода и проводить существующие тесты dav1d при каждом коммите. Между статической проверкой компилятора Rust и интеграционными тестами всего декодера мы проводили сравнительно меньше времени за отладкой, чем провели бы в случае реализации декодера с нуля. Основную часть команды этого проекта составляли эксперты по системному программированию и Rust, но у них не было столь важного опыта работы с кодеками AV. Наш эксперт по кодекам, Фрэнк Боссен, предоставил бесценные рекомендации, но в основную работу вовлечён не был.

Когда мы впервые приступили к реализации проекта, то предположили, что он займёт около 7 месяцев. Но в итоге стало ясно, что придётся вложить в него намного больше ручного труда, чем предполагалось. В общей сложности наша команда из 3 разработчиков затратила более 20 человеко-месяцев усилий. Процесс переписывания, особенно в попытке сохранить производительность, оказался намного более трудозатратным, чем ожидалось. Мы столкнулись с рядом проблем, связанных со своеобразностью базы кода dav1d и инструментом транспиляции c2rust. Например, структура кода dav1d привела к значительному дублированию трансилированного кода Rust для 8- и 16-битных изображений, который нам пришлось вручную унифицировать и удалять повторы.

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

Сложности

Мы столкнулись с различными проблемами, связанными с расхождениями между паттернами С и безопасного Rust. Управление жизненным циклом требовало подробного понимания существующей базы кода, но в итоге жизненные циклы и заимствование стали не самыми серьёзными проблемами. Безопасность потоков в Rust, которая усложняет обмен изменяемыми данными между рабочими потоками, никак не сочеталась с моделью потоков dav1d, где потоки неявно обмениваются практически всеми данными. Кроме того, ощутимыми источниками затруднений стали владение памятью и указатели буферов, а также объединения и прочие небезопасные паттерны C.

▍ Потоки

В dav1d используется пул рабочих потоков, которые конкурентно выполняют задачи, не зависящие друг от друга. Тем не менее эти задачи воздействуют на общие структуры в глобальном контексте и контексте фрейма, в том числе имеют общий изменяемый доступ к одним и тем же буферам. Например, в листинге 1а показан фрагмент корневой структуры, к которой должны обращаться все потоки. Каждый обрабатываемый фрейм содержит Dav1dFrameContext, который является общим для работающих с ним потоков.

Каждый объект Dav1dTaskContext используется только одним потоком, но в С каждый из них посредством корневого Dav1dContext доступен для всех других потоков. Наконец, этот корневой контекст содержит множество других полей состояний, некоторые из которых должны быть изменяемы основным потоком и считываемы в рабочих потоках.

Листинг 1a. Фрагмент из структуры корневого контекста:

struct Dav1dContext  {
    Dav1dFrameContext *fc;
    unsigned n_fc;

    Dav1dTaskContext *tc;
    unsigned n_tc;

    struct Dav1dTileGroup *tile;
    int n_tile_data_alloc;
    int n_tile_data;
    int n_tiles;

    // ...
}

Rust требует, чтобы все данные, совместно используемые потоками, были Sync. Это означает, что к ним должны иметь возможность безопасно обращаться конкурентно несколько потоков. Нам нужно обмениваться заимствуемым корневым контекстом между всеми потоками, поэтому все данные в этом контексте должны быть иммутабельны. Чтобы сделать возможным изменение общих данных, мы должны ввести блокировки, которые позволят сохранить безопасность потоков в среде выполнения. Для этого мы добавили Mutex и RwLock как необходимые. Если предположить, что изначальный код С не имеет состояний data race (в dav1d мы их не наблюдали), за эти новые блокировки никогда не должно возникать соперничества. Мы активно использовали Mutex::try_lock() и RwLock::try_read() / RwLock::try_write(), чтобы проверить в среде выполнения, может ли поток безопасно обращаться к данным, не внося задержек из-за ожидания освобождения блокировки.

Листинг 1b. Соответствующая выдержка контекста из rav1d:

pub struct Rav1dContext {
    pub(crate) state: Mutex<Rav1dState>,
    pub(crate) fc: Box<[Rav1dFrameContext]>,
    pub(crate) tc: Box<[Rav1dContextTaskThread]>,

    // ...
}

pub struct Rav1dState {
    pub(crate) tiles: Vec<rav1dTileGroup>,
    pub(crate) n_tiles: c_int,

    // ...
}

Как показывает листинг 1b, нам пришлось перестроить структуры dav1d, чтобы они лучше вписывались в модель потокобезопасности Rust. Мы отрефакторили изменяемое состояние в новой структуре Rav1dState и обернули его в мьютекс. Также стоит отметить, что tc больше не содержит локальных для каждого потока данных, а только хэндл потока и метаданные Sync для их координации. Все локальные данные потоков из Dav1dTaskContext теперь управляются каждым рабочим потоком независимо, так что ему не обязательно быть Sync.

Добавление дополнительных блокировок позволяет обработать случай, в котором только одному потоку нужно изменять конкретное поле структуры. Версия dav1d во многих сценариях опирается на конкурентный, но не пересекающийся доступ к одному буферу. Один поток должен считывать или записывать из одного участка буфера, в то время как другой поток обращается к другому его участку, не связанному с первым. Этот паттерн, хоть и лишён на практике состояний data race, не сопоставляется чётко с безопасными идиомами Rust. В Rust мы бы сначала разбили буфер на непересекающиеся части, после чего распределили их между разными потоками для обработки.

Такой паттерн требует знания точного размера каждого буфера данных заранее, чтобы иметь возможность правильно распределить эти участки для постановки задач потокам. В случае AV1 разделение буфера оказалось бы крайне сложным, поскольку оно не статично и даже не непрерывно. Для сохранения N-мерных массивов, таких как ndarray, которые бы позволили разделить эти буферы, существуют крейты. Но в этом случае для правильной разбивки буферов нам бы потребовалось понимание точных паттернов доступа всех задач во всех буферах. А это, в свою очередь, потребовало бы фундаментального перестраивания планировки задач rav1d.

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

Однако отслеживание каждого заимствования для релизных сборок слишком затратно, поэтому в них вся структура DisjointMut является бесплатной обёрткой вокруг внутреннего буфера. При доступе к буферу по-прежнему выполняется проверка на выход за границы, так что мы сохраняем пространственную безопасность, потенциально компрометируя безопасность потоков. Все буферы DisjointMut в rav1d представляют примитивные данные, поэтому в худшем случае этот паттерн может лишь внести неопределённость, если доступ окажется некорректно разделён.

В листинге 2 показан фрагмент структуры, которая совместно используется всеми рабочими потоками. Несколько потоков конкурентно изменяют различные блоки в поле b, поэтому мы обернули этот вектор в DisjointMut, чтобы обеспечить возможность конкурентного доступа.

Листинг 2a. Фрагмент из структуры контекста фрейма C:

struct Dav1dFrameContext_frame_thread  {
    // ...

    // Индексируется с помощью t->by * f->b4_stride + t->bx
    Av1Block *b;
    int16_t *cbi; /* bits 0-4: txtp, bits 5-15: eob */

    // ...
}  frame_thread;

Листинг 2b. Эквивалент той же структуры с листинга 2а на Rust:

pub struct Rav1dFrameContextFrameThread {
    // ...

    // Индексируется с помощью `t.b.y * f.b4_stride + t.b.x`.
    pub b: DisjointMut<Vec<Av1Block>>,

    pub cbi: Vec<RelaxedAtomic<CodedBlockInfo>>,

    // ...
}

Там, где это возможно, вместо добавления блокировки мы использовали атомарные типы. Мы опираемся на код, уже избегающий логических состояний data race, и атомарные примитивные типы обеспечивают формальную безопасность потоков. Мы не требовали конкретного атомарного упорядочивания памяти, поскольку предполагаем, что операции записи в общие поля не склонны к состояниям гонки, поэтому использовали слабое (relaxed) упорядочивание.

На интересующих нас платформах естественным образом выровненные операции загрузки и сохранения уже являются атомарными, поэтому relaxed-операции упорядочивания атомиков в Rust сводятся к тем же операциям с памятью в С без дополнительных издержек¹. Мы не могли использовать несвободные атомарные операции или методы fetch+update, поскольку они сводятся к сложным и более медленным инструкциям. Мы добавили обёртку RelaxedAtomic, чтобы упростить использование этих атомарных полей и исключить применение неэффективных паттернов. Мы также использовали крейт atomig, чтобы сделать простые структуры примитивного размера и перечисления атомарными.

В целом мы выяснили, что модель безопасности потоков Rust крайне строга. Если бы мы писали этот декодер с нуля, то спроектировали бы более разделённый и отчётливый обмен данными между потоками. Тем не менее мы смогли, по сути, сохранить производительность без внесения значительных изменений в существующую логику, используя новые структуры данных вроде DisjointMut и RelaxedAtomic, которые по-прежнему дают нам желаемые гарантии безопасности при слабой защите от состояния data race.

▍ Самореферентные структуры

Указатели на одну структуру или переключающиеся рекурсивно между разными являются типичным паттерном в С, который сложно воссоздать в безопасном Rust. Проблемные указатели, которые мы встретили при портировании dav1d, преимущественно относятся к одной из двух категорий: курсоры, отслеживающие позиции в буфере, и ссылки между структурами контекста.

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

Нам пришлось распутать структуры контекста dav1d, удалив из дочерних структур указатели на их контейнеры и затем передав дополнительные ссылки на структуры в виде параметров функций. Например, мы добавили в decode_tile_sbrow ссылочные параметры Rav1dContext и Rav1dFrameData, поскольку из структуры контекста задачи нам эти указатели пришлось удалить.

Листинг 3. Отделение структуры контекста:

// Оригинальная функция C
int dav1d_decode_tile_sbrow(Dav1dTaskContext *const t) {
  const Dav1dFrameContext *const f = t->f;
  Dav1dTileState *const ts = t->ts;
  const Dav1dContext *const c = f->c;
  // ...
}

// Её безопасная версия на Rust
pub(crate) fn rav1d_decode_tile_sbrow(
    c: &Rav1dContext,
    t: &mut Rav1dTaskContext,
    f: &Rav1dFrameData,
) -> Result<(), ()> {
    let ts = &f.ts[t.ts];
}

▍ Объединения

В dav1d используются неразмеченные объединения из языка C. В случаях, где дополнительное поле применяется в качестве тега, мы переписали эти объединения в безопасные размеченные перечисления Rust. Однако в С дискриминант некоторых объединений был неявным. Например, то, какой вариант объединения должен использоваться, определяла стадия задачи, сохранённая совершенно в другой структуре контекста. Для этих случаев вместо того, чтобы добавлять ненужный тег и изменять представление структуры и размер, мы предпочли использовать крейт zerocopy, чтобы в среде выполнения повторно интерпретировать те же байты как два разных типа. Это было единственное возможное решение, поскольку объединения состояли полностью из примитивных типов без заполнения. Предоставляемые zerocopy трейты обуславливают использование этого инварианта и делают возможным доступ к содержимому объединения без затрат и не требуя явного тега. И хотя этот паттерн менее идиоматичен, в нескольких случаях мы нашли его применение необходимым для повышения быстродействия и совместимости.

Заключение

Оказались ли транспиляция и переписывание оправданы? Мы считаем, что да, по крайней мере, для проекта rav1d. Переписывание декодера AV1 с нуля привело бы к внесению всевозможных багов и проблемам совместимости. Мы выяснили, что, несмотря на проблемы с потоками и заимствованием, переписывание существующего кода С на безопасном и производительном Rust возможно. Наша реализация rav1d на данный момент примерно на 6% медленнее текущей реализации dav1d на С. Детали оптимизации быстродействия rav1d мы более подробно разберём в следующей статье.

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

Благодарим Amazon Web Services, Sovereign Tech Fund и Alpha-Omega за поддержку этого проекта. Если вы хотите больше узнать о rav1d или начать его использовать, добро пожаловать на GitHub.

1. Единственные издержки связаны с невозможностью совместить, например, 2 последовательных, выровненных загрузки AtomicU8 в одно хранилище AtomicU16, что должно прозрачно реализовываться для u8s и u16s. Для отдельных полей это не является проблемой, но всё же вносит дополнительные издержки при работе с массивами и срезами.