Tarantool и кодогенерация на Lua

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

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

При работе с компилируемыми языками мы переносим часть работы из рантайма на стадию компиляции и получаем программу, которая способна очень быстро выполнять нашу специализированную задачу. Много материала про кодогенерацию уже написано. Сразу вспоминается язык Go, хотя это, конечно, не единственный пример — в каждом языке возникают задачи, связанные с маршаллингом, парсингом и т.д. Я буду рассказывать про то, что ближе мне — про использование кодогенерации при разработке приложений на Tarantool. При этом в отличие от компилируемых языков нам не потребуется производить перекомпиляцию кода в случае изменения каких-либо параметров — всё будет происходить прямо в рантайме.

Немного про Tarantool и LuaJIT

Tarantool — это платформа для in-memory вычислений — флакон, объединяющий сервер приложений и базу данных. Сам Tarantool написан на языке С, но пользователь может работать с ним с помощью языка Lua. А если совсем точно, то одной из его реализаций — LuaJIT — не с просто интерпретатором, а ещё и с поддержкой и JIT-компиляции. И часто при работе возникают задачи по трансформации сущностей при записи в базу или после извлечения из неё, а также их валидации на соответствие схеме, заданной пользователем. Типичный подход для решения этой и схожих задач — написание функций для преобразования данных. Эти функции не привязаны к конкретной схеме и зачастую представляют из себя набор замыканий. Однако не стоит забывать, что мы работаем с LuaJIT языком, который способен компилировать и достаточно быстро выполнять "горячие" участки кода.

Но, к сожалению, не всё подряд может быть скомпилировано, у платформы есть ряд ограничений — это так называемые NYI (Not yet implemented) функции. Кроме того, работа с данными активно использует дополнительные структуры — массивы и хэш-мапы. В Lua они представлены общим типом данных — "table" (таблица). Перед нами две основные проблемы — использование части функций серьезно влияет на производительность, а избыточное использование вспомогательных структур приводит к излишней нагрузке на GC, с которым у и Lua 5.1, и у LuaJIT проблемы. Поэтому задача — написание кода, который сможет быть скомпилирован LuaJIT, и будет приводить к минимально возможному количеству аллокаций.

К реальным задачам

Данный подход мы будем разбирать на реальном примере, на примере модуля CRUD. Задача данного модуля — это упрощение работы с шардированными данными. То есть данные распределены между несколькими стораджами (инстансами Tarantool, хранящими данные), и мы, обращаясь к ним через роутер (по сути, клиент), не хотим задумываться, на каком именно из стораджей лежат интересующие нас данные, а просто указываем условие поиска, и модуль возвращает нам уже готовые данные. Немного про хранение. Tarantool хранит данные в спейсах (spaces) — аналог таблиц в реляционных БД. Единица хранения — кортеж (tuple) — массив заданных нами значений. При этом нам привычно работать именно с Lua-таблицами — обращаться к полю по названию, а не по номеру в кортеже. В качестве аналогии можно привести формат JSON. Обычно именно в таком формате поступают данные из внешних систем — которые затем парсятся в Lua-таблицы, "сплющиваются" и сохраняются в базу. Соответственно типичными для тарантула операциями являются так называемый "флаттенинг" (flatten) и "анфлаттенинг" (unflatten) — получение из луа-таблицы плоского тапла и наоборот. И в частном случае пользователь может написать руками все эти операции.

-- Создаем space - аналог таблицы в реляционных БД
box.schema.space.create('data')

-- Создаем первичный ключ
box.space.data:create_index('primary_key')

-- Попробуем вставить в наш space следующий объект
object = { id = 1, key = "key", value = "value" }

-- Выполняем "сплющивание" объекта - flatten
tuple = {object["id"], object["key"], object["value"]}

-- Единицей хранение в Tarantool является tuple - кортеж из значений
box.space.data:insert(tuple)

-- После сохранения мы можем достать наш объект по первичному ключу
tuple = box.space.data:get({1})

-- Преобразуем объект в исходное состояние - unflatten
object = {
   id = tuple[1],
   key = tuple[2],
   value = tuple[3],
}

Здесь мы явно захардкодили порядок полей в спейсе. Однако в общем случае схема задается извне некоторым форматом, и мы пишем простенькие функции, которые занимаются трансформацией объекта в соответствии с этим форматом. Модуль CRUD, как и сам Tarantool, имеет функцию replace — она точно также вставляет кортеж в базу. Для упрощения жизни пользователям была также добавлена функция replace_object — которая принимает объект, преобразует в плоский вид в соответствии с форматом спейса, а затем уже сохраняет.

