Всем привет! Меня зовут Роман Кабаев, я инженер в отделе разработки инструментов тестирования компании YADRO. Вместе с коллегами мы разрабатываем собственную тест-менеджмент систему с открытым исходным кодом TestY.
На этапе запуска TestY в качестве фреймворка для разработки мы выбрали Django, так как он позволяет в максимально короткие сроки реализовать MVP. Однако развивать такой продукт — добавлять фичи, наращивать число пользователей и объем хранимых данных в системе — бывает сложно.
Мы действительно быстро запустили MVP, перевезли данные из TestRail с помощью плагинов, и команды тестирования YADRO уже более года пользуются системой. Но есть одно «но»: пользовательские сценарии разных команд сильно отличаются. Так, добавление в систему более полумиллиона тестов привело к просадке скорости работы определенных эндпоинтов, завязанных на древовидных структурах.
Спойлер: камнем преткновения для нас стали CPU-bound задачи с большим количеством данных, о том, как я это выяснил, расскажу ниже. Изучив, как можно ускорить выполнение таких задач в Python, я протестировал несколько решений и нашел оптимальное. Если вы разрабатываете веб-приложение на Django или Python и так же, как я, хотите ускорить работу сервиса, читайте эту статью.
В качестве альтернативы Rust можно посмотреть и на Cython, и на чистый C. Возможно, стоит подумать о микросервисной архитектуре и отдавать тяжелые с точки зрения CPU задачи микросервису, написанному, например, на Go. Я же последовал модным веяниям и решил изучить, как написать Python-пакет с использованием Rust.
Отправная точка — 10,93 секунды
Изначально система отображения данных в виде древовидной структуры нас более чем устраивала. Но из-за неожиданно большого количества данных в TestY стали видны узкие места. Например, интерфейс для создания плана тестирования, который можно увидеть на скриншоте ниже.
Чтобы построить красивое дерево выбора, все тестовые сьюты для выбора и кейсы внутри них отдаются эндпоинтом сразу. А значит, если у вас 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:
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. Выкладка нашего модуля будет выглядеть так:
В 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, о которой мы недавно писали, можно по ссылке.
