Скаляризация Map.get(new CompositeKey(key1, key2, ...))

Один из самых интересных примеров из моего последнего доклада на jug-ekb (на JPoint этот пример не влез) это скаляризация композитных ключей в поиске по словарю:

final Map<CompositeKey, V> map = ...;

...
    final CompositeKey key = new CompositeKey( key1, key2, ... );
    final V value = map.get( key );
...

Это весьма частый паттерн: нам нужно отображение вида [key1, key2, key3...] -> value. Наиболее простой и очевидный способ это сделать: как раз таки использовать композитный ключ (кортеж из ключей). Решение отличное по всем показателям, кроме одного: для каждого поиска в словаре придется создавать новый объект.

Иногда хочется этого избежать. Тогда приходится организовывать иерархические словари-словарей-словарей... — я видел много таких примеров (поиск по текущему проекту находит аж 5 различных реализаций чего-то вроде TwoKeyMap/ThreeKeyMap). Это решение мне самому не нравится (хотя 2 из 5 упомянутых реализаций — мои). Например, словарь-словарей это совсем другой тип, нежели словарь-словарей-словарей — расширение ключа на одну позицию не является прозрачным, окружающий код специализирован под конкретное число ключей. У этого решения, как правило, хуже производительность: во-первых, потому что приходится делать 2-3-4... поиска вместо одного, а во-вторых потому, что соответствующие словари обычно получаются очень неравномерно заполненными. Скажем, key1 может иметь всего 3-4 разных значения, а лежать они будут в таблице размером по-умолчанию (16). При этом, скажем, ключей вида ("A", *) всего 2, а вида ("B", *) — 10000. То есть часто образуется много сильно недозаполненных таблиц, вперемешку с несколькими нормальными. Будь все эти ключи в одной таблице — все эти неоднородности сильно сгладились бы.

В общем, хотелось бы и чтобы все было хорошо, и за это не приходилось бы ничего платить. Может ли JIT сам скаляризовать аллокацию ключа в сценарии типа map.get( new CompositeKey( key1, key2, ... ) )?

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

Как выглядит сам бенчмарк: я генерирую пул строк, из них образую парные композитные ключи (StringKey), которыми наполняю хэш-таблицу. Внутри главного цикла бенчмарка я создаю ключи таким образом, чтобы они присутствовали или отсутствовали в словаре с какой-то вероятностью (hit probability — один из параметров сценария), и запрашиваю у словаря соответствующие значения.

private final Pool<String> keys = poolOf( randomUniqueStrings( 2 * SIZE ) );

private final Pool<String> nonExistentKeys = poolOf( randomUniqueStrings( 2 * SIZE ) );

//...

public long run() {
 final boolean successful = ( rnd.nextDouble() <= SUCCESSFUL_LOOKUPS_PROBABILITY );
 
 final String key1;
 final String key2;
 if( successful ) {
  key1 = keys.next();
  key2 = keys.next();
 } else {
  key1 = nonExistentKeys.next();
  key2 = nonExistentKeys.next();
 }

 final String value = map.get( new StringKey( key1, key2 ) );

 if( value == null ) {
  return 0;
 } else {
  return 1;
 }
}

Что получается на выходе: иногда ключи скаляризуются почти стабильно, для целого спектра размеров таблицы (1,2,4...65) и спектра hit probability=(0%, 50%, 100%). А иногда — почти стабильно не скаляризуются. А иногда вообще какая-то ересь получается: сценарий аллоцирует что-то вроде 120 байт на вызов, вместо обычных 24-х — что уже вообще ни в какие ворота. Поведение меняется от небольших изменений в коде, и от того, использую ли я в качестве ключей свой собственный класс StringKey, или заменяю его на ImmutablePair из apache commons.

Дальше я буду идти от теории к практике — порассуждаю про код, и подтвержу свои рассуждения примерами. В реальности все происходило ровно наоборот: я сначала поставил много экспериментов, результаты которых были довольно запутывающими, а потом уже построил модель, что и как происходит. Такая история, конечно, гораздо увлекательнее, но пересказывать ее уж очень долго выходит.

Начнем с кода HashMap.get(). В 1.8.0_102 он выглядит так:

public V get(Object key) {
        Node<K,V> e;
        return (e = getNode(hash(key), key)) == null ? null : e.value;
    }

    final Node<K,V> getNode(int hash, Object key) {
        Node<K,V>[] tab; Node<K,V> first, e; int n; K k;
        if ((tab = table) != null && (n = tab.length) > 0 &&
            (first = tab[(n - 1) & hash]) != null) {
            if (first.hash == hash && // always check first node
                ((k = first.key) == key || (key != null && key.equals(k))))
                return first;
            if ((e = first.next) != null) {
                if (first instanceof TreeNode)
                    return ((TreeNode<K,V>)first).getTreeNode(hash, key);
                do {
                    if (e.hash == hash &&
                        ((k = e.key) == key || (key != null && key.equals(k))))
                        return e;
                } while ((e = e.next) != null);
            }
        }
        return null;
    }

