Оптимизация, которая невозможна в Rust

в 9:01, , рубрики: dst, german string, Rust, umbra string, unsafe

Поскольку я изучаю системы баз данных для получения степени магистра в Германии, статья с названием «Why German Strings are Everywhere» сразу привлекла мое внимание. Мне было интересно узнать, что речь идет о структуре данных, описанной в статье «Umbra: A Disk‑Based System with In‑Memory Performance», о которой я узнал из другой статьи о её движке хранилища — «LeanStore: In‑Memory Data Management Beyond Main Memory». Еще более интересным было то, что она используется во многих других решениях для хранения данных, таких как DuckDB, Apache Arrow и Polars.

This string implementation also allows for the very important “short string optimization”: A short enough string can be stored “in place,” i.e., we set a specific bit in the capacity field, and the remainder of capacity, as well as size and ptr, become the string itself. This way, we save on allocating a buffer and a pointer dereference each time we access the string. An optimization that's impossible in Rust, by the way;).

Если я правильно понял данное утверждение, то оно объективно некорректно, поскольку уже существует большое количество пакетов, поддерживающих эту опцию, например compact_str, smartstring и smol_str. Более того, polars, пример из статьи, написан на Rust и реализует Umbra‑style строки. Я стал жертвой охоты на нердов и начал реализовывать German string, чтобы доказать себе, что это возможно.

Охота на нердов.by xkcd licensed under CC BY-NC 2.5

Охота на нердов.
by xkcd licensed under CC BY-NC 2.5

В этой статье я опишу, как я реализовывал German string и с какими трудностями столкнулся в Rust. В особенности я рассмотрю, как добавить общее владение для подобной структуры данных. Предоставление уникального владения достаточно тривиально и описано в хорошо написанном туториале в Rustonomicon, где рассматривается реализация Vec<T>, которая не сильно отличается от String. Но сначала поговорим о концепции German string.

Разве строковых типов данных недостаточно в Rust?

The term “German-style string” was coined by Andy Pavlo in one of his lectures on Advanced Database Systems. “German string,” “German-style string,” and “Umbra-style string” all refer to the same concept.

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

  • String — закодированная в UTF-8 расширяемая строка.

  • CString — владеемая, C‑совместимая, нуль‑терминированная строка без нулевых байтов в середине.

  • OSString — владеемая, изменяемая строка, нативная для платформы, дешево преобразуемая в строки Rust.

  • PathBuf — тонкая обертка над OSString, представляющая собой владеемый изменяемый путь.

Строки в Rust

String — один из наиболее используемых типов в Rust, и поэтому он будет предметом нашей дискуссии. Под капотом String представляет собой Vec<u8>, то есть последовательность байт. Также эта последовательность должна быть закодирована в UTF-8, поэтому один символ не всегда будет использовать один байт и может занимать от одного до четырех байт. Исключая некоторые детали, реализация String выглядит следующим образом:

pub struct String {
    vec: Vec<u8>,
}

pub struct Vec<T> {
    buf: RawVec<T>,
    len: usize,
}

pub(crate) struct RawVec<T, A: Allocator = Global> {
    ptr: Unique<T>,
    cap: Cap,
    alloc: A,
}

pub struct Unique<T: ?Sized> {
    pointer: NonNull<T>,
    _marker: PhantomData<T>,
}

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

  • len — длина строки в байтах.

  • cap — размер буфера на куче.

  • ptr — указатель на буфер на куче.

Схема строк в Rust

Схема строк в Rust

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

German strings

В теории строка может содержать бесконечно много видов текста бесконечных размеров. Тем не менее, в реальных примерах, особенно в базах данных, можно выделить несколько общих характеристик:

  • Чаще всего строки короткие.

  • Большая часть строк не меняется.

  • Лишь малая часть строк используется для сравнения или сортировки.

Основываясь на этих наблюдениях, были изобретены German string для использования в таких специфичных случаях с целью ускорения обработки строк. Основная идея заключается в «оптимизации коротких строк». Чтобы избежать дорогостоящего процесса выделения памяти, для достаточно коротких строк вместо выделения буфера на куче можно хранить байты непосредственно на стеке, используя для этого cap и ptr. Отсутствие буфера также убирает необходимость разыменования указателя при операциях с короткими строками.

Под капотом German string занимает всего 16 байт на стеке и имеет два представления в зависимости от длины строки. Также она неизменяема, что убирает необходимость отслеживания размера буфера.

Короткие строки

Схема German string для коротких строк

Схема German string для коротких строк

