Арены своими руками

Post Thumbnail

Перевод статьи "Cheating the Reaper in Go". Автор рассказывает про кастомную реализацию арен (Golang memory arena) с ручным управлением памятью. Не думаю, что как Golang разработчик вы будете применять это в рабочих задачах, но это очень интересный эксперимент, который раскрывает некоторые нюансы работы GC. Читаемся и наслаждаемся этим не простым, но интересным Golang How To

Хотя в душе я программист на C++, язык Golang восхищает меня. Правда, совсем не по тем причинам, о которых вы могли подумать. В Go было принято несколько интересных дизайнерских решений:

  • В языке практически отсутствует неопределённое поведение (Undefined Behavior)
  • У него очень простая семантика сборки мусора (GC), которую сложно изменить из-за особенностей языка.

Благодаря этому, несмотря на наличие GC, в Golang можно вручную управлять памятью даже без использования runtime, на чистом Go, но в согласии со сборщиком мусора. Чтобы продемонстрировать это, мы создадим нетипизированную arena-абстракцию с поддержкой GC и все это опираясь на внутренние детали реализации сборщика.

В Rust или C++ я бы не стал так рисковать: LLVM крайне умён и с каждой новой версией компилятора находит новые способы "сломать" подобные трюки. Go, к счастью, работает немного иначе. Хотя вам никто не гарантирует обратной совместимости для кода, использующего unsafe, на практике этому мешают два фактора:

  • Go не пытается строго определить, что разрешено, а что нет: у unsafe просто нет формальной семантики.
  • Go ставит во главу угла сохранение работоспособности экосистемы, а значит, можно предположить, что закон Хайрума защитит некоторые наблюдаемые поведения рантайма. Из этого мы сможем сделать выводы о том, какие изменения могут легко сломать код, а какие — нет.

Это контрастирует с высокопроизводительными компиляторами вроде LLVM, где чётко определены границы неопределённого поведения, что позволяет безнаказанно "ломать" программы, выходящие за эти рамки.

Итак, давайте "обманем смерть" и погрузимся в детали.

Что мы будем разрабатывать?

Наша цель — создать арену (arena). Ареной называется структура данных для эффективного выделения памяти, в которой все объекты имеют одинаковое время жизни. Это снижает нагрузку на общий аллокатор: вместо множества мелких запросов память запрашивается крупными блоками и освобождается сразу вся.

Для рассмотрим следующую программу на Golang:

package main

import "fmt"

func main() {
  var s []int
  for i := range 1000 {
    prev := cap(s)
    s = append(s, i)
    if cap(s) != prev {
      fmt.Println(cap(s))
    }
  }
}

Эта программа выводит последовательные степени двойки. Это происходит потому, что append реализован примерно так:

func append[S ~[]T, T any](a, b S) S {
  // Если нужно, увеличиваем выделенную память.
  if cap(a) - len(a) < len(b) {
    // Либо удваиваем размер, либо выделяем ровно больше памяти, 
    // если удвоение недостаточно.
    newCap := max(2*cap(a), len(a)+len(b))

    // Расширяем `a`.
    a2 := make([]T, len(a), newCap)
    copy(a2, a)
    a = a2
  }

  // Увеличиваем длину `a`, чтобы вместить `b`, 
  // и записываем `b` в новую область.
  a = a[:len(a)+len(b)]
  copy(a[len(a)-len(b):], b)
  return a
}

При добавлении небольших элементов make вызывается всего O(log n) раз — это намного эффективнее, чем выделять память при каждом вызове append. Практически во всех языках динамические массивы используют эту оптимизацию.

Арена обобщает этот подход, но вместо экспоненциального увеличения размера она выделяет новые блоки памяти и раздаёт указатели на них. Желаемый интерфейс выглядит так:

type Allocator interface {
  Alloc(size, align uintptr) unsafe.Pointer
}

В Go нет "неинициализированной" памяти, видимой пользователю, поэтому возвращаемая область должна быть обнулена. Также требуется, чтобы align был степенью двойки.

Мы можем создать типобезопасный интерфейс, добавив обобщённую функцию New:

// New выделяет память под нулевое значение golang типа T через указанный аллокатор
// и возвращает указатель на него.
func New[T any](a Allocator) *T {
  var t T
  p := a.Alloc(unsafe.Sizeof(t), unsafe.Alignof(t))
  return (*T)(p)
}

Это выглядит вполне разумно для тех, кто привык работать с malloc или operator new в C++, но здесь есть важная проблема. Что произойдёт, если мы выделим память для указателя *int через этот аллокатор?

// Выделяем память для указателя через наш аллокатор,
// затем инициализируем его указателем в куче Go
p := New[*int](myAlloc)
*p = new(int)

runtime.GC()
**p = 42  // Использование после освобождения!

Хотя Allocator.Alloc принимает только размер и выравнивание (чего достаточно для описания любого golang типа в памяти), сборщику мусора Go требуется дополнительная информация. Например, на 64-битных системах int и *int имеют одинаковый layout(как golang типы располагаются в памяти): 8 байт размера, 8 байт выравнивания. Те в памяти они расположено одинаково. А значит, что в программе выше GC не поймет что у нас тут указатель и освободит память слишком рано.

Но для корректной работы GC нужно знать разницу между лайяутом значения(как расположен в памяти) и типом значения(информация о содержимом). Чтобы разобраться с этим, нужно понять как работает сборщик мусора.

Алгоритм "Mark and Sweep"

Для полного понимания принципов работы сборщика мусора рекомендую ознакомиться с учебным примером GC, который я разработал ранее: The Alkyne GC.

Основная задача сборщика мусора это управление аллокатором памяти и учет двух ключевых аспектов:

  • Какая память была выделена
  • Какая память всё ещё используется

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

Наиболее распространённая реализация использует архитектуру "mark and sweep". Работает это следующим образом:

Наиболее популярный способ достижения этого - через архитектуру "маркировки и очистки". Сборщик мусора периодически обходит весь граф объектов программы, начиная с определённых корневых точек (корни стека, глобальные переменные и т.д.). Каждый найденный живой объект помечается как активный. После завершения маркировки, вся память, не помеченная как активная, считается мусором. Эта память помечается как свободная для будущего использования (если планируется повторное использование) или, в случае значительного избытка памяти, возвращается операционной системе.

