Como trocar de framework nos salvou 8GB de memória

No meu time na Gazin, um grande problema que tivemos a medida que nossos projetos cresceram foi o consumo de memória RAM em testes. Com o uso de uma stack baseada em Node/React, optamos por usar o Jest como framework de testes para todos os projetos. No entanto, a medida que o projeto cresceu observamos um consumo grande de memória ao rodar testes, especialmente nos projetos usando React. E em um primeiro momento tentamos otimizar os testes atuais, depois exploramos usar o @swc/jest como transpilador, ambos esforços sem sucesso. Parte desses problemas podem estar relacionados ao bug 11956. Assim, me propus a procurar uma alternativa e, mais importante, mensurar a eficácia dela.

Procpath

Primeiro eu precisava de uma ferramenta capaz de medir com precisão a memória usada por um processo, ou uma série de processos. Como eu desenvolvo usando linux optei por uma ferramenta chamada Procpath capaz de realizar queries na árvore de processos do Linux, e retornar métricas de um mais processos, como consumo de CPU e várias medidas de memória como Resident Set Size (RSS), Proportional Set Size (PSS) e Unique Set Size (USS).

Mas qual medida usar?

Como eu disse, é possível gravar várias medidas de memória usando o Procpath, mas qual delas é a melhor?

RSS: Mede o consumo total de memória alocada na RAM incluindo o tamanho total das bibliotecas compartilhadas usadas pelo processo
PSS: Medida de consumo de memória alocada incluindo um tamanho proporcional de páginas de bibliotecas compartilhadas entre processos
USS: Consumo de memória alocada na RAM com contando apenas as páginas únicas ao processo.

Aqui escolhi usar a métrica de Unique Set Size para os testes.

Medindo o Jest

Com a Procpath, podemos usar o comando procpath query <query string> para inspecionar a árvore de processos, eu queria identificar qualquer processo que tivesse jest na sua linha de comando, então com um pouco de pesquisa cheguei na seguinte query string:

procpath query "$..children[?("jest" in @.cmdline)]"

Rodando esse comando com os testes em execução identifiquei uma série de processos, o que me indicou que eu estava no caminho certo.

Agora eu precisava gravar as informações desses processos enquanto os testes estavam em execução, para isso, usei o seguinte comando:

sudo procpath record -i 1 -d jest.sqlite "$..children[?("jest" in @.cmdline)]" -f smaps_rollup,stat,cmdline

A flag -i define o intervalo de pooling da query, aqui usei 1 segundo. Já a flag -d nomeia o arquivo sqlite onde serão gravados os dados. E a flag -f define alguns atributos extras necessários para medir USS depois.

Por fim, para visualizar os dados usei o comando:

procpath plot -q uss -d jest.sqlite -f jest.svg

Esse comando plota um gráfico do Unique Set Size (USS) ao longo do tempo, e como primeiro resultado obtive esse gráfico: Consumo de memória por processo do jest ao longo do tempo

Aqui vemos que o Jest cria onze processos (um deles é o procpath gravando os outros processos), e quase todos os processos chegam a consumir 1Gib no pico de uso.

Agregando processos

Mas como agregar esses processos para ter uma medida e gráfico mais exatos? Para isso, usei a query abaixo, uma modificação da query padrão do procpath, no banco jest.sqlite, que agrega linhas com a mesma timestamp e soma os valores de USS, e a exportei como JSON.

.mode json
.once jest-aggr.json

select
    ts,
    SUM((smaps_rollup_private_clean + smaps_rollup_private_dirty) / 1024.0) USS
from record group by ts;

Com esses dados usei o PyPlot para desenhar um gráfico do consumo agregado: Consumo agregado do jest

Fiquei assustado com esse gráfico, embora ao observar o consumo de memória do sistema em ferramentas como htop pude confirmar que ele refletia a realidade.

Agora eu tinha um processo bem definido para poder medir e comparar a eficácia de outras ferramentas ou otimizações.

Vitest

A alternativa mais adequada que eu queria testar, e que acabou sendo a única, era o Vitest, um framework de testes que usa o Vite como bundler e tem a sua API de testes baseada no Jest, oferecendo assim uma migração relativamente simples. Além disso, ele também tem suporte e preferência por módulos ESM, uma feature que ainda anda a passos lentos no Jest.

Seguindo o guia de migração, ativei as funções globais do Vitest e troquei todas as referências ao namespace jest para vi, como por exemplo:

- const bar = jest.fn()
+ const bar = vi.fn()

Também alterei todas as chamadas de jest.requireActual para vi.importActual

vi.mock('react-router-dom', async () => {
-  const actual = jest.requireActual('react-router-dom')
+  const actual: object = await vi.importActual('react-router-dom')
  
  return {
    ...actual,
    useHistory: () => vi.fn()
  }
})

A biblioteca @testing-library/jest-dom que usamos tem o problema de depender do pacote @types/jest que sobrescreve as constantes globais do Vitest. Para corrigir isso, re-exportei os tipos do Vitest em um arquivo .d.ts

const { SuiteAPI, TestAPI, SuiteHooks } = require('vitest')

declare const suite: SuiteAPI<{}>
declare const test: TestAPI<{}>
declare const describe: SuiteAPI<{}>
declare const it: TestAPI<{}>

declare const beforeAll: (
  fn: SuiteHooks['beforeAll'][0],
  timeout?: number | undefined
) => void
declare const afterAll: (
  fn: SuiteHooks['afterAll'][0],
  timeout?: number | undefined
) => void
declare const beforeEach: (
  fn: SuiteHooks['beforeEach'][0],
  timeout?: number | undefined
) => void
declare const afterEach: (
  fn: SuiteHooks['afterEach'][0],
  timeout?: number | undefined
) => void

Pronto, com essas mudanças já era possível executar todos os testes com êxito.

Comparando frameworks

Agora com o Vitest executando todos os nossos testes eu podia medir o consumo de memória usando o mesmo processo descrito acima. E aqui estão os resultados:

Consumo de memória por processo do Vitest

Como o Vitest usa threads para paralelismo ao invés de processos, nós só temos três processos no gráfico, sendo eles um test runner, o esbuild e o próprio procpath. Mas nesse caso podemos considerar o consumo de memória do esbuild desprezível.

Aqui usei novamente o PyPlot para criar um gráfico comparativo do consumo agregado de memória: Comparativo de memória entre Jest e Vitest

Como podemos ver, apesar do Vitest demorar um pouco mais para executar os testes, ele mantém um consumo de memória estável ao longo do tempo, enquanto o Jest tem um consumo crescente.

E então?

Como observado, não somente o Vitest se mostrou extremamente mais performático quando comparado ao Jest, mas o esforço necessário para fazer a migração é pequeno devido à API compatível. Além disso, o Vitest entrega suporte aos novos módulos ESM nativamente, algo que já nos impediu de usar algumas bibliotecas no passado.