Строки размером до 12 байт инлайнятся и хранятся непосредственно на стеке. Из 16 байт 4 используются для хранения длины, а 12 — для хранения содержимого.

Длинные строки

Схема German string для длинных строк

Схема German string для длинных строк

Строки длиной более 12 байт хранятся на куче как обычные String и содержат указатель на буфер на куче. Однако есть несколько отличий:

  • Только 4 байта используются для хранения длины, что ограничивает размер строки до 4 GB. Это приемлемо, поскольку большинство строк в реальном мире короткие.

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

another one

Теперь, когда мы знаем, как должна выглядеть German string и с какой целью, давайте разбираться, как мы можем реализовать это в Rust. Для начала, наша цель состоит в том, чтобы реализовать владеемую неизменяемую строку, закодированную в UTF-8, и поддерживающую ранее описанные оптимизации. Более того, буфер на куче будет поддерживать атомарный подсчет ссылок, позволяя использовать его в разных потоках.

Есть два способа структурирования нашего типа. Первый подход, который приходит в голову, выглядит следующим образом:

pub struct UmbraString {
    len: u32,
    data: Data,
}

union Data {
    buf: [u8; 12],
    ptr: ManuallyDrop<SharedDynBytes>,
}

struct SharedDynBytes {
    prefix: [u8; 4],
    ptr: NonNull<SharedDynBytesInner>
    phantom: PhantomData<SharedDynBytesInner>
}

Согласно Rust layout rules, SharedDynBytes имеет размер 16 байт и выравнивание 8 байт, что приводит к тому, что размер UmbraString становится равен 24 байтам с выравниванием 8 байт. Это отличается от желаемого нами расположения. Чтобы уменьшить размер до 16 байт, можно снизить выравнивание до 4 байт и задать корректный размер с помощью #[repr(packed(4))]. Однако линтер Rust Clippy не будет доволен, когда мы будем ссылаться на поле ptr в SharedDynBytes, поскольку это может создать ссылку на невыровненный адрес, что является неопределенным поведением в Rust. И хотя нам никогда не понадобится создавать подобную ссылку, мы не сможем использовать часть API на указателях, которые было бы полезно иметь.

Схема UmbraString для длинных строк с полями, расширенными до их выравнивания.

Схема UmbraString для длинных строк с полями, расширенными до их выравнивания.

Гораздо неприятнее то, что это усложняет реализацию сравнения и упорядочивания строк. При сравнении или сортировке строк мы хотим взаимодействовать в первую очередь с префиксом. С данным расположением нам нужно сначала проверить длину строки, чтобы определить, что представляет собой Data — buf или ptr, прежде чем перейти к префиксу. Чтобы обойти это, мы можем пометить всё как #[repr(C)] и полагаться на инвариант, что префикс всегда является первым полем в SharedDynBytes. В результате мы можем интерпретировать данные как buf, не проверяя длину, и всё будет нормально, если мы обращаемся только к первым 4 байтам. Более корректный подход представлен ниже; он не полагается на подобный инвариант и освобождает нас от необходимости беспокоиться о порядке полей, предоставляя структуры для различных стратегий владения выделенного буфера.

#[repr(C)]
pub struct UmbraString {
    len: u32,
    prefix: [u8; 4],
    trailing: Trailing,
}

#[repr(C)]
union Trailing {
    buf: [u8; 8],
    ptr: ManuallyDrop<SharedDynBytes>,
}

#[repr(C)]
struct SharedDynBytes {
    ptr: NonNull<SharedDynBytesInner>
    phantom: PhantomData<SharedDynBytesInner>
}

С точки зрения памяти нет никакой разницы по сравнению с версией #[repr(C, packed(4))]. Вместо использования 12-байтового объединения наше новое объединение Trailing занимает всего 8 байт и содержит либо последние 8 байт строки, либо указатель на буфер на куче. Также нам необходимо добавить аннотацию #[repr(C)] к нашей структуре, чтобы порядок полей соответствовал их объявлению. Для UmbraString обязательно, чтобы после поля prefix следовало поле trailing. Префикс при этом убирается из объединения для более простого взаимодействия и позволяет правилам распределения Rust работать эффективно без излишнего вмешательства с нашей стороны. Представление префикса в виде массива из четырех байт также улучшает производительность, поскольку сравнение массивов фиксированного размера работает значительно быстрее, чем сравнение слайсов.

Байты с атомарным подсчетом ссылок

SharedDynBytes, слайс с атомарным подсчетом ссылок, ведет себя как Arc. Из этого следует логичный вопрос: почему бы просто не использовать Arc<[u8]>?