Ближе к коду и измерению производительности

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

-- test_data.lua
-- Формат - 8 строковых полей + bucket_id
-- (специальное поле, необходимое при шардировании данных).
local format = {
    {name = 'field1', type = 'string', is_nullable = false},
    {name = 'field2', type = 'string', is_nullable = false},
    {name = 'field3', type = 'string', is_nullable = false},
    {name = 'field4', type = 'string', is_nullable = false},
    {name = 'field5', type = 'string', is_nullable = false},
    {name = 'field6', type = 'string', is_nullable = false},
    {name = 'field7', type = 'string', is_nullable = false},
    {name = 'field8', type = 'string', is_nullable = false},
    {name = 'bucket_id', type = 'unsigned', is_nullable = false},
}

-- Объект необходимого формата
local data = {
    field1 = 'string1',
    field2 = 'string2',
    field3 = 'string3',
    field4 = 'string4',
    field5 = 'string5',
    field6 = 'string6',
    field7 = 'string7',
    field8 = 'string8',
    bucket_id = nil,
}

return {
    format = format,
    data = data,
}

Функция, замеряющая время выполнения нашего кода.

-- bench.lua
-- Замеряем, сколько времени займет 1 миллион итераций
local clock = require('clock')

local count = 1e6

local function run(f, ...)
    local start = clock.time()
    for _ = 1, count do
        f(...)
    end
    return clock.time() - start
end

return {
    run = run,
}

И входная точка, исполняемый файл, в который поочередно можно будет добавлять тесты:

#!/usr/bin/env tarantool

-- init.lua
local bench = require('bench')
local test_data = require('test_data')

-- Это наш первый тест
local naive = require('naive')
local res = bench.run(naive.flatten, test_data.data, test_data.format, 1)
print(string.format('Naive result: %0.3f s', res))

-- После добавления нужного модуля, мы раскомментируем каждый фрагмент.
-- local code_gen_v1 = require('code_gen_v1')
-- local res = bench.run(code_gen_v1.flatten, test_data.data, test_data.format, 1)
-- print(string.format('code_gen_v1 result: %0.3f s', res))

-- local code_gen_v2 = require('code_gen_v2')
-- local res = bench.run(code_gen_v2.flatten, test_data.data, test_data.format, 1)
-- print(string.format('code_gen_v2 result: %0.3f s', res))

Давайте рассмотрим, как раньше выполнялось данное преобразование — самый наивный подход:

-- naive.lua
local system_fields = { bucket_id = true }

local function flatten(object, space_format, bucket_id)
    if object == nil then return nil end

    local tuple = {}

    local fieldnames = {}

    for fieldno, field_format in ipairs(space_format) do
        local fieldname = field_format.name
        local value = object[fieldname]

        if not system_fields[fieldname] then
            if not field_format.is_nullable and value == nil then
                return nil, string.format("Field %q isn't nullable", fieldname)
            end
        end

        if bucket_id ~= nil and fieldname == 'bucket_id' then
            value = bucket_id
        end

        tuple[fieldno] = value
        fieldnames[fieldname] = true
    end

    for fieldname in pairs(object) do
        if not fieldnames[fieldname] then
            return nil, string.format("Unknown field %q is specified", fieldname)
        end
    end

    return tuple
end

return {
    flatten = flatten,
}

Пример слегка упрощен. Но стоит заметить несколько вещей:

• При каждом запросе мы пересоздаем таблицу "fieldnames", содержащую поля формата. Необходима для проверки, что наш объект не содержит лишних полей (и без кодогенерации её следовало бы как-то закэшировать).

• Обходим весь объект в соответствии с форматом. При этом формат нам известен и меняется достаточно редко. Это одна из предпосылок для использования кодогенерации.

Запускаем:

➜  tarantool init.lua 
Naive result: 1.109 s

На моём ноутбуке этот тест выполнился за 1 секунду.

Теперь попробуем написать функцию, которая развернет нам цикл. Объект же продолжит обрабатываться в соответствии с форматом.

-- code_gen_v1.lua

-- Небольшой хелпер для работы со строками
local function append(lines, s, ...)
    table.insert(lines, string.format(s, ...))
end

-- Кэш, где ключ - таблица с "форматом", а значение - функция флаттенинга.
-- Для простоты считаем, что формат не меняется, не занимаемся инвалидацией кэша.
local cache = {}

