Интеграционные тесты для проблемы N+1 в Java

Проблема N+1 — распространенная проблема во многих корпоративных проектах. Хуже всего то, что вы не замечаете этого, пока объем данных не станет огромным. К сожалению, код может дойти до стадии, когда решение проблемы N + 1 станет непосильной задачей.

В этой статье я говорю вам:

Как автоматически отслеживать проблему N + 1?
Как написать тест, чтобы проверить, что количество запросов не превышает ожидаемого значения?

Технический стек состоит из Java, Spring Boot, Spring Data JPA и PostgreSQL. Посмотреть репозиторий с примерами кода можно по этой ссылке.

Нет никаких ограничений на применение Spring Boot или Hibernate. Если вы взаимодействуете с javax.sql.DataSource в своей кодовой базе, то решение вам поможет. Даже если вы вообще не используете Spring.

Пример задачи N+1

Предположим, мы работаем над приложением, управляющим зоопарками. В этом случае есть две основные сущности: Zoo и Animal. Посмотрите на фрагмент кода ниже:

@Entity
@Table(name = "zoo")
public class Zoo {
    @Id
    @GeneratedValue(strategy = IDENTITY)
    private Long id;

    private String name;

    @OneToMany(mappedBy = "zoo", cascade = PERSIST)
    private List<Animal> animals = new ArrayList<>();
}

@Entity
@Table(name = "animal")
public class Animal {
    @Id
    @GeneratedValue(strategy = IDENTITY)
    private Long id;

    @ManyToOne(fetch = LAZY)
    @JoinColumn(name = "zoo_id")
    private Zoo zoo;

    private String name;
}

Теперь мы хотим получить все существующие зоопарки с их животными. Посмотрите на код ZooService ниже.

@Service
@RequiredArgsConstructor
public class ZooService {
    private final ZooRepository zooRepository;

    @Transactional(readOnly = true)
    public List<ZooResponse> findAllZoos() {
        final var zoos = zooRepository.findAll();
        return zoos.stream()
                   .map(ZooResponse::new)
                   .toList();
    }
}

Кроме того, мы хотим проверить, что все работает гладко. Итак, вот простой интеграционный тест:

@DataJpaTest
@AutoConfigureTestDatabase(replace = NONE)
@Transactional(propagation = NOT_SUPPORTED)
@Testcontainers
@Import(ZooService.class)
class ZooServiceTest {
    @Container
    static final PostgreSQLContainer<?> POSTGRES = new PostgreSQLContainer<>("postgres:13");

    @DynamicPropertySource
    static void setProperties(DynamicPropertyRegistry registry) {
        registry.add("spring.datasource.url", POSTGRES::getJdbcUrl);
        registry.add("spring.datasource.username", POSTGRES::getUsername);
        registry.add("spring.datasource.password", POSTGRES::getPassword);
    }

    @Autowired
    private ZooService zooService;
    @Autowired
    private ZooRepository zooRepository;

    @Test
    void shouldReturnAllZoos() {
        /* data initialization... */
        zooRepository.saveAll(List.of(zoo1, zoo2));

        final var allZoos = assertQueryCount(
            () -> zooService.findAllZoos(),
            ofSelects(1)
        );

        /* assertions... */
        assertThat(
            ...
        );
    }
}

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

У меня есть отдельная статья о тестировании уровня данных в приложении Spring Boot с помощью Testcontainers. Если вы не знакомы с темой, обязательно просмотрите ее.

Тест проходит успешно. Однако, если вы регистрируете операторы SQL, вы заметите кое-что, что может вас обеспокоить. Посмотрите на вывод ниже:

-- selecting all zoos
select z1_0.id,z1_0.name from zoo z1_0
-- selecting animals for the first zoo
select a1_0.zoo_id,a1_0.id,a1_0.name from animal a1_0 where a1_0.zoo_id=?
-- selecting animals for the second zoo
select a1_0.zoo_id,a1_0.id,a1_0.name from animal a1_0 where a1_0.zoo_id=?

Как видите, у нас есть отдельный select запрос для каждого присутствующего Zoo. Общее количество запросов равно количеству выбранных зоопарков + 1. Следовательно, это задача N + 1.

Это может привести к серьезным потерям производительности. Особенно на больших объемах данных.

Автоматическое отслеживание проблемы N+1

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

Есть классная библиотека под названием datasource-proxy. Он предоставляет удобный API для обертывания javax.sql.DataSource интерфейса прокси-сервером, содержащим определенную логику. Например, мы можем зарегистрировать обратные вызовы, вызываемые до и после выполнения запроса. Что интересно, библиотека также содержит готовое решение для подсчета выполненных запросов. Мы собираемся немного изменить его, чтобы удовлетворить наши потребности.

Служба подсчета запросов

Во-первых, добавьте библиотеку в зависимости:

implementation "net.ttddyy:datasource-proxy:1.8"

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

@UtilityClass
public class QueryCountService {
    static final SingleQueryCountHolder QUERY_COUNT_HOLDER = new SingleQueryCountHolder();

    public static void clear() {
        final var map = QUERY_COUNT_HOLDER.getQueryCountMap();
        map.putIfAbsent(keyName(map), new QueryCount());
    }

    public static QueryCount get() {
        final var map = QUERY_COUNT_HOLDER.getQueryCountMap();
        return ofNullable(map.get(keyName(map))).orElseThrow();
    }

    private static String keyName(Map<String, QueryCount> map) {
        if (map.size() == 1) {
            return map.entrySet()
                       .stream()
                       .findFirst()
                       .orElseThrow()
                       .getKey();
        }
        throw new IllegalArgumentException("Query counts map should consists of one key: " + map);
    }
}