Типы с динамическим размером

Большая часть типов в Rust имеет статически известный размер и выравнивание. Типы с динамическим размером (DST) являются исключением из этого правила, и в Rust доступны два основных DST:

  • Объекты трейтов: dyn Trait

  • Слайсы: [T], str и подобные

Поскольку DST не имеют статически известного размера, они не могут размещаться на стеке и могут быть доступны только по указателю. Каждый указатель на DST является fat pointer, состоящим из самого указателя и дополнительной информации о данных, на которые он указывает. В случае с слайсом, который нас интересует, этой информацией является длина. Это приводит к тому, что Arc<[T]> имеет размер 16 байт вместо 8, из‑за чего UmbraString занимает 24 байта вместо 16.

Схема UmbraString с использованием Arc<[u8]>

Схема UmbraString с использованием Arc<[u8]>

Мы знаем, что UmbraString уже отслеживает длину строки, и дополнительный размер от толстого указателя можно сократить. Нам нужно преобразовать толстый указатель в тонкий, что не поддерживается Arc<[u8]>. На данный момент нет способа создать DST напрямую, и мы можем только создать инстанс кастомного DST, сделав его дженериком и принудительно уменьшив размер.

layout

После описания проблемы наши структуры, описывающие байты с подсчетом ссылок, выглядят следующим образом, немного отличаясь от ранее представленного примера.

#[repr(C)]
pub struct SharedDynBytes {
    ptr: NonNull<SharedDynBytesInner<[u8; 0]>>,
    phantom: PhantomData<SharedDynBytesInner<[u8; 0]>>,
}

#[repr(C)]
struct SharedDynBytesInner<T: ?Sized> {
    count: AtomicUsize,
    data: T,
}

Структура может быть превращена в DST добавлением DST в качестве последнего поля. Здесь SharedDynBytesInner содержит счетчик ссылок и универсальное значение типа без размера. SharedDynBytesInner всегда располагается на куче, и атомарный счетчик отслеживает количество созданных ссылок на данный момент. Выделяя память как для счетчика, так и для данных, мы можем избежать лишнего выделения памяти и дополнительной абстракции. Хотя он и является дженериком, нас интересуют две разновидности типа с явной структурой и выравниванием полей:

  • SharedDynBytesInner<[u8]> имеет динамический размер, и указатели на него являются толстым указателем.

  • SharedDynBytesInner<[u8; 0]> имеет статический размер, и указатель на него является тонким указателем.

Представление SharedDynBytes в памяти

Представление SharedDynBytes в памяти

Один интересный трюк, который мы можем использовать, это преобразование между толстым указателем типа *mut SharedDynBytesInner<[u8]> и тонким указателем типа *mut SharedDynBytesInner<[u8; 0]>. Это преобразование позволяет нам добавлять или удалять информацию о длине слайса из указателя в зависимости от того, как мы хотим его использовать. Преобразование толстого указателя в тонкий осуществляется легко — достаточно одного каста.

let fat_ptr: *mut SharedDynBytesInner<[u8]> = _;
let thin_ptr = fat_ptr as *mut SharedDynBytesInner<[u8; 0]>;

А вот обратное преобразование требует немного больше усилий. Вместо прямого каста мы сначала должны преобразовать тонкий указатель в указатель на слайс заданной длины, что позволяет Rust включить информацию о длине в указатель. Лишь после этого мы можем преобразовать указатель на слайс в наш кастомный DST.

let thin_ptr: *mut SharedDynBytesInner<[u8; 0]> = _;
let fake_slice = std::ptr::slice_from_raw_parts_mut(thin_ptr.cast::<u8>(), len);
let fat_ptr = fake_slice as *mut SharedDynBytesInner<[u8]>;

Выделение памяти

Теперь, когда мы можем преобразовывать толстый указатель в тонкий и обратно, мы можем начать выделять память под него. Если мы не можем привести тип со статическим размером к типу с динамическим размером, то мы можем создать значение DST только путем ручного выделения памяти. Перед этим нам необходимо задать представление для выделяемого значения.

fn shared_dyn_bytes_inner_layout(len: usize) -> Layout {
    Layout::new::<SharedDynBytesInner<()>>()
        .extend(array_layout::<u8>(len))
        .expect("A valid layout for a SharedDynBytesInner")
        .0
        .pad_to_align()
}

Представление в памяти организовано следующим образом: сначала создается представление для SharedDynBytesInner<()>, где поля не занимают места в памяти. Затем мы расширяем это представление представлением массива байт, получая SharedDynBytesInner<[u8]>. Получив его, мы можем выделить явно известное количество памяти для счетчика ссылок и массива байт заданной длины.

