Всем привет! Меня зовут Роман Кабаев, я инженер в отделе разработки инструментов тестирования компании YADRO. Вместе с коллегами мы разрабатываем собственную тест-менеджмент систему с открытым исходным кодом TestY.
На этапе запуска TestY в качестве фреймворка для разработки мы выбрали Django, так как он позволяет в максимально короткие сроки реализовать MVP. Однако развивать такой продукт — добавлять фичи, наращивать число пользователей и объем хранимых данных в системе — бывает сложно.
Мы действительно быстро запустили MVP, перевезли данные из TestRail с помощью плагинов, и команды тестирования YADRO уже более года пользуются системой. Но есть одно «но»: пользовательские сценарии разных команд сильно отличаются. Так, добавление в систему более полумиллиона тестов привело к просадке скорости работы определенных эндпоинтов, завязанных на древовидных структурах.
Спойлер: камнем преткновения для нас стали CPU-bound задачи с большим количеством данных, о том, как я это выяснил, расскажу ниже. Изучив, как можно ускорить выполнение таких задач в Python, я протестировал несколько решений и нашел оптимальное. Если вы разрабатываете веб-приложение на Django или Python и так же, как я, хотите ускорить работу сервиса, читайте эту статью.
В качестве альтернативы Rust можно посмотреть и на Cython, и на чистый C. Возможно, стоит подумать о микросервисной архитектуре и отдавать тяжелые с точки зрения CPU задачи микросервису, написанному, например, на Go. Я же последовал модным веяниям и решил изучить, как написать Python-пакет с использованием Rust.
Отправная точка — 10,93 секунды
Изначально система отображения данных в виде древовидной структуры нас более чем устраивала. Но из-за неожиданно большого количества данных в TestY стали видны узкие места. Например, интерфейс для создания плана тестирования, который можно увидеть на скриншоте ниже.
![](https://habrastorage.org/getpro/habr/upload_files/58f/ff4/64d/58fff464de4f8cdc00364c05c919f3b9.png)
Чтобы построить красивое дерево выбора, все тестовые сьюты для выбора и кейсы внутри них отдаются эндпоинтом сразу. А значит, если у вас 10 000 кейсов в одной из сьют, они подгружаются одновременно. Отсюда и такая огромная просадка по скорости отдачи данных.
Проблему можно исправить с помощью хорошего UX-дизайна и небольшого изменения архитектуры, но так как TestY — не единственный наш проект, на это не было ресурсов. Поэтому, чтобы сделать быстрый фикс без изменений UI, нам нужно было ускорить отдачу древовидных структур именно с серверной части приложения.
Первый проект, на котором мы обнаружили проблему, — один из самых важных у команды тестирования, поэтому для оценки производительности возьмем его:
16 419 тестовых кейсов.
445 тестовых сьют, разных степеней вложенности.
Размер payload (полезной нагрузки) с 2 МБ данных.
Для работы с базой данных используем Django ORM.
Запускаем это все в Docker-контейнере на MacBook Pro (M2) 16 ГБ ОЗУ.
Версия Python 3.9.13, операционная система контейнера Alpine.
Изначально у нас есть код с использованием Django и Django-rest-framework (DRF) следующего вида:
@cases_search_schema
@action(methods=[_GET], url_path='search', url_name='search', detail=False)
def cases_search(self, request):
cases = self.get_queryset()
cases = self.filter_queryset(cases)
suites_selector = TestSuiteSelector()
suites = TestSuiteSelector().suites_by_ids_list(
cases.values_list('pk', flat=True),
field_name='test_cases__pk',
)
suites_depth = suites_selector.get_max_level()
suites = (
suites
.get_ancestors(include_self=True)
.filter(parent=None)
.prefetch_related(
*suites_selector.suites_tree_prefetch_children(suites_depth),
*form_tree_prefetch_objects(
nested_prefetch_field='child_test_suites',
prefetch_field='test_cases',
tree_depth=suites_depth,
queryset=cases,
),
).annotate(**suites_selector.cases_count_annotation())
)
return Response(
data=TestSuiteTreeCasesSerializer(
suites,
context=self.get_serializer_context(),
many=True,
).data,
)
Производительность не очень впечатляющая, эндпоинт отдал данные за 10,93 секунды. Попробуем ускорить решение, чтобы добиться более высокой скорости работы.
Поиск узкого места
Чтобы решить какую-либо задачу, сначала нужно ее проанализировать и понять корень проблемы. Так как мы говорим про веб-приложение, первое, с чего я начал, — это посмотрел, как быстро работает запрос в базу данных с помощью django-debug-toolbar. Он съедает часть производительности системы, поэтому мы включаем его только единожды, чтобы посмотреть, как много времени у нас занимают запросы в базу данных. Все запросы в БД занимают 371 миллисекунду.
Упрощаем DRF-сериализаторы
Сериализаторы DRF очень тяжеловесны и содержат много скрытой логики и CPU-bound нагрузки. Попробуем их облегчить:
избавимся от ModelSerializer в пользу Serializer,
сделаем все поля read_only.
Получаем сериализаторы следующего вида:
from rest_framework import serializers
class TestSuiteTestSerializer(serializers.Serializer):
id = serializers.IntegerField(read_only=True)
name = serializers.CharField(read_only=True)
parent_id = serializers.IntegerField(read_only=True)
title = serializers.CharField(read_only=True)
test_cases = GenericIdNameSerializer(many=True, read_only=True)
parent = GenericIdNameSerializer(read_only=True)
children = serializers.SerializerMethodField()
def get_children(self, value):
return type(self)(
value.child_test_suites.all(),
many=True,
context=self.context,
).data
Получаем результат в 9,12 секунды. Такой медленный ответ пользователя не устроит, продолжим поиски.
Используем стороннюю библиотеку для перевода Python-объектов в JSON
Так как конвертирование Python-словаря в JSON — не самая «легкая» cpu-bound операция из-за рекурсивного обхода вложенных Python-объектов и манипуляций со строками, ищем альтернативу.
Вспоминаем про библиотеку Pydantic, которая часто используется для сериализации объектов в связке с FastAPI. Нам не нужен весь функционал Pydantic, интересует только сериализация в JSON, а для этого Pydantic использует Orjson — библиотеку, написанную на Rust как раз для ускорения сериализации в JSON.
Используем Orjson для сериализации в JSON и получаем результат в 8,32 секунды. Пользователь в ярости от такой нерасторопности системы: возможно, уже обновил страницу кулаком и не один раз повысил нагрузку на сервер. Результат — плохо всем.
Как и большинство веб-проектов на Django, для запуска в production-окружении мы используем application server gunicorn. Так как мы не можем обрабатывать один запрос более двух минут, чтобы не сломать систему и тем более виртуальную машину, на которой у нас развернуто приложение, мы выставляем timeout. Если запрос исполняется больше двух минут, gunicorn worker ликвидируется и выдает неприятную ошибку в системе мониторинга Sentry.
Стадия принятия
Каким бы прекрасным ни был Python, когда разработчики сталкиваются с ограничением производительности, то обращаются к более высокопроизводительным языкам, а затем делают биндинги в Python. Воспользуемся этим. Orjson использует Rust, и мы решили написать на этом языке свой пакет.
Так как сборка древовидных структур происходит по ID его родителя, вся нагрузка ложится на Python. Но мы можем получить наборы данных в плоской структуре — в виде списков словарей, а затем передать их в пакет, написанный на Rust, и обработать данные там.
Что мы делали раньше:
Сформировали запрос в базу данных с помощью django-orm.
Отдали набор данных QuerySet в DRF сериализатор.
Достали из него Python-словарь.
Конвертировали Python словарь в JSON (данный шаг скрыт под капотом у класса Response из библиотеки DRF).
Отдали ответ пользователю.
Как это выглядит теперь:
Получаем QuerySet всех тестовых кейсов, используем метод values (). Он вернет QuerySet, содержащий словари. Обернем все это в список — так мы получим список словарей, с которым проще всего будет работать в Rust. Таким же образом получаем все тестовые кейсы, в values указываем все нужные нам поля, чтобы не тянуть лишние данные из базы.
Ищем тестовые сьюты, к которым относятся эти тестовые кейсы, и получаем их предков (родители родителя и так до 0 уровня). Получаем плоский набор данных по принципу, описанному выше.
Передаем наборы данных в функцию из нашей Rust-библиотеки и производим необходимую обработку.
Возвращаем список словарей. Словари, в свою очередь, представляют иерархическое дерево.
Сериализуем полученный список в JSON с помощью Orjson.
Отдаем ответ пользователю.
Для написания пакета нам нужен сам Rust, с инструкцией по его установке можно ознакомится на официальном сайте. Также понадобится Python-пакет maturin.
Maturin создает проект с необходимым содержимым для того, чтобы мы начали писать свой Python-пакет на Rust:
![](https://habrastorage.org/getpro/habr/upload_files/904/d10/5f4/904d105f4a72ec95b88baa3c97bb3cbe.png)
Cargo.toml — конфигурационный файл пакетного менеджера Rust Cargo
pyproject.toml — конфигурация нашего Python-пакета,
lib.rs — точка входа Rust-кода.
Сначала напишем функцию, которую будет вызывать Python-код. Это функция с декоратором #[pyfunction]:
// src/serialization.rs
#[pyfunction]
pub fn serialize_tree(
py: Python,
data_set_object: TreeObject,
prefetch_objects: Vec<Prefetch>,
is_tree: bool,
) -> Vec<PythonInstance> {
let mut dict_map = BTreeMap::new();
let mut root_refs: Vec<PythonInstanceRef> = Vec::new();
for mut dict in data_set_object.instances {
for prefetch_object in &prefetch_objects {
dict.related_objects.insert(prefetch_object.group_key.clone(), vec![]);
}
dict_map.insert(dict.id, Rc::new(RefCell::new(dict)));
}
add_related_objects(py, prefetch_objects, &mut dict_map);
if !is_tree {
return dict_map.values().map(|obj| obj.borrow().clone()).collect();
}
for dict in dict_map.values() {
let borrowed_dict = dict.borrow();
if let Some(parent_id) = borrowed_dict.parent_id {
if let Some(parent) = dict_map.get(&parent_id) {
map_parent(parent, parent_id, dict, borrowed_dict, py)
}
} else {
root_refs.push(Rc::clone(dict))
}
}
root_refs.iter().map(|root| root.borrow().clone()).collect()
}
fn map_parent(
parent: &Rc<RefCell<PythonInstance>>,
parent_id: u32,
dict: &Rc<RefCell<PythonInstance>>,
borrowed_dict: Ref<PythonInstance>,
py: Python,
) {
let py_dict = borrowed_dict
.dict
.downcast::<PyDict>(py)
.expect("Could not downcast PyObject to PyDict");
let mut borrowed_parent = parent.borrow_mut();
borrowed_parent.children.push(Rc::clone(dict));
let name: String = borrowed_parent
.dict
.downcast::<PyDict>(py)
.expect("Could not cast PyAny to PyDict")
.get_item("name")
.expect("name was not found in PyDict")
.expect("unwrapping")
.extract()
.expect("Could not extract String from PyObject");
py_dict.set_item(
"parent",
Parent {
id: parent_id,
name,
}.into_py(py),
).expect("Could not set item in dict");
}
fn add_related_objects(
py: Python,
prefetch_objects: Vec<Prefetch>,
dict_map: &mut BTreeMap<u32, PythonInstanceRef>,
) {
for prefetch_object in prefetch_objects {
for instance in prefetch_object.instances {
let related_dict = instance
.downcast::<PyDict>(py)
.expect("Could not cast PyObject to PyDict");
let fk: u32 = related_dict
.get_item(&prefetch_object.fk_key)
.expect("Could not find fk by provided key from dict")
.expect("fk was not found")
.extract()
.expect("Could not convert PyObject to u32");
let relation_map = &mut dict_map.get_mut(&fk).expect("Could not get map").borrow_mut().related_objects;
if let Some(related_objects) = relation_map.get_mut(&prefetch_object.group_key) {
related_objects.push(instance);
}
}
}
}
Регистрируем ее в наш модуль в файле lib.rs:
// src/lib.rs
use pyo3::prelude::*;
mod serialization;
mod sql;
#[pymodule]
fn rusty(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(serialization::serialize_tree, m)?)?;
m.add_function(wrap_pyfunction!(sql::cases_search, m)?)?;
Ok(())
}
Не буду вдаваться в детали алгоритма построения дерева и группирования дочерних объектов на основных элементах, так как эта задача достаточно тривиальна. В нашем случае основные элементы — это тестовые сьюты и сценарии, которые в них содержатся.
Поговорим немного о преобразовании Python-объектов в Rust-объекты. Нам нужна какая-то структура, которая помогла бы преобразовывать динамически типизированные Python-объекты в статически типизированные Rust-объекты, поэтому напишем следующие структуры:
// src/serialization.rs
use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
use pyo3::{pyfunction, types::PyDict, FromPyObject, IntoPy, PyAny, PyObject, PyResult, Python};
type PythonInstanceRef = Rc<RefCell<PythonInstance>>;
#[derive(Clone, Debug)]
pub struct PythonInstance {
id: u32,
parent_id: Option<u32>,
dict: PyObject,
children: Vec<PythonInstanceRef>,
related_objects: BTreeMap<String, Vec<PyObject>>,
}
#[derive(Clone, Debug, FromPyObject)]
pub struct Prefetch {
group_key: String,
fk_key: String,
instances: Vec<PyObject>,
}
struct Parent {
id: u32,
name: String,
}
pub struct TreeObject {
instances: Vec<PythonInstance>,
}
Built-in типы PyO3 может преобразовывать и сам. Например, list[int]) преобразуется в Vec без явного вмешательства с нашей стороны. Однако нам нужно передать определенную метаинформацию, чтобы потенциально переиспользовать функцию из Python-Rust-пакета для обработки других структур, отличных от тестовых сьют с кейсами. А еще мы бы хотели работать с объектами Python как с не строго типизированными, потому что мы можем получить разное количество ключей в словарях или передать словари с другими полями. Поэтому у нас есть Python-объекты (дата-классы), которые несут определенную метаинформацию.
Объекты, которые мы хотим обработать, оставляем в формате PyObject. В процессе формирования иерархической структуры мы достаем из них только ключи, необходимые нам для формирования древовидных структур. Чтобы добиться этого, мы реализуем Rust traits из PyO3, такие как FromPyObject и IntoPy.
// src/serialization.rs
impl<'source> FromPyObject<'source> for TreeObject {
fn extract(obj: &'source PyAny) -> PyResult<TreeObject> {
let parent_key: Option<String> = obj
.getattr("parent_key")
.expect("Could not get parent key from python object")
.extract()?;
let pk_key: String = obj
.getattr("pk_key")
.expect("Could not_get pk key from python object")
.extract()?;
let mut instances = Vec::new();
let py_instances: Vec<PyObject> = obj
.getattr("instances")
.unwrap()
.extract()?;
for py_instance in py_instances {
let py_dict = py_instance.downcast::<PyDict>(obj.py())?;
let id = py_dict
.get_item(&pk_key)
.unwrap()
.unwrap()
.extract()?;
let mut parent_id = None;
if let Some(parent_id_from_py) = py_dict
.get_item(&parent_key)
.expect("Could not get element from python dict")
{
parent_id = parent_id_from_py
.extract()
.expect("Could not parse parent id value as u32");
}
instances.push(PythonInstance {
id,
dict: py_instance,
parent_id,
related_objects: BTreeMap::new(),
children: Vec::new(),
});
}
Ok(TreeObject { instances })
}
}
impl IntoPy<PyObject> for Parent {
fn into_py(self, py: Python) -> PyObject {
let dict = PyDict::new(py);
dict.set_item("name", self.name)
.expect("Could not set name for parent");
dict.set_item("id", self.id)
.expect("Could not set id for parent");
dict.into()
}
}
FromPyObject конвертирует Python-объект в Rust-структуру, IntoPy — это обратная операция. Чтобы с модулем было удобнее работать из основного приложения, нужно добавить тайп-хинты и объекты группировки. Последние подскажут пользователю, как работать с нашим модулем.
Так как мы хотим использовать и Python, и Rust в одном проекте, сделаем папку pysrc. Выкладка нашего модуля будет выглядеть так:
![](https://habrastorage.org/getpro/habr/upload_files/7f8/abd/2d6/7f8abd2d6f4affb953d6d3c60f4d6d3b.png)
В pyproject.toml, который сгенерировал пакет maturin, обозначим, что Python-исходник — это пункт python-source:
# Cargo.toml
[build-system]
requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[tool.maturin]
# "extension-module" tells pyo3 we want to build an extension module (skips linking against libpython.so)
features = ["pyo3/extension-module"]
python-source = "pysrc"
[project]
name = "TestY-rust"
requires-python = ">=3.8"
classifiers = [
"Programming Language :: Rust",
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Python :: Implementation :: PyPy",
]
dynamic = ["version"]
Важный момент: в init.py файле в pysrc/rusty нужно будет импортировать все из пакета rusty (тот пакет, который мы обозначили в lib.rs), иначе в пакете мы не найдем наши функции.
# pysrc/rusty/__init__.py
from typing import Any
from .constants import *
from .rusty import *
from .types import *
__all__ = (
"CASES_SEARCH_SUITE_FIELDS",
"CASES_SEARCH_CASE_FIELDS",
"CaseSearchQueryParams",
"Prefetch",
"DataSetObject",
"serialize_tree",
"cases_search",
)
В pysrc/rusty/types сделаем пару простых дата-классов, которые будут содержать необходимую для форматирования данных метаинформацию. Prefetch — это объекты, которые мы будем группировать на другие объекты, в нашем случае тестовые кейсы на сьюты. Нам нужен ключ, по которому мы будем складывать наши тестовые кейсы в сьюты. fk_key — это имя ключа родительского элемента, по нему будем искать нужный объект, instances — это наши тестовые кейсы.
DataSetObject — это наши cьюты, pk_key — это ключ первичного ключа, parent_key — ключ родительского элемента, instances — это сами сьюты:
# pysrc/rusty/types.py
from dataclasses import dataclass
from typing import Any, Optional
@dataclass
class Prefetch:
group_key: str
fk_key: str
instances: list[dict[str, Any]]
@dataclass
class DataSetObject:
parent_key: Optional[str]
pk_key: str
instances: list[dict[str, Any]]
И чтобы однозначно обозначить, что наши функции принимают, и чтобы IDE подсказывала нам, что эти функции существуют, сделаем .pyi-файл и положим его рядом с кодом на Python в pysrc/rusty:
# pysrc/rusty/__init__.pyi
def serialize_tree(
data_set_object: DataSetObject,
prefetch_objects: list[Prefetch],
is_tree: bool,
) -> list[dict[str, Any]]:
"""
Retrieve python objects retrieved from queries and build dict for json.
Args:
data_set_object: description type of provided data
prefetch_objects: description type of data that should be "prefetched" on main object
is_tree: boolean flag to build the main data set as a tree.
Returns:
list of dictionaries ready to be given to the final user.
"""
Так мы приходим к следующему коду в нашем view поиска тестовых кейсов:
@cases_search_schema
@action(methods=[_GET], url_path='search', url_name='search', detail=False)
def cases_search(self, request):
cases = self.get_queryset()
cases = self.filter_queryset(cases)
suites = list(
TestSuite.objects.filter(
id__in=cases.values_list('suite_id', flat=True),
).get_ancestors(
include_self=True,
).annotate(
title=F(_NAME),
).values(_ID, _NAME, 'title', 'parent_id').order_by('name'),
)
data = rusty.serialize_tree(
rusty.DataSetObject('parent_id', _ID, instances=suites),
prefetch_objects=[
rusty.Prefetch(
'test_cases',
'suite_id',
instances=list(cases.values(_ID, _NAME, 'suite_id', 'labels').order_by('name')),
),
],
is_tree=True,
)
data = orjson.dumps(data)
return HttpResponse(content=data, content_type='application/json')
Результатом работы становится ответ сервера в 300 мс — это отлично, если сравнивать с исходным результатом в 10,93 секунды
Но мы пойдем дальше и посмотрим, что будет, если сделать этот же запрос на самый нагруженный проект в системе. В проекте 221 203 тестовых кейса и 101 тестовый сьют. В результате получим ответ от сервера за 722 мс с пейлоадом в 22.3 МБ. Это более чем удовлетворительный результат, подходящий пользователям.
Раз мы удовлетворены перформансом эндпоинта, поговорим о подводных камнях при установке пакета.
Установка Rust/Python-модуля
Так как мы хотим, чтобы наш пакет работал на большом количестве конфигураций систем, было бы хорошо включить в наш Dockerfile Rust и собирать пакет на лету. Но это увеличивает и так не маленькое время установки зависимостей проекта и запуска контейнеров. Примерно на 4−8 минут, так как сначала нужно установить Rust, а затем уже скомпилировать нашу зависимость. Как разработчиков системы это нас очень печалит.
Вспоминаем, что есть предсобранные файлы wheels или же .whl Python-пакеты, но есть одно «но». Даже у разработчиков в нашей команде разные конфигурации систем MacOS (arm/x86), Debian (x86) и Windows (x86), то есть мы не можем собрать .whl-файл для какой-то конкретной системы, даже для нашего внутреннего использования. Также надо помнить, что TestY — это проект с открытым исходным кодом, а значит, он должен собираться на большом количестве разных систем.
Что мы можем сделать, чтобы уменьшить шансы ошибки сборки проекта? Собрать .whl-файлы для самых часто встречающихся конфигураций машин. Поэтому мы решили автоматизировать сборку .whl-файлов и их загрузку на pypi. И когда пользователь будет делать pip-install, pypi сам будет определять необходимую пользователю зависимость. Для этого мы воспользуемся GitHub Workflows, в частности maturin-action. Как только мы запушим новый tag в наш репозиторий на GitHub, maturin-action соберет необходимые нам .whl-файлы и зальет их на pypi, если не найдет идентичные .whl-файлы.
Пример кода пайплана для сборки .whl-файлов для различных дистрибутивов Linux на разных архитектурах:
# .github/workflows/CI.yml
name: CI
on:
push:
branches:
- main
- master
tags:
- '*'
pull_request:
workflow_dispatch:
permissions:
contents: read
jobs:
linux:
runs-on: ubuntu-latest
strategy:
matrix:
arch: [ i686, x86_64, aarch64 ]
image: [ manylinux2014, musllinux_1_1 ]
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: Build wheels - manylinux
uses: PyO3/maturin-action@v1
with:
target: ${{ matrix.arch }}
args: --release --out dist -i python3.9
sccache: 'true'
manylinux: ${{ matrix.image }}
- name: Upload wheels
uses: actions/upload-artifact@v4
with:
name: wheels-linux-${{ matrix.image }}-${{ matrix.arch }}
path: dist
Часть пайплайна, которая загружает собранные .whl-файлы в pypi, опущена, также как и сборка под другие операционные системы. Теперь наш пакет полностью готов к использованию.
Выводы по задаче и альтернативные решения
В ходе разработки я понял, что интеграция Rust-модуля в Python — не самая тривиальная задача. Хоть на эту тему есть много гайдов от других инженеров, всегда возникают подводные камни, которые никто не описывал. Как по мне, документация Rust-пакетов, в частности PyO3, могла бы быть лучше, но все знают, что наличие документации в целом — уже чудо.
Такая оптимизация — это обходной путь для нехватки человеческого ресурса и времени на разработку. Тем не менее, мы всеми силами стараемся сделать систему лучше, не отказываемся от Django и не хотим переписывать всю систему на Rust, так как это сильно замедлит скорость разработки. Мы используем все возможные средства для улучшения пользовательского опыта, не замыкаемся в стандартных решениях для Django или ограничениях языка.
Если тоже хотите поучаствовать в улучшении тест-менеджмент системы TestY, пишите на почту testy@yadro.com. А скачать репозиторий с версией TestY 1.3, о которой мы недавно писали, можно по ссылке.