local function flatten(object, space_format, bucket_id)
    -- В случае если функция уже сгенерирована,
    -- берем её из кэша. Иначе приступаем к кодогенерации.
    local fun = cache[space_format]
    if fun ~= nil then
        return fun(object, bucket_id)
    end
    -- Будем "готовить" наш код построчно и сохранять в массив lines.
    local lines = {}
    append(lines, 'local object, bucket_id = ...')

    append(lines, 'local result = {}')
    for i, field in ipairs(space_format) do
        if field.name ~= 'bucket_id' then
            append(lines, 'result[%d] = object[%q]', i, field.name)
        else
            append(lines, 'result[%d] = bucket_id', i)
        end
    end
    append(lines, 'return result')
    -- Конкатенируем элементы массива, чтобы получить полный текст функции.
    local code = table.concat(lines, '\n')
    
    -- Раскомментриуйте, чтобы увидеть результат
    -- print(code)
    
    -- С помощью функции "load" преобразуем текст функции в саму функцию
    fun = assert(load(code))

    cache[space_format] = fun
    return fun(object, bucket_id)
end

return {
    flatten = flatten,
}

В результате выполнения получим следующий код:

local object, bucket_id = ...
local result = {}
result[1] = object["field1"]
result[2] = object["field2"]
result[3] = object["field3"]
result[4] = object["field4"]
result[5] = object["field5"]
result[6] = object["field6"]
result[7] = object["field7"]
result[8] = object["field8"]
result[9] = bucket_id
return result

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