Корни графа обычно являются сущностями с которыми программа активно работает. В случае Go это любые данные, находящиеся в данный момент на стеке какой-либо горутины (G), глобальные переменные и все что находиться в глобальной области памяти (их набор известен на этапе компиляции).

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

Фаза маркировки начинается со сканирования стеков: сборщик проверяет стек каждой горутины (G) и находит все указатели, содержащиеся в нём. Компилятор Go генерирует метаданные для каждой функции, с пометкой какие ячейки стека в её фрейме содержат указатели. Все такие указатели по определению считаются "живыми".

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

Критический шаг здесь — взять адрес некоторого выделенного участка памяти и преобразовать его во все значения указателей внутри него. В Go используется точная сборка мусора (precise garbage collection), что означает, что сборщик учитывает только те значения, которые явно объявлены как указатели в языке: целое число, которое случайно выглядит как адрес, не приведёт к его обработке как указателя. Это обеспечивает более эффективное использование памяти, но добавляет некоторую сложность в работу сборщика мусора.

Например, типы *int, map[int]byte, string, struct {A int; B *int} содержат по крайней мере один указатель, в то время как int, [1000]byte, struct {X bool; F uintptr} не содержат. Последние называются типами без указателей.

Go улучшает макет типа до формы, добавляя битовую карту, которая указывает, какие слова выровнены по указателю и имеют размер указателя в области памяти типа, содержат указатель. Они называются битами указателя. Например, вот формы нескольких типов Go на 64-битной системе.

Go расширяет layout типа до shape типа(я хз как тут перевести layout и shape). Добавляет битовую маску, которая указывает, какие выровненные под указатель слова (размером с указатель) в памяти содержат указатели. Это называется битами указателей. Например, вот shape для нескольких типов Go в 64-битной системе.

Type Size/Align Pointer Bits
byte 1/1 0
int 8/8 0
rune 4/4 0
*int 8/8 1
unsafe.Pointer 8/8 1
string 16/8 10
[]int 24/8 100
[3]string 48/8 101010
map[int]byte 8/8 1
map[int]string 8/8 1
any 16/8 014
error 16/8 01
func(int) int 8/8 1
runtime.hchan5 104/8 0010110011110

 

В сборщике мусора Go каждое выделение помечается своей формой (это делается различными способами в сборщике мусора, либо через явный заголовок на самом выделении («заголовок `malloc`»), либо через тип времени выполнения, хранящийся в `runtime.mspan` выделения, или другой механизм). При сканировании значения он использует эту информацию для определения местоположения указателей, которые нужно просканировать.

В сборщике мусора (GC) Go каждое выделение памяти помечается своим shape. Это реализуется разными способами: через явный заголовок выделения (так называемый "malloc header"), через тип в runtime.mspan или другие механизмы. При сканировании значения GC использует эту информацию, чтобы определить, где находятся указатели, которые нужно проверить.

Самая очевидная проблема с типом Allocator.Alloc заключается в том, что он не различает shape (формы памяти), а значит, не может корректно выделять память, содержащую указатели. Сборщик мусора (GC) не сможет обнаружить эти указатели и освободит их преждевременно

В нашем примере, где мы выделили память под *int с помощью нашего кастомного аллокатора, в итоге получаем **int на стеке. Можно подумать, что Go просто пройдёт через первый указатель *, чтобы найти *int и пометить его как живой, но это не так. Вместо этого Go обнаруживает указатель на какой-то блок памяти, который пользовательский аллокатор получил из кучи, но при этом у этого блока отсутствует информация о типе (pointer bits) в его структуре (shape)

Почему Go не проверяет тип указателя, через который он проходит? Две причины.

  • Все указатели в Go с точки зрения рантайма — нетипизированные: каждый *T преобразуется в unsafe.Pointer. Это позволяет большей части рантайма Go быть «универсальной» без использования настоящих дженериков.
  • Метаданные о типе значения (pointee) можно агрегировать, так что каждый указатель на объект не обязан хранить информацию о своём типе во время выполнения.

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

Получается, мы не можем безопасно размещать указатели в нашей арене. Из-за этого наш API и конструктор New становятся небезопасными - особенно потому, что в Go нет стандартного способа ограничить дженерики, запретив им содержать указатели (очевидно, большинству пользователей такие детали не важны).

Технически, можно вычислить биты указателей (pointer bits) для типа через рефлексию, но это очень медленно, а главная цель использования арен - скорость. Однако по мере проектирования нашей арены станет ясно, что безопасное хранение указателей в ней всё же возможно.

Дизайн Арены

Теперь, когда мы хорошо разобрались в том, как работает сборщик мусора Go, можно приступить к проектированию быстрой структуры арены.

Идеальный случай — когда вызов Alloc выполняется очень быстро: в обычном сценарии это просто смещение указателя. Первое допущение, которое мы можем сразу сделать — вся память будет выровнена по максимальному граничному значению: большинство объектов имеют размер как минимум с указатель, а Go и так задаёт максимальное выравнивание для пользовательских типов. Поэтому мы можем просто игнорировать параметр align и всегда выравнивать, скажем, по 8 байтам. Это означает, что указатель на следующий нераспределённый блок всегда будет правильно выровнен.

Таким образом, можно реализовать такую структуру:

type Arena struct {
    next unsafe.Pointer // указатель на следующий свободный блок
    left uintptr        // оставшееся место (в машинных словах)
    cap  uintptr        // текущая ёмкость арены (в машинных словах)
}

const (
 // Размер минимального блока выделения (степень двойки).
  wordBytes = 8  // Зависит от архитектуры, здесь для 64-бит.
  minWords  = 8 //  Минимальный размер выделения (в словах).
)