Что мы можем здесь увидеть — помимо того, что автору неплохо бы пересмотреть свой стиль именования переменных? Мы можем увидеть, что здесь две ветки — "прямое попадание", и процедура разрешения коллизий, и внутри второй ветки — еще две ветки (поиск в связном списке или дереве). Чтобы скаляризация имела шанс сработать, нам необходимо (но не достаточно), чтобы заинлайнились все методы, в которые уходит интересующий нас объект (key), до самого низа. В данном случае, как минимум, должны заинлайниться hashCode(), оба вызова .equals() и .getTreeNode().

А что нужно, чтобы метод заинлайнился? О, это большая тема :) Инлайнинг — одна из самых потенциально эффективных оптимизаций в JIT-ах, так что эвристик у нее немало. К счастью, все они нам обычно не нужны, достаточно сильно сокращенной версии:

1. Если метод-кандидат виртуальный — он должен быть девиртуализован. Для девиртуализации необходимо либо чтобы CHA показал, что сейчас в JVM в принципе загружена только одна реализация, либо чтобы профиль, собранный во время исполнения (TypeProfile) показал, что в данном конкретном месте (вероятно) вызывается не более двух реализаций виртуального метода.
2. Если метод-кандидат достаточно горячий (горячесть вычисляется как по абсолютному количеству вызовов, так и по тому, насколько часто метод-кандидат вызывается из того метода, в который мы его собираемся вклеивать), то его размер в байт-кодах должен быть не больше FreqInlineSize(=325), и размер в машинных кодах (если он уже был скомпилирован изолированно) не больше SmallInlineSize(=2000).
3. Если метод недостаточно горячий, то, возможно, он хотя бы маленький (< MaxInlineSize(=35))? Но, при этом, он должен все-таки выполниться хоть сколько-то: как минимум, MinInliningThreshold(=250) раз. А иначе какой смысл вклеивать?
4. ...Но если какая-то ветка не выполнилась вообще ни разу (за время профилирования), то C2 выкинет ее код вообще, а на случай "а вдруг" поставит вместо нее т.н. uncommon trap: перезапуск выполнения метода в режиме интерпретатора. Эта часть уже не является частью политики инлайнинга, это отдельная, и очень мощная спекулятивная оптимизация, но она очень помогает скаляризации: ведь если весь код ветки выброшен, то он не анализируется на предмет убегания ссылок.

(Обычному java-программисту непривычно думать о размере метода в байт-кодах, поэтому цифры типа 325 и 35 могут выглядеть довольно абстрактно. Точного соотношения между размером джава-кода и размером байт-кода, разумеется, нет, но по порядку величины можно сказать так: 325 байт это довольно большой метод, примерно на полстраницы-страницу, а 35 байт это коротенький метод, на 3-4 строчки)

...Итак, у нас есть код поиска ключа в хэш-таблице, который состоит из нескольких этапов (веток): поиск нужной корзинки (bucket), сравнение (по хэш-коду и по .equals()), поиск по списку коллизий (по хэш-коду и по .equals()) или поиск по дереву коллизий (там вообще мрак и ужас). И поиск может завершиться на любом этапе, но вся идея хэш-таблицы как структуры данных в том, что как правило у нас списки коллизий короткие, или отсутствуют вовсе. То есть, если раскладка данных в хэш-таблице достаточно равномерная (==хэш-функция хорошая), то вероятность завершить поиск на первом же этапе довольно высока, вероятность попасть в ветку разрешения коллизий — заметно меньше, а вероятность наткнуться на список коллизий настолько длинный, чтобы он конвертировался в дерево — совсем маленькая. Более точно, можно выделить такие сценарии:

1. Попали в пустую корзинку: этот сценарий довольно часто (с вероятностью = 1 - load factor = 25%..60%) будет реализовываться, когда мы будем искать несуществующие в таблице ключи (hit probability << 100%). Здесь вызывается только key.hashCode(), ни один из 2х .equals(), ни .getTreeNode() не вызывается вообще
2. Попали в корзинку с единственной парой ключ-значение: этих сценариев будет большинство среди тех, когда мы вообще что-то найдем (==ключ присутствует в таблице). В этом случае мы дернем key.hashCode(), и первый .equals(), а второй .equals() и .getTreeNode() все так же не выполняются
3. Попали в корзинку с несколькими парами ключ-значение, но нашли нужный нам ключ с первой же попытки — наш маршрут через код совпадает с предыдущим
4. Наконец, мы попали в корзинку с несколькими парами, и нашли нужный ключ не с первой попытки, либо вообще его не нашли, просмотрев весь список коллизий. В этом случае мы успеем дернуть key.hashCode(), и оба .equals(), (.getTreeNode()все так же не у дел)
5. Для редкостных неудачников: мы попали в корзинку с >8 парами ключ-значение, при общем размере таблицы >64, и обнаружили, что корзинка сконвертирована в TreeBin — тут, наконец, мы заглянем во врата ада.getTreeNode()

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

