Всем привет! Меня зовут Роман Кабаев, я инженер в отделе разработки инструментов тестирования компании 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 секунды. Попробуем ускорить решение, чтобы добиться более высокой скорости работы. 

Упрощаем 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, и обработать данные там.

Что мы делали раньше:

1. Сформировали запрос в базу данных с помощью django-orm.

2. Отдали набор данных QuerySet в DRF сериализатор.

3. Достали из него Python-словарь.

4. Конвертировали Python словарь в JSON (данный шаг скрыт под капотом у класса Response из библиотеки DRF).

5. Отдали ответ пользователю.

Как это выглядит теперь:

1. Получаем QuerySet всех тестовых кейсов, используем метод values (). Он вернет QuerySet, содержащий словари. Обернем все это в список — так мы получим список словарей, с которым проще всего будет работать в Rust. Таким же образом получаем все тестовые кейсы, в values указываем все нужные нам поля, чтобы не тянуть лишние данные из базы.

2. Ищем тестовые сьюты, к которым относятся эти тестовые кейсы, и получаем их предков (родители родителя и так до 0 уровня). Получаем плоский набор данных по принципу, описанному выше.

3. Передаем наборы данных в функцию из нашей Rust-библиотеки и производим необходимую обработку.

4. Возвращаем список словарей. Словари, в свою очередь, представляют иерархическое дерево.

5. Сериализуем полученный список в JSON с помощью Orjson.

6. Отдаем ответ пользователю.

Для написания пакета нам нужен сам 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, о которой мы недавно писали, можно по ссылке.