func (a *Arena) Alloc(size, align uintptr) unsafe.Pointer {
  // Сначала округляем размер до выравнивания объектов в арене.
  mask := wordBytes - 1
  size = (size + mask) &^ mask
  // Затем переводим размер в машинные слова (uintptr).
  // Это не приводит к потере данных, так как size теперь кратен wordBytes.
  words := size / wordBytes

  // Проверяем, хватит ли места. Если нет — расширяем арену.
  if a.left < words {
    // Выбираем наибольшее из: минимальный размер, удвоенная текущая ёмкость
    // или ближайшая степень двойки, превышающая words.
    a.cap = max(minWords, a.cap*2, nextPow2(words))
    a.next = unsafe.Pointer(unsafe.SliceData(make([]uintptr, a.cap)))
    a.left = a.cap
  }

  // Выделяем блок, смещая указатель.
  p := a.next
  a.left -= words
  if a.left > 0 {
    a.next = unsafe.Add(a.next, size)
  } else {
    // Смещение за пределы выделенной памяти в Go явно запрещено.
    a.next = nil
  }

  return p
}

// nextPow2 возвращает ближайшую степень двойки, большую или равную n.
func nextPow2(n uintptr) uintptr {
  return uintptr(1) << bits.Len(uint(n))
}

Как быстро этот код работает? Вот небольшой бенчмарк

func BenchmarkArena(b *testing.B) {
  bench[int](b)
  bench[[2]int](b)
  bench[[64]int](b)
  bench[[1024]int](b)
}

const runs = 100000

var sink any

func bench[T any](b *testing.B) {
  var z T
  n := int64(runs * unsafe.Sizeof(z))
  name := fmt.Sprintf("%v", reflect.TypeFor[T]())

  b.Run(name, func(b *testing.B) {
    b.Run("arena", func(b *testing.B) {
      b.SetBytes(n)
      for b.Loop() {
        a := new(arena.Arena)
        for range runs {
          sink = arena.New[T](a)
        }
      }
    })

    b.Run("new", func(b *testing.B) {
      b.SetBytes(n)
      for b.Loop() {
        for range runs {
          sink = new(T)
        }
      }
    })
  })
}

Цель этого бенчмарка — измерить стоимость выделения множества объектов одинакового размера. Количество выполнений цикла for b.Loop() заранее неизвестно и определяется фреймворком для тестирования, чтобы уменьшить статистические аномалии. Если бы мы тестировали всего одно выделение памяти, результат будет сильно зависеть от количества прогонов, а в нашем бенче все будет максимально честно.

Мы также используем b.SetBytes, чтобы оценить пропускную способность в бенчмарке. Это немного удобнее для интерпретации, чем простое измерение в ns/op (наносекунды на операцию), которое выдает тест по умолчанию. Метрика bytes/sec показывает, сколько памяти каждый аллокатор может выделить за единицу времени.

Нам нужно сравнить производительность с обычным new, но если написать просто _ = new(T), компилятор оптимизирует этот вызов, так как указатель никуда не "убегает". Запись в глобальную переменную гарантирует, что указатель считается "убегающим", и выделение памяти не будет оптимизировано - то что нужно.

Ниже приведены результаты (в сокращённом виде, только байт в секунду). Все тесты проводились на AMD Ryzen Threadripper 3960X. Чем больше значение — тем лучше.

BenchmarkArena/int/arena-48         794.84 MB/s
BenchmarkArena/int/new-48           390.59 MB/s
BenchmarkArena/[2]int/arena-48      1263.58 MB/s
BenchmarkArena/[2]int/new-48        528.06 MB/s
BenchmarkArena/[64]int/arena-48     7370.08 MB/s
BenchmarkArena/[64]int/new-48       2865.24 MB/s
BenchmarkArena/[1024]int/arena-48   9889.20 MB/s
BenchmarkArena/[1024]int/new-48     2875.75 MB/s

Круто, похоже мы двигаемся в нужную сторону. Арена быстрее аллоцирует больше памяти. Увеличение производительности, похоже, растет с количеством выделенной памяти - от 2 до 4 раз.

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

Не теряем память зря

В методе (*Arena).Alloc, когда мы выделяем новый блок памяти и перезаписываем a.next, сборщик мусора (GC) теоретически может освободить старый блок. Однако это не проблема: пока существуют "живые" указатели на этот блок арены, GC не освободит его, независимо от состояния самой арены. Так что, кажется, нам не о чем беспокоиться? Не совсем.

Вся суть арены — выделять большое количество памяти с одинаковым временем жизни. Это особенно полезно для графовых структур данных, например, AST (абстрактного синтаксического дерева) или промежуточного представления (IR) в компиляторах, где происходит множество операций с выделением памяти, а затем результат просто выбрасывается.

Мы не можем размещать указатели в арене, потому что они исчезнут из поля зрения сборщика мусора (GC) и будут освобождены слишком рано. Если указатель должен находиться в арене, он по определению должен пережить саму арену, поскольку он существует дольше, чем часть арены, а арена по задумке имеет единое время жизни для всех своих объектов.

Ключевая идея: если сделать так, чтобы любой указатель, возвращённый Alloc, предотвращал освобождение всей арены сборщиком мусора (GC), то арена сможет безопасно содержать указатели в себе. Рассмотрим пример:

  • У нас есть указатель p **int, выделенный в арене a.
  • GC видит этот указатель (как unsafe.Pointer без информации о типе) и помечает его блок памяти как "живой".
  • Каким-то образом GC также должен помечать саму арену a как "живую" вслед за указателем.
  • Далее GC помечает все блоки памяти, выделенные ареной a, как "живые".
  • В результате блок, на который указывает *p, тоже остаётся "живым", поэтому сам *p не нужно помечать отдельно — и он не будет освобождён prematurely.

Шаг 3 критически важен. Если заставить GC помечать всю арену как активную, то любые указатели внутри неё, ссылающиеся на другие объекты арены, будут автоматически сохраняться — без необходимости для GC знать, как их сканировать.

Таким образом, выражение *New[*int](a) = new(int) по-прежнему приведёт к use-after-free (использованию после освобождения). Но *New[*int](a) = New[int](a) уже не вызовет проблем.

Это небольшое улучшение не делает нашу арену абсолютно безопасной, но структура данных с внутренней ареной может быть полностью безопасной — при условии, что в арену помещаются только указатели на объекты, которые алоцированы в этой же арене.