Во-первых, очевидно, у нас могут доминировать сценарии №1. В этом варианте может оказаться, что ни один из двух .equals() не набрал достаточного количества вызовов, чтобы оказаться горячим, и тогда код .equals() должен быть ну очень коротким, чтобы пролезть под порогом инлайнинга для не горячих методов (35 байт). Например, такой вот, стандартный

public boolean equals( final Object o ) {
 if( this == o ) {
  return true;
 }
 if( !( o instanceof StringKey ) ) {
  return false;
 }

 final StringKey key = ( StringKey ) o;

 return item1.equals( key.item1 )
   && item2.equals( key.item2 );
}

компилируется в 55 байт, и не пролезает.

Как я уже упоминал, на этот сценарий реальнее всего попасть когда большинство искомых ключей в таблице отсутствуют. В самом деле, запуская сценарий с (HashMap.size=64, hit probability=0%) мы получаем грустную историю дискриминации толстяков:

@ 84   j.u.HashMap::get (23 bytes)   inline (hot)
 \-> TypeProfile (11045/11045 counts) = java/util/HashMap
  @ 2   j.u.HashMap::hash (20 bytes)   inline (hot)
    @ 9   ...StringKey::hashCode (19 bytes)   inline (hot)
      @ 4     j.l.String::hashCode (55 bytes)   inline (hot)
      @ 14    j.l.String::hashCode (55 bytes)   inline (hot)
  @ 6   j.u.HashMap::getNode (148 bytes)   inline (hot)
    @ 59   ...StringKey::equals (55 bytes)   too big       <--- 
    @ 126  ...StringKey::equals (55 bytes)   too big       <--- 

Видно, что ни один из двух .equals() не вклеился. too big это лишь один из вариантов, я наблюдал так же и executed < MinInliningThreshold

Если мы начнем увеличивать hit probability, то очень быстро "разогреем" первый .equals(): (size=64, hit probability=1%, 2%, ... )

@ 84   j.u.HashMap::get (23 bytes)   inline (hot)
 \-> TypeProfile (11045/11045 counts) = java/util/HashMap
  @ 2   j.u.HashMap::hash (20 bytes)   inline (hot)
    @ 9   ...StringKey::hashCode (19 bytes)   inline (hot)
      @ 4     j.l.String::hashCode (55 bytes)   inline (hot)
      @ 14    j.l.String::hashCode (55 bytes)   inline (hot)
  @ 6   j.u.HashMap::getNode (148 bytes)   inline (hot)
    @ 59   ...StringKey::equals (55 bytes)   inline (hot)    <---
      @ 29     j.l.String::equals (81 bytes)   (intrinsic)
      @ 43     j.l.String::equals (81 bytes)   (intrinsic)
    @ 126  ...StringKey::equals (55 bytes)   too big         <---

Увеличивая hit probability дальше, мы в какой-то момент разогреем и второй .equals(). Точный момент указать сложно, потому что ключи каждый раз генерируются заново, и раскладка по корзинкам получается в каждом запуске немного другая (а она зависит еще и от размера таблицы...), да и сбор профиля внутри JVM — тоже не детерминированный процесс. Вот сейчас у меня получилось достичь этого при hit probability=5%:

...
@ 84   j.u.HashMap::get (23 bytes)   inline (hot)
 \-> TypeProfile (112592/112592 counts) = java/util/HashMap
  @ 2   j.u.HashMap::hash (20 bytes)   inline (hot)
    @ 9   ...StringKey::hashCode (19 bytes)   inline (hot)
      @ 4     j.l.String::hashCode (55 bytes)   inline (hot)
      @ 14    j.l.String::hashCode (55 bytes)   inline (hot)
  @ 6   j.u.HashMap::getNode (148 bytes)   inline (hot)
    @ 59   ...StringKey::equals (55 bytes)   inline (hot)   <---
      @ 29     j.l.String::equals (81 bytes)   (intrinsic)
      @ 43     j.l.String::equals (81 bytes)   (intrinsic)
    @ 126  ...StringKey::equals (55 bytes)   inline (hot)   <---
      @ 29     j.l.String::equals (81 bytes)   (intrinsic)
      @ 43     j.l.String::equals (81 bytes)   (intrinsic)
