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

Поскольку я изучаю системы баз данных для получения степени магистра в Германии, статья с названием «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, чтобы доказать себе, что это возможно.

В этой статье я опишу, как я реализовывал 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 — указатель на буфер на куче.

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

German strings

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

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

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

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

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

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

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

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

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

Строки длиной более 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 на указателях, которые было бы полезно иметь.

Гораздо неприятнее то, что это усложняет реализацию сравнения и упорядочивания строк. При сравнении или сортировке строк мы хотим взаимодействовать в первую очередь с префиксом. С данным расположением нам нужно сначала проверить длину строки, чтобы определить, что представляет собой 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]>. На данный момент нет способа создать 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]> имеет статический размер, и указатель на него является тонким указателем.

Один интересный трюк, который мы можем использовать, это преобразование между толстым указателем типа *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.