Как это реализовать? Ответ на шаге 4: добавляем в арену поле []unsafe.Pointer и записываем туда каждый выделенный указатель. Это гарантирует, что GC будет видеть все указатели арены при сканировании.

type Arena struct {
  next      unsafe.Pointer
  left, cap uintptr

  chunks []unsafe.Pointer  // Новое поле: список всех выделенных блоков
}

func (a *Arena) Alloc(size, align uintptr) unsafe.Pointer {
  // ... пропущенный код ...
    
  if a.left < words {
      // Выбираем наибольшее из: минимальный размер блока,
      // удвоенная текущая емкость или следующая степень двойки
      a.cap = max(minWords, a.cap*2, nextPow2(words))
      a.next = unsafe.Pointer(unsafe.SliceData(make([]uintptr, a.cap)))
      a.left = a.cap
      a.chunks = append(a.chunks, a.next)  // Добавляем новый блок в список
  }
    
    // ... пропущенный код ...
}

Стоимость операции append амортизирована: для выделения n байт мы в итоге выполняем дополнительно O(log log n) выделений. Но как это повлияет на наши бенчмарки?

BenchmarkArena/int/arena-48         800.08 MB/s
BenchmarkArena/int/new-48           386.81 MB/s
BenchmarkArena/[2]int/arena-48      1236.00 MB/s
BenchmarkArena/[2]int/new-48        520.84 MB/s
BenchmarkArena/[64]int/arena-48     7999.71 MB/s
BenchmarkArena/[64]int/new-48       2706.68 MB/s
BenchmarkArena/[1024]int/arena-48   9998.00 MB/s
BenchmarkArena/[1024]int/new-48     2816.28 MB/s

Результаты практически не изменились, и это хороший знак.

Обратные указатели

Теперь, когда арена надёжно хранит всю выделенную память, нужно сконцентрироваться на шаге 3: сделать так, чтобы существование любого действующего указателя из Alloc автоматически защищало всю арену от сборки мусора.

Здесь мы можем использовать важное свойство работы GC в Go: любой указатель на выделенную память предотвращает её освобождение, как и всё, что достижимо через этот указатель. Однако наши блоки памяти (chunks) — это срезы []uintptr, которые GC не сканирует. Но если бы в таком срезе хотя бы один указатель подлежит сканированию, мы могли бы разместить там указатель *Arena на саму арену. Тогда, сканирование GC любого объекта, возвращённого Alloc, приводило бы к пометке всей арены a как "живой".

До сих пор мы выделяли память под [N]uintptr с помощью make([]T), но на самом деле нам нужно выделить память под структуру вида:

struct { 
    A [N]uintptr
    P unsafe.Pointer 
}

где N — это некоторое динамическое значение.

В своей бесконечной мудрости, стандартная библиотека Go предоставляет специальный механизм для этого: reflect.StructOf. Его можно использовать для создания произвольных анонимных структур во время выполнения, которые затем можно разместить в куче.

Вместо вызова make мы можем использовать такую функцию:

func (a *Arena) allocChunk(words uintptr) unsafe.Pointer {
    chunk := reflect.New(reflect.StructOf([]reflect.StructField{
        {
            Name: "X0",
            Type: reflect.ArrayOf(int(words), reflect.TypeFor[uintptr]()),
        },
        {Name: "X1", Type: reflect.TypeFor[unsafe.Pointer]()},
    })).UnsafePointer()

    // Смещаемся в конец чанка и записываем туда `a`.
    end := unsafe.Add(chunk, words * unsafe.Sizeof(uintptr(0)))
    *(**Arena)(end) = a

    return chunk
}

Пояснение:

  • reflect.StructOf позволяет создать новый тип структуры динамически.
  • reflect.ArrayOf создаёт тип массива заданного размера words с элементами типа uintptr.
  • reflect.New выделяет память под новую структуру и возвращает указатель (в виде reflect.Value).
  • UnsafePointer() преобразует reflect.Value в unsafe.Pointer.
  • unsafe.Add вычисляет адрес конца чанка, куда записывается указатель на арену a.

Это оказывает небольшое, но заметное влияние на производительность

BenchmarkArena/int/arena-48         763.91 MB/s
BenchmarkArena/int/new-48           385.49 MB/s
BenchmarkArena/[2]int/arena-48      1174.00 MB/s
BenchmarkArena/[2]int/new-48        524.32 MB/s
BenchmarkArena/[64]int/arena-48     7563.54 MB/s
BenchmarkArena/[64]int/new-48       2649.63 MB/s
BenchmarkArena/[1024]int/arena-48   8668.02 MB/s
BenchmarkArena/[1024]int/new-48     2648.10 MB/s

Дополнительные оптимизации

Если вернуться к Arena.Alloc, в конце этой функции есть if:

func (a *Arena) Alloc(size, align uintptr) unsafe.Pointer {
    // ... другой код ...

    // Выделяем память, увеличивая указатель.
    p := a.next
    a.left -= words
    if a.left > 0 {
        a.next = unsafe.Add(a.next, size)
    } else {
        // Важно: смещение на 'один за концом' — одна из немногих вещей,
        // которые Go явно запрещает.
        a.next = nil
    }

    return p
}

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

В C++, если у нас есть массив int из n элементов, и int* p - указатель на начало массива, то p + n является корректным указателем, даже если его нельзя разыменовать: он указывает на позицию сразу после последнего элемента массива. Это полезная конструкция, так как позволяет, например, избавиться от переменной индукции в цикле:

// Наивный цикл с переменной индукции:
for (int i = 0; i < n; i++) {
  do_something(p[i]);
}

// Оптимизированный вариант без переменной индукции:
for (auto end = p + n; p < end; p++) {
  do_something(*p);
}

Однако в Go такой подход вызывает проблемы, поскольку он сбивает с толку сборщик мусора. GC не может отличить указатель 'one-past-the-end' для аллокации A от указателя на начало следующей аллокации B, расположенной сразу после неё. В лучшем случае это приводит к тому, что память остаётся занятой дольше необходимого, а в худшем - срабатывают защитные механизмы GC, вызывающие панику, если GC попытается просканировать указатель на адрес, который, как ему известно, уже был освобождён.