fn from(bytes: &[u8]) -> Self {
    let ptr = if bytes.is_empty() {
        NonNull::dangling()
    } else {
        let layout = shared_dyn_bytes_inner_layout(bytes.len());
        let nullable = unsafe { std::alloc::alloc(layout) };
        let nullable_fat_ptr = SharedDynBytesInner::<[u8]>::cast(nullable, bytes.len());
        let Some(fat_ptr) = NonNull::new(nullable_fat_ptr) else {
            std::alloc::handle_alloc_error(layout)
        };
        unsafe {
            let inner = &mut (*fat_ptr.as_ptr());
            std::ptr::write(&mut inner.count, AtomicUsize::new(1));
            std::ptr::copy_nonoverlapping(bytes.as_ptr(), inner.data.as_mut_ptr(), bytes.len());
        }
        fat_ptr.cast()
    };
    Self {
        ptr,
        phantom: PhantomData,
    }
}

Придерживаемся ленивого подхода: если слайс пустой, память выделяться не будет. Если он не пустой, выполняем следующие действия:

  1. Выделяем память, используя заданное представление.

  2. Преобразуем полученный указатель в толстый указатель на наш DST.

  3. Копируем данные из полученного слайса в поле data.

  4. Преобразуем толстый указатель обратно в тонкий и сохраняем его.

Освобождение памяти

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

unsafe fn dealloc_unchecked(&self, len: usize) {
    if len > 0 {
        let inner = unsafe { &*self.ptr.as_ptr() };
        if inner.count.fetch_sub(1, atomic::Ordering::Release) == 1 {
            inner.count.load(atomic::Ordering::Acquire);
            unsafe {
                std::alloc::dealloc(
                    self.ptr.as_ptr().cast::<u8>(),
                    shared_dyn_bytes_inner_layout(len),
                );
            };
        }
    }
}

При вызове метод в первую очередь проверяет, не равна ли переданная длина нулю. Если она равна нулю, значит, память не выделялась, и ничего делать не нужно. Когда память выделялась, мы уменьшаем счетчик ссылок и проверяем, являемся ли мы единственным владельцем. Если мы — единственный владелец, освобождаем память в соответствии с заданной моделью.

Клонирование

Аналогичным образом, мы не можем реализовать Clone для нашего типа и вынуждены полагаться на переданное пользователем значение длины, чтобы определить, выделялась ли память. Клонирование не требует копирования данных; оно просто увеличивает счетчик ссылок на 1.

unsafe fn clone_unchecked(&self, len: usize) -> Self {
    let ptr = if len == 0 {
        NonNull::dangling()
    } else {
        let inner = unsafe { &*self.ptr.as_ptr() };
        inner.count.fetch_add(1, atomic::Ordering::Relaxed);
        self.ptr
    };

    Self {
        ptr,
        phantom: PhantomData,
    }
}

Получение слайса

Получение слайса из хранимого массива также полагается на предоставленную пользователем длину и выполняется в два этапа:

  1. Преобразование тонкого указателя в толстый.

  2. Создание слайса с использованием указателя на поле data и предоставленной пользователем длины.

#[inline]
unsafe fn as_bytes_unchecked(&self, len: usize) -> &[u8] {
    if len == 0 {
        Default::default()
    } else {
        let fat_ptr = SharedDynBytesInner::<[u8]>::cast(self.ptr.as_ptr().cast::<u8>(), len);
        let ptr = unsafe { (*fat_ptr).data.as_ptr() };
        unsafe { std::slice::from_raw_parts(ptr, len) }
    }
}

Собираем все вместе

Теперь, когда у нас есть все необходимое, мы можем приступить к реализации German string. Хотя для полноценной строки требуется много компонентов, мы покажем только базовый функционал.

Выделение памяти

Создание строки — простой процесс. Мы копируем префикс, обрабатывая два случая:

  1. Если строка занимает не больше 12 байт, суффикс инлайнится.

  2. В противном случае выделяется новый SharedDynBytes с содержимым строки.

fn new(s: &str) -> Result<UmbraString, Error> {
    let bytes = s.as_bytes();
    let len = bytes.len();
    if len > u32::MAX as usize {
        return Err(Error::TooLong);
    }
    let mut prefix = [0u8; 4];
    let trailing = if len <= 12 {
        let mut buf = [0u8; 8];
        if len <= 4 {
            prefix[..len].copy_from_slice(&bytes[..len]);
        } else {
            prefix.copy_from_slice(&bytes[..4]);
            buf[..len - 4].copy_from_slice(&bytes[4..]);
        }
        Trailing { buf }
    } else {
        prefix.copy_from_slice(&bytes[..4]);
        Trailing {
            ptr: ManuallyDrop::new(SharedDynBytes::from(bytes)),
        }
    };
    #[allow(clippy::cast_possible_truncation)]
    Ok(Self {
        len: len as u32,
        prefix,
        trailing,
    })
}

