Снижение аллокации при замыкании (closure)

Недавно у меня появилась задача по снижению аллокации в очень горячем месте кода. Там происходит тривиальное: запускаются Task'и в которых заранее известным набором handler'ов обрабатываются объекты. Вооружившись профайлером, я с удивлением обнаружил, что много памяти (и много времени GC) затрачивается на удаление объектов-замыканий.

Что такое замыкание в C#?

Замыкания (closure) это очень крутая штука, которая помогает писать более лаконичный код на C#. Под капотом, замыкание это более-менее обычный класс, который "захватывает" ссылки на переменные, которые участвуют в замыкании.

Думаю вы видели, как многие IDE честно подсказывают, что в месте использования замыкания возникает захват переменных:

Что же происходит "под капотом"? В этом же классе создаётся класс, который представляет из себя то самое замыкание. Класс специально называется хитрым образом (в моём случае он называется DisplayClass4_0) и помечается атрибутом CompilerGenerated.

В специальном классе создаётся набор полей по количеству переменных, захватываемых замыканием. Также, создаётся метод, ссылка на который передаётся в метод Task.Run.

В декомпилированном коде (я использую dotPeek) это выглядит примерно вот так:

Почему происходит именно так? Потому что "замыкание" это не механизм платформы .NET, а языка C#. Если угодно, это синтаксический сахар, который делает язык красивым и выразительным. Однако, на более низком уровне, любой синтаксический сахар требует низкоуровневой реализации - и это она и есть. Подробнее о замыканиях написал Сергей Тепляков и никому не известный Стефан Тауб. У их написано много, объясняется сравнительно легко, в том числе затрагиваются особенности работы с замыканиями.

Аллокация при замыкании

Можно заметить, что при каждом замыкании "под капотом" создаётся инстанс класса замыкания, в его поля помещаются захватываемые значения, а в нужный нам метод передаётся ссылка на метод closure-класса, где и происходит выполнение логики, указанной в замыкании. Напомню, что инстанс класса размещается в куче, откуда его потом удалит GC.

Кажется, что это совсем не страшно, так как речь в подавляющем большинстве случаев идёт о помещении инстанса в Gen0, откуда он будет быстро удалён. Более того, сам класс замыкания предельно лёгкий и не занимает много места.

Однако, если место использования замыкания горячее (часто вызывается), то GC может не успеть удалить инстансы closure-класса. При самых печальных сценариях, это может привести к "выживанию" классов вплоть до Gen2, с последующим stop the world для проведения вдумчивой очистки кучи.

Более того, не надо забывать, что не все имплементации платформы работают одинаково. Например, игровой движок Unity имеет особенный GC с одним поколением. Это требует от разработчиков очень внимательно относиться к тому, кто и что аллоцирует и в каких количествах.

Имплементация собственного замыкания

Чтобы снизить нагрузку на GC, в некоторых сценариях можно попытаться написать собственную имплементацию замыкания. Кажется, что это просто, так как мы знаем как работает closure.

private sealed class Closure<T> {
  private readonly Action<T> _action;
  private readonly Action _closure;
  private T _value;

  public Closure(Action<T> action) {
    _action = action;
    _closure = Execute;
  }

  [MethodImpl(MethodImplOptions.AggressiveInlining)]
  public void Clear() => _value = default;

  [MethodImpl(MethodImplOptions.AggressiveInlining)]
  public Action Prepare(T value) {
    _value = value;
    return _closure;
  }

  private void Execute() => _action(_value);
}

Если мы знаем количество вызовов этого класса, то мы можем запросто преаллоцировать все замыкания. При использовании мы просто записываем значение аргумента в поле класса (точно также, как сделали бы за нас "под капотом "), а в качестве метода используем заранее созданную ссылку на метод этого же класса.

for (var i = 0; i < _objects.Length; i++) {
    _tasks[i] = Task.Run(_closures[i].Prepare(in _objects[i]))
}

Task.WaitAll(_tasks);

foreach (var closure in _closures) {
    closure.Clear();
}

При использовании собственного класса-замыкания при декомпиляции кода мы можем наблюдать более понятную картину без "магии под капотом". Бонусом, мы избавились от создания new Action при передачи ссылки на метод замыкания в нужный нам метод.

Минус подобного использования - инстанс класса замыкания нужно очищать от значения, которое в него было передано ранее. Сделать это необходимо, поскольку это место становится местом потенциальной утечки памяти, так как замыкание будет хранить ссылку на "захваченное" значение вечно.

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

Уменьшение аллокации при замыкании

Какие же конкретные числа мы можем получить при замене стандартного механизма замыкания на собственный велосипед? Были произведены замеры с использованием известного фреймворка для микробенчмаркинга BenchmarkDotNet. Код бенчмарка находится тут.

Приятно, что скорость осталась примерно прежней. Это говорит о том, что сделано более менее правильно.

Столбец "Allocated" бодро рапортует нам о том, что аллокация меньше почти в два раза. Но, собственно, почему же она есть? Если вы посмотрите код бенчмарка, то вы заметите, что я пытаюсь минимизировать аллокацию при запуске Task'ов. Это достаточно распространенный случай использования замыкания. Цифры, которые можно увидеть в столбце Allocated включают в себя затраты платформы на создание Task'ов.

Parallel.For и Parallel.ForEach

В бенчмарке, также, можно найти результаты для Parallel.For и Parallel.ForEach. Их использование значительно повышает скорость работы и, к сожалению, существенно увеличивают аллокацию. Дьявол кроется в деталях: Parallel.ForEach принимает в качестве аргумента IEnumerable<T>, который возвращает IEnumerator<T>. Это объект, который будет расположен в куче, а значит будет нагружать GC. Ну а Parallel.For принимает делегат, где снова создаётся объект-замыкания, что также влияет на аллокацию. Спасибо комментатору @Deosis.

И снова дьявол кроется в деталях. При увеличении количества заданий (изначально их было 10) начинает стремительно выигрывать имплементация на Parallel.For и Parallel.ForEach. Во-первых, она просто быстрая, а во-вторых, создание enumerator'a и аллокация замыкания - это фиксированная плата, никак не зависящая от количества. Бенчмарк это явно показывает. Спасибо @soalexmn. В его комментарии количество заданий увеличено до 400 и цифры там совсем драматичные.

P.S.: Начал писать в телегу и, внезапно, Дзен. Заглядывайте, если интересно.