Но в нашем коде каждый чанк теперь содержит дополнительный элемент в самом конце, который не используется для аллокации. Это позволяет нам иметь указатель 'one-past-the-end' для массива [N]uintptr, из которого мы выделяем память.

Обновлённая функция аллокации будет выглядеть так:

func (a *Arena) Alloc(size, align uintptr) unsafe.Pointer {
    // ... другой код ...

    // Выделяем память, увеличивая указатель
    p := a.next
    a.next = unsafe.Add(a.next, size)
    a.left -= words

    return p
}

Примечательно, что мы не заменяем a.left на указатель конца, из-за сравнения if a.left < words. И мы не можем избежать вычитания a.left -= words, так как оно необходимо для корректной работы этого сравнения

Насколько стало лучше?

BenchmarkArena/int/arena-48         780.07 MB/s
BenchmarkArena/int/new-48           383.16 MB/s
BenchmarkArena/[2]int/arena-48      1245.73 MB/s
BenchmarkArena/[2]int/new-48        530.39 MB/s
BenchmarkArena/[64]int/arena-48     7684.39 MB/s
BenchmarkArena/[64]int/new-48       2679.94 MB/s
BenchmarkArena/[1024]int/arena-48   8859.99 MB/s
BenchmarkArena/[1024]int/new-48     2611.33 MB/s

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

Оказывается, есть куда более эффективный способ улучшить производительность.

Барьеры записи

Вот ассемблерный код, который Go сгенерировал для этой функции (сильно сокращенный), с аннотациями соответствующих строк исходного кода на Go.

TEXT (*Arena).Alloc(SB)
  CMPQ    SP, 0x10(R14)
  JBE     moreStack  ; Stack growth prologue.
  PUSHQ   BP
  MOVQ    SP, BP
  SUBQ    $0x58, SP

  ; size = (size + mask) &^ mask
  LEAQ    0x7(BX), DX
  ANDQ    $-0x8, DX
  ; words := size / wordBytes
  MOVQ    DX, SI
  SHRQ    $0x3, DX

  ; if a.left < words
  CMPQ    0x8(AX), DX
  JAE     alloc

  MOVQ    AX, 0x68(SP)
  MOVQ    SI, 0x48(SP)
  MOVQ    DX, 0x40(SP)

  ; nextPow2(words)
  MOVZX   runtime.x86HasPOPCNT(SB), DI
  TESTL   DI, DI
  JE      1f
  XORL    DI, DI
  POPCNTQ DX, DI
  JMP     2f
1:
  MOVQ    DX, AX
  CALL    math/bits.OnesCount(SB)
  MOVQ    0x40(SP), DX
  MOVQ    0x48(SP), SI
  MOVQ    AX, DI
  MOVQ    0x68(SP), AX
2:
  CMPQ    DI, $0x1
  JE      1f
  BSRQ    DX, CX
  MOVQ    $-0x1, DI
  CMOVE   DI, CX
  INCQ    CX
  MOVL    $0x1, DI
  SHLQ    CL, DI
  CMPQ    CX, $0x40
  SBBQ    R8, R8
  ANDQ    R8, DI
  MOVQ    DI, DX
1:
  MOVQ    0x10(AX), CX
  SHLQ    $0x1, CX

  ; a.cap = max(minWords, a.cap*2, nextPow2(words))
  CMPQ    CX, $0x8
  MOVL    $0x8, BX
  CMOVA   CX, BX
  CMPQ    DX, BX
  CMOVA   DX, BX
  MOVQ    BX, 0x10(AX)

  ; a.next = a.allocChunk(a.cap)
  CALL    github.com/mcy/go-arena.(*Arena).allocChunk(SB)
  CMPL    runtime.writeBarrier(SB), $0x0
  JNE     1f
  MOVQ    0x68(SP), DX
  JMP     2f
1:
  CALL    runtime.gcWriteBarrier2(SB)
  MOVQ    AX, 0(R11)
  MOVQ    0x68(SP), DX
  MOVQ    0(DX), R8
  MOVQ    R8, 0x8(R11)
2:
  MOVQ    AX, 0(DX)

  ; a.left = a.cap
  MOVQ    0x10(DX), R8
  MOVQ    R8, 0x8(DX)
  MOVQ    0x28(DX), CX
  MOVQ    0x20(DX), BX
  INCQ    BX
  MOVQ    0x18(DX), R8
  CMPQ    CX, BX
  JAE     2f

  ; a.chunks = append(a.chunks, a.next)
  MOVQ    AX, 0x50(SP)
  MOVQ    R8, AX
  MOVL    $0x1, DI
  LEAQ    0x28f70(IP), SI
  CALL    runtime.growslice(SB)
  MOVQ    0x68(SP), DX
  MOVQ    CX, 0x28(DX)
  CMPL    runtime.writeBarrier(SB), $0x0
  JE      1f
  CALL    runtime.gcWriteBarrier2(SB)
  MOVQ    AX, 0(R11)
  MOVQ    0x18(DX), CX
  MOVQ    CX, 0x8(R11)
1:
  MOVQ    AX, 0x18(DX)
  MOVQ    AX, R8
  MOVQ    0x50(SP), AX
2:
  MOVQ    BX, 0x20(DX)
  CMPL    runtime.writeBarrier(SB), $0x0
  JE      1f
  CALL    runtime.gcWriteBarrier2(SB)
  MOVQ    AX, 0(R11)
  MOVQ    -0x8(R8)(BX*8), CX
  MOVQ    CX, 0x8(R11)
1:
  MOVQ    AX, -0x8(R8)(BX*8)
  MOVQ    DX, AX
  MOVQ    0x40(SP), DX
  MOVQ    0x48(SP), SI

alloc:
  ; p := a.next
  MOVQ    0(AX), CX

  ; a.next = unsafe.Add(a.next, size)
  LEAQ    0(CX)(SI*1), BX
  CMPL    runtime.writeBarrier(SB), $0x0
  JE      1f
  CALL    runtime.gcWriteBarrier2(SB)
  MOVQ    BX, 0(R11)
  MOVQ    0(AX), SI
  MOVQ    SI, 0x8(R11)