Drop

Освобождение памяти может быть выполнено путем реализации Drop для UmbraString, поскольку у нас есть вся необходимая информация для этого. Мы проверяем, нужно ли освобождать память, основываясь на длине. Если данные находятся на куче, мы используем метод dealloc_unchecked, показанный ранее.

impl Drop for UmbraString {
    fn drop(&mut self) {
        if self.len > 12 {
            unsafe {
                self.trailing.ptr.dealloc_unchecked(len);
            }
        }
    }
}

Clone

Реализация Clone для UmbraString происходит аналогичным образом. Мы копируем байты в новый экземпляр, когда строка заинлайнена. В противном случае мы копируем префикс и используем метод clone_unchecked у SharedDynBytes.

impl<B> Clone for UmbraString<B>
where
    B: DynBytes,
{
    fn clone(&self) -> Self {
        let trailing = if self.len <= 12 {
            unsafe {
                Trailing {
                    buf: self.trailing.buf,
                }
            }
        } else {
            unsafe {
                Trailing {
                    ptr: ManuallyDrop::new(self.trailing.ptr.clone_unchecked(self.len as usize)),
                }
            }
        };
        Self {
            len: self.len,
            prefix: self.prefix,
            trailing,
        }
    }
}

PartialEq

Сравнение начинается с первых 8 байт, состоящих из длины и префикса. Эта проверка обычно проходит успешно для большинства строк, и метод сразу возвращает результат. Таким образом, большая часть сравнений не требует обращения к остатку строки на куче. Сравнения между массивами фиксированного размера работают очень быстро и хорошо оптимизированы. Если обе строки короткие и заинлайнены, мы можем сравнить последние 8 байт без разыменования указателя. Разыменование указателей происходит только внутри вызова suffix(), когда строки не заинлайнены.

impl PartialEq<UmbraString> for UmbraString {
    fn eq(&self, other: &UmbraString) -> bool {
        let lhs_first_qword = std::ptr::from_ref(self).cast::<u64>();
        let rhs_first_qword = std::ptr::from_ref(other).cast::<u64>();
        if unsafe { *lhs_first_qword != *rhs_first_qword } {
            return false;
        }
        if self.len <= 12 {
            return unsafe { self.trailing.buf == other.trailing.buf };
        }
        self.suffix() == other.suffix()
    }
}

PartialOrd

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

impl PartialOrd<UmbraString> for UmbraString {
    fn partial_cmp(&self, other: &UmbraString) -> Option<cmp::Ordering> {
        let ordering = Ord::cmp(&self.prefix, &other.prefix).then_with(|| {
            if self.len <= 4 && other.len <= 4 {
                return Ord::cmp(&self.len, &other.len);
            }
            if self.len <= 12 && other.len <= 12 {
                let ordering = unsafe { Ord::cmp(&self.trailing.buf, &other.trailing.buf) };
                return ordering.then_with(|| Ord::cmp(&self.len, &other.len));
            }
            Ord::cmp(self.suffix(), other.suffix())
        });
        Some(ordering)
    }
}

Послесловие

Я сделал невозможное возможным! Реализация German string в Rust требует знания модели памяти Rust и расположения типов в памяти, но сама по себе достаточно проста. Сложности возникают из‑за того, что я решил включить поддержку общего владения для данного типа для переиспользования буферов. В процессе я обнаружил несколько интересных деталей про Rust:

  • Не вся информация о типе явно дается в Rust. Посмотрев на определение Arc<T>, я бы никогда не осознал, что он может занимать 16 байт вместо 8.

  • DST не очень хорошо поддерживаются и могут быть созданы только через ручное выделение памяти или принудительное изменение размера.

  • Мы можем преобразовывать толстые указатели в тонкие и обратно, используя трюк со слайсами.

На этом всё, спасибо за прочтение! Надеюсь, что, как и я, вы узнали что‑то новое! Увидеть больше кода вы можете на следующей странице: https://crates.io/crates/strumbra.

Автор: MrPizzly

Источник

* - обязательные к заполнению поля


https://ajax.googleapis.com/ajax/libs/jquery/3.4.1/jquery.min.js