Что еще мы можем улучшить? В самом начале сгенерированного кода мы создаем таблицу result и постепенно её заполняем. Такой подход приводит к неоднократным реаллокациям, что плохо и довольно бессмысленно — ведь размер таблицы известен заранее. Давайте учтём это и поменяем строку append(lines, 'local result = {}') на append(lines, 'local result = {%s}', string.rep('box.NULL,', #space_format)). Так мы сразу создадим массив нужного нам размера — local result = {box.NULL, ..., box.NULL}. Запуск бенчмарка выдает 0.2 секунды.

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

-- code_gen_v2.lua
local function append(lines, s, ...)
    table.insert(lines, string.format(s, ...))
end

local cache = setmetatable({}, {__mode = 'k'})

local function flatten(object, space_format, bucket_id)
    local fun = cache[space_format]
    if fun ~= nil then
        return fun(object, bucket_id)
    end

    local lines = {}
    append(lines, 'local object, bucket_id = ...')

    append(lines, 'for k in pairs(object) do')
    append(lines, '    if fieldmap[k] == nil then')
    append(lines, '        return nil, format(\'Unknown field %%q is specified\', k)')
    append(lines, '    end')
    append(lines, 'end')

    local len = #space_format
    append(lines, 'local result = {%s}', string.rep('NULL,', len))

    local fieldmap = {}

    for i, field in ipairs(space_format) do
        fieldmap[field.name] = true
        if field.name ~= 'bucket_id' then
            if field.is_nullable ~= true then
                append(lines, 'if object[%q] == nil then', field.name)
                append(lines, '    return nil, \'Field %q isn\\\'t nullable\'', field.name)
                append(lines, 'end')
            end
            append(lines, 'result[%d] = object[%q]', i, field.name)
        else
            append(lines, 'if bucket_id ~= nil then')
            append(lines, '    result[%d] = bucket_id', i, field.name)
            append(lines, 'else')
            append(lines, '    result[%d] = object[%q]', i, field.name)
            append(lines, 'end')
        end
    end
    append(lines, 'return result')

    local code = table.concat(lines, '\n')
    local env = {
        pairs = pairs,
        format = string.format,
        fieldmap = fieldmap,
        NULL = box.NULL,
    }
    fun = assert(load(code, '@flatten', 't', env))

    cache[space_format] = fun
    return fun(object, bucket_id)
end

return {
    flatten = flatten,
}

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

Бенчмарк показал 0.3 секунды.

➜  tarantool init.lua 
Naive result: 1.109 s
code_gen_v1 result: 0.210 s
code_gen_v2 result: 0.299 s

Также стоит отметить, что у функции load появились дополнительные аргументы, а именно chunkname — название нашей функции (может быть полезным при отладке), mode — t — мы создаем функцию на основе обычного текста, а не байткода и env — окружение, доступное внутри нашей функции. На последний аргумент стоит обратить особое внимание. Кроме возможности создавать удобные песочницы для выполнения пользовательского кода (обычно не давать доступа к "опасным" функциям), данная опция позволяет передавать в глобальное окружение нужные нам функции и аргументы. В нашем случае это pairs, format, fieldmap и NULL. Отдельно стоит отметить, что load — это функция из Lua 5.2 — расширение LuaJIT. Тот, кто работает с чистым Lua 5.1, может использовать функции loadstring для создания функции и setfenv для установки окружения у этой функции.

Не обязательно использовать кодогенерацию для каждой функции. Например, если бы я бы захотел добавить ещё и валидацию типов полей, то мне бы не потребовалось генерировать функции для проверки (is_number, is_string, ...) — их достаточно просто передать через окружение.

Небольшой пример:

local function is_string(value)
    return type(value) == 'string'
end

-- Функции is_string нет в языке Lua,
-- но с помощью окружения мы можем добавить в нужные нам функции
-- и убрать лишние.
local code = [[
local value = ...
local result = {NULL}

if not is_string(value) then
    error("value is not a string")
end

result[1] = value

return result
]]

local fun = load(code, '@test', 't', {
    error = error,
    -- Функция is_string будет доступна внутри
    -- загружаемого нами кода
    is_string = is_string,
    NULL = box.NULL,
})

Как в будущем всё не сломать

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

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

Более качественная оценка результатов

В статье я привел не так много результатов тестирования — просто сравнил несколько цифр: время выполнения программы до и после. При этом я говорил, что мы уменьшили количество аллокаций и написали код, пригодный для компиляции LuaJIT'ом. Но как это можно проверить, как в этом можно убедиться?

Не хочется превращать статью в гайд о том, как профилировать код на Tarantool. Но всё-таки мы слегка затронем эту тему.

Во-первых, это memory profiler, который появился в версии 2.7.1 — инструмент, который покажет в каких именно местах и в каких количествах выделяется/реаллоцируется память. Как по мне, вывод довольно удобен — а в будущем станет ещё удобнее. Воспользовавшись этим инструментом, можно показать количественную разницу между кодом до и кодом после. В нашем случае мы получили бы вывод в формате @<filename>:<function_line>, line <line where event was detected>: <number of events> <allocated> <freed>. Для наглядности напротив некоторых строк я помещу фрагменты кода, которые находятся на этих строках:

Для кода "до" (naive.lua):

ALLOCATIONS
INTERNAL: 3999953	360000380	0
@../naive.lua:4, line 26: 1000038	384003936	0           // fieldnames[fieldname] = true
@../naive.lua:4, line 7: 1000000	64000000	0           // local tuple = {}
@../naive.lua:4, line 9: 1000000	64000000	0           // local fieldnames = {}
@../naive.lua:4, line 25: 16	384	0
@../naive.lua:4, line 0: 4	672	0

REALLOCATIONS
INTERNAL: 1999982	112000560	64000288
	Overrides:
		@../naive.lua:4, line 0
		@../naive.lua:4, line 25
		INTERNAL

@../naive.lua:4, line 25: 1000022	136001232	72000704
	Overrides:
		@../naive.lua:4, line 25
		INTERNAL


DEALLOCATIONS
INTERNAL: 5953572	0	784628243
	Overrides:
		@../naive.lua:4, line 0
		@../naive.lua:4, line 25
		@../naive.lua:4, line 26
		@../naive.lua:4, line 7
		@../naive.lua:4, line 9
		INTERNAL

@../naive.lua:4, line 26: 1000022	0	192001584
	Overrides:
		@../naive.lua:4, line 26
		INTERNAL

Для кода "после" (code_gen_v2.lua):

ALLOCATIONS
@flatten:0, line 7: 1000000	144000000	0 // local result = {NULL,NULL,NULL,NULL,NULL,NULL,NULL,NULL,NULL,}
@../code_gen_v3.lua:7, line 55: 1	48	0

REALLOCATIONS
INTERNAL: 5	1984	3968
	Overrides:
		INTERNAL


DEALLOCATIONS
INTERNAL: 974298	0	140298062
	Overrides:
		@flatten:0, line 7

Во-вторых, сам LuaJIT поставляется с профилировщиком — require('jit.p')

Для кода "до":

52%  ../naive.lua:11  // for fieldno, field_format in ipairs(space_format) do
30%  ../naive.lua:26  // fieldnames[fieldname] = true
12%  ../naive.lua:9   // local fieldnames = {}

Для кода "после":

36%  flatten:3  // if fieldmap[k] == nil then
36%  flatten:7  // local result = {NULL,NULL,NULL,NULL,NULL,NULL,NULL,NULL,NULL,}
11%  ../code_gen_v3.lua:9 // выбираем значение из кэша
 4%  flatten:39
 4%  flatten:2
 4%  ../code_gen_v3.lua:8
 4%  flatten:8
 4%  ../code_gen_v3.lua:10

А также для тех, кто хочет копнуть совсем глубоко, есть возможность дампа байткода, который LuaJIT генерирует и выполняет — require('jit.dump')

Заключение

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

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