1:
  MOVQ    BX, 0(AX)

  ; a.left -= words
  LEAQ    0(CX)(SI*1), BX
  SUBQ    DX, 0x8(AX)

  ; return p
  MOVQ    CX, AX
  ADDQ    $0x58, SP
  POPQ    BP
  RET

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

Барьер записи - это механизм синхронизации пользовательского кода со сборщиком мусора. Go генерирует код барьера записи каждый раз, когда сохраняется тип, содержащий указатели. Например, запись в **int, string или []int требует барьера записи.

Барьеры записи реализованы следующим образом:

  • Проверяется runtime.writeBarrier, который определяет, необходим ли барьер записи (только когда GC находится в фазе пометки). В противном случае выполняется переход для пропуска барьера записи.
  • Вызывается одна из функций runtime.gcWriteBarrierN, где N - количество указателей, о которых нужно уведомить GC.
  • Эта функция вызывает runtime.gcWriteBarrier, которая возвращает буфер, в который следует записывать указатели для последующего отслеживания GC.
  • Происходит фактическая запись.

Барьер записи требуется в случаях, подобных следующему. Рассмотрим такой код:

func alloc(n **int) {
  *n = new(int)
}

Эта функция вызовет runtime.newobject для выделения восьми байт памяти. Полученный указатель будет возвращён в регистре rax. Затем функция сохранит значение rax в переменную n и завершит работу. Если мы исследуем эту функцию через Godbolt, то обнаружим, что она действительно генерирует барьер записи:

TEXT x.alloc
  CMPQ    SP, 16(R14)
  JLS     growStack
  PUSHQ   BP
  MOVQ    SP, BP
  SUBQ    $16, SP

  MOVQ    AX, main.n+32(SP)

  ; new(int)
  LEAQ    type:int(SB), AX
  CALL    runtime.newobject(SB)

  MOVQ    main.n+32(SP), CX
  TESTB   AL, (CX)

  ; This is the write barrier.
  CMPL    runtime.writeBarrier(SB), $0
  JEQ     skip
  MOVQ    (CX), DX
  CALL    runtime.gcWriteBarrier2(SB)
  MOVQ    AX, (R11)
  MOVQ    DX, 8(R11)
skip:
  MOVQ    AX, (CX)  ; The actual store.

  ADDQ    $16, SP
  POPQ    BP
  RET

growStack:
  NOP
  MOVQ    AX, 8(SP)
  CALL    runtime.morestack_noctxt(SB)
  MOVQ    8(SP), AX
  JMP     x.alloc

Обратите внимание, что записываются два указателя: указатель, возвращённый new(int), и предыдущее значение n. Это гарантирует, что независимо от того, в какой момент сборщик мусора проверяет n, он увидит оба значения во время фазы пометки.

Однако это не требуется, если соответствующие указатели уже доступны другим способом, именно это мы и реализовали в нашей арене (благодаря срезу chunks). Поэтому барьер записи избыточен.

Но как от него избавиться? Существует директива //go:nowritebarrier, но её нельзя использовать вне пакетов, внесённых в белый список компилятора. К тому же она не отключает барьеры записи, а лишь выводит диагностическое сообщение, если они генерируются.

Вспомним, что барьеры записи возникают только при сохранении указателей, поэтому мы можем просто заменить next unsafe.Pointer на next uintptr.

type Arena struct {
  next      uintptr // Реальный указатель
  left, cap uintptr

  chunks []unsafe.Pointer
}

func (a *Arena) Alloc(size, align uintptr) unsafe.Pointer {
  mask := wordBytes - 1
  size = (size + mask) &^ mask
  words := size / wordBytes

  if a.left < words {
    a.cap = max(minWords, a.cap*2, nextPow2(words))

    p := a.allocChunk(a.cap)
    a.next = uintptr(p)
    a.left = a.cap
    a.chunks = append(a.chunks, p)
  }

  p := a.next
  a.next += size
  a.left -= words

  return unsafe.Pointer(p)
}

go vet не одобряет такой подход, потому что не понимает, что мы умнее него. Но делает ли это код быстрее? Чтобы приблизить тесты к реальным условиям, я создал отдельный вариант бенчмарков, который интенсивно нагружает сборщик мусора в отдельной горутине (G):

go func() {
  for { runtime.GC() }
}()

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

# Before
BenchmarkArena/int/arena-48         169.09 MB/s
BenchmarkArena/int/new-48           84.73 MB/s
BenchmarkArena/[2]int/arena-48      309.40 MB/s
BenchmarkArena/[2]int/new-48        120.23 MB/s
BenchmarkArena/[64]int/arena-48     1954.16 MB/s
BenchmarkArena/[64]int/new-48       950.48 MB/s
BenchmarkArena/[1024]int/arena-48   3341.13 MB/s
BenchmarkArena/[1024]int/new-48     1413.26 MB/s

# After
BenchmarkArena/int/arena-48         195.58 MB/s
BenchmarkArena/int/new-48           83.67 MB/s
BenchmarkArena/[2]int/arena-48      352.49 MB/s
BenchmarkArena/[2]int/new-48        120.13 MB/s
BenchmarkArena/[64]int/arena-48     1987.22 MB/s
BenchmarkArena/[64]int/new-48       903.78 MB/s
BenchmarkArena/[1024]int/arena-48   3342.67 MB/s
BenchmarkArena/[1024]int/new-48     1439.99 MB/s

Отделяемся от кучи

Еще одним источником замедления является то, что при каждой аллокации в куче Go вынужден явно очищать весь выделенный блок памяти, поскольку он содержит указатели. Если профилировать этот код, окажется, что значительное время тратится на выполнение runtime.memclrNoHeapPointers. Поскольку мы выделяем блоки памяти фиксированного размера, мы можем использовать массив sync.Pool для распределения затрат на аллокацию и очистку блоков.