...

И да — бинго! — в этом сценарии ключи начали скаляризоваться, код больше ничего не аллоцирует.

Среди прочего, интересно заметить, что в логе компиляции .getTreeNode() вообще не упоминается. Это все потому, что соответствующая ветка вообще не используется, и JIT заменил ее на uncommon trap — большая любезность с его стороны, а то у нас были бы околонулевые шансы добиться скаляризации в этом примере хоть когда-нибудь.

До сих пор все просто и логично. Вот только я точно помню, что в каких-то старых вариантах бенчмарка наблюдал стабильную скаляризацию для всего спектра hit probability, в том числе и = 0. Как такое могло происходить?

Оказывается (это "оказывается" стоило мне нескольких часов недосыпа), я неаккуратно генерировал ключи. В генерации пула "несуществующих" ключей была ошибка, из-за которой с некоторой заметной вероятностью такие ключи все же находились в словаре — а как мы видели, достаточно небольшой примеси "найденных" ключей, чтобы разогреть холодные ветки. Но этого мало, у меня еще и сами ключи генерировались с очень плохим распределением. Точнее, отдельные-то ключи имели вполне равномерное распределение, но в пары они собирались по порядку, и вот тут-то хэш-код становился буквально кратным 32-м. Это, конечно, создавало очень неравномерное распределение по корзинкам, и длинные списки коллизий, а длинные списки коллизий еще сильнее ускоряли разогрев соответствующей ветки.

Эти же длинные списки коллизий сыграли при переходе на ключи ImmutablePair из apache commons. С ImmutablePair я наблюдал аллокацию аж 120 байт на вызов, которым, вроде, совсем неоткуда там было взяться. А все дело в том, что, в отличие от простого и дубового StringKey, навороченный ImmutablePair implements Comparable, и при превращении цепочки коллизий в дерево, HashMap это пытается использовать, чтобы сделать дерево более красивым. И вот такой код, где-то внутри .getTreeNode():

static Class<?> comparableClassFor(Object x) {
    if (x instanceof Comparable) {
        Class<?> c; Type[] ts, as; Type t; ParameterizedType p;
        if ((c = x.getClass()) == String.class) // bypass checks
            return c;
        if ((ts = c.getGenericInterfaces()) != null) {   // <---- 
            for (int i = 0; i < ts.length; ++i) {
                if (((t = ts[i]) instanceof ParameterizedType) &&
                    ((p = (ParameterizedType)t).getRawType() ==
                     Comparable.class) &&
                    (as = p.getActualTypeArguments()) != null &&
                    as.length == 1 && as[0] == c) // type arg is c
                    return c;
            }
        }
    }
    return null;
}

что он делает? Дело в том, что даже если key1 instanceof Comparable && key2 instanceof Comparable, нельзя так просто взять, и вызвать key1.compareTo(key2). Нельзя потому, что, возможно, эти Comparable промеж собой не совместимы: может, один из них Comparable<A>, а другой Comparable<B>. Процитированный код как раз пытается понять, чем конкретно параметризован этот Comparable. И вот вызов c.getGenericInterfaces() как раз и порождает загадочные 120 байт

Что получается в итоге? Получается, что успех или неуспех скаляризации зависит от раскладки ключей в хэш-таблице. Интересно что: чем ближе раскладка к идеальной (с точки зрения производительности) "один ключ на одну корзину", тем более вероятно, что ветка разрешения коллизий окажется холодной, и скаляризация не произойдет. Более того, профилирование происходит в начале исполнения кода, и после окончательной компиляции (C2) без веских причин перепрофилирования не будет, а это значит, что версия, скомпилированная по "неудачному" профилю, собранному в начале исполнения, так и будет исполняться дальше, уже независимо от актуального профиля.

P.S. Первопричина многих подобных нестабильностей в том, эвристики инлайнинга жестко заточены на производительность: мы стараемся вклеивать те методы, от вклеивания которых ожидается наибольший прирост скорости исполнения, и не рвемся вклеивать те, от которых заметного прироста не ожидаем. Это не совсем то, что нужно для скаляризации: скаляризация-то срабатывает по принципу "все или ничего": все маршруты конкретного объекта должны быть доступны для анализа, и если хотя бы один, пусть и редкий, маршрут уходит за горизонт (в невклеенный метод), то скаляризации не сможет работать. С точки зрения производительности же вклеивать редкий метод нет смысла: прирост производительности минимальный, а код раздувается. Скаляризация выступает тут бедным родственником, который вынужден довольствоваться тем, что дают: иначе было бы логично явно учитывать и ее интересы в эвристиках инлайнинга (такой EA-guided inlining предлагался, кажется, еще в Choi99).