В этом случае мы предполагаем, что в нашем приложении есть один DataSource. Вот почему функция keyName в противном случае выдает исключение. Однако код не будет сильно отличаться при использовании нескольких источников данных.

SingleQueryCountHolder хранит все объекты QueryCount в обычном файле ConcurrentHashMap.

Наоборот, ThreadQueryCountHolder сохраняет значения в объекте ThreadLocal. Но для нашего случая достаточно SingleQueryCountHolder.

API предоставляет два метода. Метод get возвращает текущее количество выполненных запросов, а метод clear устанавливает счетчик на ноль.

BeanPostProccessor и прокси-сервер DataSource

Теперь нам нужно зарегистрировать QueryCountService, чтобы он собирал данные из DataSource. В этом случае на помощь приходит интерфейс BeanPostProcessor. Посмотрите на пример кода ниже.

@TestComponent
public class DatasourceProxyBeanPostProcessor implements BeanPostProcessor {
    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) {
        if (bean instanceof DataSource dataSource) {
            return ProxyDataSourceBuilder.create(dataSource)
                       .countQuery(QUERY_COUNT_HOLDER)
                       .build();
        }
        return bean;
    }
}

Я помечаю класс аннотацией @TestComponent и помещаю его в каталог src/test, потому что мне не нужно подсчитывать запросы вне области проверки.

Как видите, идея тривиальна. Если bean-компонент DataSource, оберните его ProxyDataSourceBuilder и поместите значение QUERY_COUNT_HOLDER в качестве QueryCountStrategy.

Наконец, мы хотим подтвердить количество выполненных запросов для конкретного метода. Посмотрите на фрагмент кода с пользовательскими утверждениями ниже:

Пользовательские утверждения

@UtilityClass
public class QueryCountAssertions {
    @SneakyThrows
    public static <T> T assertQueryCount(Supplier<T> supplier, Expectation expectation) {
        QueryCountService.clear();
        final var result = supplier.get();
        final var queryCount = QueryCountService.get();
        assertAll(
            () -> {
                if (expectation.selects >= 0) {
                    assertEquals(expectation.selects, queryCount.getSelect(), "Unexpected selects count");
                }
            },
            () -> {
                if (expectation.inserts >= 0) {
                    assertEquals(expectation.inserts, queryCount.getInsert(), "Unexpected inserts count");
                }
            },
            () -> {
                if (expectation.deletes >= 0) {
                    assertEquals(expectation.deletes, queryCount.getDelete(), "Unexpected deletes count");
                }
            },
            () -> {
                if (expectation.updates >= 0) {
                    assertEquals(expectation.updates, queryCount.getUpdate(), "Unexpected updates count");
                }
            }
        );
        return result;
    }
}

Алгоритм прост:

Установите счетчик текущих запросов на ноль.
Выполнил предоставленную лямбду.
Подтвердите количество запросов к данному объекту Expectation.
Если все пройдет успешно, вернуть результат выполнения.

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

Класс Expectation — это обычная структура данных. Посмотрите на его объявление ниже:

@With
@AllArgsConstructor
@NoArgsConstructor
public static class Expectation {
    private int selects = -1;
    private int inserts = -1;
    private int deletes = -1;
    private int updates = -1;

    public static Expectation ofSelects(int selects) {
        return new Expectation().withSelects(selects);
    }

    public static Expectation ofInserts(int inserts) {
        return new Expectation().withInserts(inserts);
    }

    public static Expectation ofDeletes(int deletes) {
        return new Expectation().withDeletes(deletes);
    }

    public static Expectation ofUpdates(int updates) {
        return new Expectation().withUpdates(updates);
    }
}

Последний пример

Посмотрим, как это работает. Сначала я добавляю утверждения запроса в предыдущем случае с проблемой N + 1. Посмотрите на блок кода ниже:

final var allZoos = assertQueryCount(
    () -> zooService.findAllZoos(),
    ofSelects(1)
);

Не забудьте импортировать DatasourceProxyBeanPostProcessor как компонент Spring в свои тесты.

Если мы повторно запустим тест, мы получим вывод ниже.

Multiple Failures (1 failure)
    org.opentest4j.AssertionFailedError: Unexpected selects count ==> expected: <1> but was: <3>
Expected :1
Actual   :3

Значит, утверждение работает. Нам удалось отследить проблему N+1 автоматически. Пришло время заменить обычное выделение на JOIN FETCH. Посмотрите на фрагмент кода ниже.

public interface ZooRepository extends JpaRepository<Zoo, Long> {
    @Query("FROM Zoo z LEFT JOIN FETCH z.animals")
    List<Zoo> findAllWithAnimalsJoined();
}

@Service
@RequiredArgsConstructor
public class ZooService {
    private final ZooRepository zooRepository;

    @Transactional(readOnly = true)
    public List<ZooResponse> findAllZoos() {
        final var zoos = zooRepository.findAllWithAnimalsJoined();
        return zoos.stream()
                   .map(ZooResponse::new)
                   .toList();
    }
}

Давайте снова запустим тест и посмотрим на результат:

Это означает, что утверждение правильно отслеживает N + 1 проблему. Кроме того, он проходит успешно, если количество запросов равно ожидаемому. Большой!

Заключение

На самом деле можно предотвратить N + 1 проблем с помощью регулярных тестов. Я думаю, что это отличная возможность установить защиту для тех частей кода, которые имеют решающее значение для производительности.

Это все, что я хотел рассказать вам о решении задачи N + 1 автоматическим способом. Если у вас есть какие-либо вопросы или предложения, оставьте свои комментарии ниже. Кроме того, если вам понравилась эта статья, поделитесь ею с друзьями и коллегами. Возможно, они тоже найдут это полезным. Спасибо за прочтение!