Прежде всего, нам понадобится некоторый элемент в этом массиве пулов для каждого размера памяти, который мы выделяем. Затем, необходимо установить финализатор для арены, чтобы освободить память после завершения работы. Наконец, мы можем изменить контракт метода Alloc, потребовав от вызывающей стороны очищать значение, и изменить New, принимая значение в качестве аргумента:

func New[T any](a Allocator, v T) *T {
  p := (*T)(a.Alloc(unsafe.Sizeof(v), unsafe.Alignof(v)))
  *p = v
  return p
}

Прелесть этого подхода в том, что он позволяет избежать очистки значения, если вместо него будет записано ненулевое значение.

Собрав всё вместе, получим следующую реализацию:

var pools [64]sync.Pool

func init() {
  for i := range pools {
    pools[i].New = func() any {
      return reflect.New(reflect.StructOf([]reflect.StructField{
        {
          Name: "A",
          Type: reflect.ArrayOf(1<<i, reflect.TypeFor[uintptr]()),
        },
        {Name: "P", Type: reflect.TypeFor[unsafe.Pointer]()},
      })).UnsafePointer()
    }
  }
}

func (a *Arena) allocChunk(words uintptr) unsafe.Pointer {
  log := bits.TrailingZeros(uint(words))
  chunk := pools[log].Get().(unsafe.Pointer)

  // Смещаемся к концу блока и записываем указатель на арену
  end := unsafe.Add(chunk, words*unsafe.Sizeof(uintptr(0)))
  *(**Arena)(end) = a

  // Если это первый выделенный блок, устанавливаем финализатор
  if a.chunks == nil {
    runtime.SetFinalizer(a, (*Arena).finalize)
  }

  // Сохраняем блок в соответствующей позиции a.chunks,
  // соответствующей его размеру (log) для простоты идентификации
  a.chunks = append(a.chunks, make([]unsafe.Pointer, log+1-len(a.chunks))...)
  a.chunks[log] = chunk

  return chunk
}

func (a *Arena) finalize() {
  for log, chunk := range a.chunks {
    if chunk == nil {
      continue
    }

    words := uintptr(1) << log
    end := unsafe.Add(chunk, words*unsafe.Sizeof(uintptr(0)))
    *(**Arena)(end) = nil // Гарантируем отсутствие утечки арены

    pools[log].Put(chunk)
  }
}

Как наши изменения повлияли на производительность?

BenchmarkArena/int/arena-48        1260.73 MB/s
BenchmarkArena/int/new-48          712.94 MB/s
BenchmarkArena/[2]int/arena-48     2457.27 MB/s
BenchmarkArena/[2]int/new-48       1167.57 MB/s
BenchmarkArena/[64]int/arena-48    4491.49 MB/s
BenchmarkArena/[64]int/new-48      6800.76 MB/s
BenchmarkArena/[1024]int/arena-48  3992.32 MB/s
BenchmarkArena/[1024]int/new-48    4320.65 MB/s

Что ж. Вот это неожиданность. Производительность значительно улучшилась для небольших аллокаций, но стала хуже для очень больших. Пока не совсем понятно, почему так происходит, но обратите внимание, что обычный new тоже стал работать быстрее. Это происходит поскольку аллокации из арены живут дольше и сборщик мусора теперб ведёт себя несколько иначе, распределяя часть затрат на выделение очень больших объектов через new.

Чтобы понять, имеет ли смысл такая оптимизация, потребуется дополнительное профилирование. Альтернативный вариант - вручную управлять повторным использованием арены, добавив крайне небезопасную функцию Reset(), которая заставляет арену вести себя так, как будто она только что создана, но сохраняет все выделенные блоки памяти. Это аналогично обнулению среза: x = x[:0].

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

Реализация этого подхода очень проста:

func (a *Arena) Reset() {
  a.next, a.left, a.cap = 0, 0, 0  // Сбрасываем указатели, оставляя chunks
}

func (a *Arena) allocChunk(words uintptr) unsafe.Pointer {
  log := bits.TrailingZeros(uint(words))
  if len(a.chunks) > log {
    // Если мы уже выделяли блок такого размера в предыдущем цикле арены,
    // просто возвращаем его.
    //
    // Важно: арена никогда не пытается выделить блок одного размера дважды
    // между вызовами Reset().
    return a.chunks[log]
  }

  // ... достально код ...
}

Немного модифицируем бенчмарк

b.Run("arena", func(b *testing.B) {
  b.SetBytes(n)
  a := new(arena.Arena)
  for b.Loop() {
    a.Reset()  // Важное изменение
    for range runs {
      sink = arena.New[T](a)
    }
  }
})

И смотрим результаты

BenchmarkArena/int/arena-48         2376.01 MB/s
BenchmarkArena/int/new-48           377.64 MB/s
BenchmarkArena/[2]int/arena-48      4314.98 MB/s
BenchmarkArena/[2]int/new-48        530.62 MB/s
BenchmarkArena/[64]int/arena-48     10496.49 MB/s
BenchmarkArena/[64]int/new-48       3959.85 MB/s
BenchmarkArena/[1024]int/arena-48   9735.19 MB/s
BenchmarkArena/[1024]int/new-48     6160.73 MB/s

Это колоссальное улучшение производительности! Есть несколько причин такой скорости. Во-первых, нам больше не нужно ждать, пока GC освободит память старых арен для повторного использования. Во-вторых, быстрый путь выполнения теперь крайне эффективен и не требует синхронизации.

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

Realloc (переаллокация)

Go не предоставляет простого механизма для "перераспределения" памяти, подобного функции realloc() в C. Это связано с отсутствием механизма явного освобождения указателей, который необходим для абстракции переаллокации.

Но нам уже не важна безопасность, поэтому мы можем реализовать переаллокацию в нашей арене. Правда, доступная нам переаллокация будет довольно примитивной: если блок оказался последним выделенным, мы можем увеличить его размер. В противном случае просто выделяем новый блок, не освобождая старый.

Это позволяет реализовать "срезы арены" (arena slices), которые можно наращивать добавлением элементов - такие срезы не будут вызывать переаллокацию при росте, пока в арену не поместят другие объекты.

Реализация Realloc будет выглядеть примерно так:

func (a *Arena) Realloc(
  ptr unsafe.Pointer,         // Указатель на перераспределяемую память
  oldSize, newSize, align uintptr, // Старый и новый размеры, выравнивание
) unsafe.Pointer {
  // Выравнивание размеров
  mask := wordBytes - 1
  oldSize = (oldSize + mask) &^ mask
  newSize = (newSize + mask) &^ mask

  // Если новый размер меньше или равен старому - возвращаем исходный указатель
  if newSize <= oldSize {
    return ptr
  }

  // Проверяем, является ли этот блок последним выделенным
  if a.next - oldSize == uintptr(ptr) {
    // Проверяем, достаточно ли места для расширения
    need := (newSize - oldSize) / wordBytes
    if a.left >= need {
      // Расширяем на месте
      a.left -= need
      return ptr
    }
  }

  // Если расширение на месте невозможно - выделяем новую память и копируем данные
  new := a.Alloc(newSize, align)
  copy(
    unsafe.Slice((*byte)(new), newSize),  // Новый блок
    unsafe.Slice((*byte)(ptr), oldSize),   // Старый блок
  )

  return new
}

Таким образом, при добавлении элементов в наш срез арены мы можем использовать a.Realloc() для его увеличения. Однако это не будет работать, если базовый указатель среза отличается от оригинального адреса, возвращённого Alloc или Realloc.

В качестве домашнего задания предлагается:

  • Реализовать тип Slice[T], использующий арену для аллокации памяти.
  • Обеспечить работу не только с базовым указателем, но и с любым адресом внутри последней аллокации. Это потребует дополнительного учёта метаданных.

Собираем все вместе

Вот полный исходный код, всего что мы реализовали за исключением функции для переаллокации

package arena

import (
  "math/bits"
  "reflect"
  "unsafe"
)

// New создает новый объект типа T в арене
func New[T any](a *Arena, v T) *T {
  p := (*T)(a.Alloc(unsafe.Sizeof(v), unsafe.Alignof(v)))
  *p = v
  return p
}

// Arena представляет аллокатор памяти
type Arena struct {
  next   unsafe.Pointer  // Следующий свободный указатель
  left   uintptr         // Осталось места в текущем блоке (в словах)
  cap    uintptr         // Емкость текущего блока (в словах)
  chunks []unsafe.Pointer // Выделенные блоки памяти
}

const (
  maxAlign uintptr = 8 // Для 64-битных систем
  minWords uintptr = 8 // Минимальный размер выделения
)

// Alloc выделяет память в арене
func (a *Arena) Alloc(size, align uintptr) unsafe.Pointer {
  // Выравнивание размера
  mask := maxAlign - 1
  size = (size + mask) &^ mask
  words := size / maxAlign

  // Проверка наличия свободного места
  if a.left < words {
    // Определение нового размера блока
    a.cap = max(minWords, a.cap*2, nextPow2(words))
    a.next = a.allocChunk(a.cap)
    a.left = a.cap
    a.chunks = append(a.chunks, a.next)
  }

  // Выделение памяти
  p := a.next
  a.next = unsafe.Add(a.next, size)
  a.left -= words

  return p
}

// Reset сбрасывает состояние арены
func (a *Arena) Reset() {
  a.next, a.left, a.cap = 0, 0, 0
}

// Пул блоков памяти
var pools [64]sync.Pool

func init() {
  for i := range pools {
    pools[i].New = func() any {
      return reflect.New(reflect.StructOf([]reflect.StructField{
        {
          Name: "X0",
          Type: reflect.ArrayOf(1<<i, reflect.TypeFor[uintptr]()),
        },
        { Name: "X1", Type: reflect.TypeFor[unsafe.Pointer]() },
      })).UnsafePointer()
    }
  }
}

// allocChunk выделяет новый блок памяти
func (a *Arena) allocChunk(words uintptr) unsafe.Pointer {
  log := bits.TrailingZeros(uint(words))
  if len(a.chunks) > log {
    return a.chunks[log]
  }

  chunk := pools[log].Get().(unsafe.Pointer)

  // Запись ссылки на арену в конец блока
  end := unsafe.Add(chunk, words*unsafe.Sizeof(uintptr(0)))
  *(**Arena)(end) = a

  // Установка финализатора для первого блока
  if a.chunks == nil {
    runtime.SetFinalizer(a, (*Arena).finalize)
  }

  // Сохранение блока
  a.chunks = append(a.chunks, make([]unsafe.Pointer, log+1-len(a.chunks))...)
  a.chunks[log] = chunk

  return chunk
}

// finalize освобождает память при сборке мусора
func (a *Arena) finalize() {
  for log, chunk := range a.chunks {
    if chunk == nil {
      continue
    }

    words := uintptr(1) << log
    end := unsafe.Add(chunk, words*unsafe.Sizeof(uintptr(0)))
    *(**Arena)(end) = nil // Очистка ссылки

    pools[log].Put(chunk) // Возврат в пул
  }
}

// nextPow2 возвращает следующую степень двойки
func nextPow2(n uintptr) uintptr {
    return uintptr(1) << bits.Len(uint(n))
}

Существуют и другие возможные оптимизации, которые я не рассматривал. Например, арены можно повторно использовать: после завершения работы с ареной её можно "сбросить" (reset) и поместить в sync.Pool. Такая арена не будет требовать новых блоков от сборщика мусора, повторно используя ранее выделенные (что потенциально снижает затраты на постоянное обнуление памяти).

Выше я говорил, что эта реализация сильно зависит от внутренних особенностей языка Golang. Насколько вероятно, что они изменятся в будущем? Требование, чтобы аллокации "знали" свою структуру, обусловлено существованием unsafe.Pointer, а необходимость сохранять всю аллокацию живой при наличии указателя на любую её часть проистекает из возможности разделения и изменения срезов (slices). Если срез попадает в кучу (и, следовательно, становится доступным нескольким горутинам), координация копий для уменьшения размера потребовала бы гораздо более сложной реализации, чем текущая система барьеров записи.

И, на мой взгляд, можно с уверенностью сказать, что здесь нас защищает закон Хайрума. ;)