Байт-код — это просто! Как сделать DI по-настоящему быстрым

Привет! Меня зовут Григорий Юрков, и я уже несколько лет работаю в инфраструктурной команде Яндекс Маркета. Два года назад мы начали разрабатывать свой легковесный DI-фреймворк Scout, который предоставляет выразительный Kotlin DSL. Он не генерирует код, а делает всю работу в рантайме.

Недавний переход с compile-time-библиотеки Dagger 2 на нашу привёл к замедлению старта приложения. Подробнее об этом и о том пути, который мы прошли от идеи до публикации в опенсорс, можно прочитать в статье моего коллеги Александра Миронычева.

В этой статье мы будем подробно рассматривать то, как применение байт-кода помогло сохранить скорость на том же уровне и спасти проект по миграции на Scout.

Весь код в этой статье упрощён для понимания. Также не стоит забывать, что библиотека постоянно изменяется и улучшается. Если вы хотите посмотреть настоящий код, который представлен в библиотеке, то прошу в наш GitHub.

Почему стало медленнее

Чтобы вникнуть в проблему, которая возникла перед нами, предлагаю немного познакомиться с библиотекой. Вот так выглядит самый простой пример кода на Scout:

class MyClass  val scope = scope("my-scope") {     factory<String> { "my-string" }     factory<MyClass> { MyClass() } }  fun main() {     println(scope.get<String>()) // "my-string"     println(scope.get<MyClass>()) // MyClass }

Мы создаём factory для двух типов — String и MyClass. Эти factory лежат в общем Scope под названием my-scope. С помощью get()мы можем получить указанные типы, используя зарегистрированные в Scope фабрики.

Простейшая реализация Scope выглядела бы так:

typealias Key = Class<*>  object Keys {     @JvmStatic fun create(keyClass: Class<*>): Key = keyClass }  class Scope {     internal val factories = HashMap<Key, () -> Any?>()      inline fun <reified T> factory(noinline factory: () -> T) {         val key: Key = Keys.create(T::class.java)         factories[key] = factory     }              inline fun <reified T> get() {         val key: Key = Keys.create(T::class.java)         return scope.factories[key].invoke()     } }

Этот код сильно упрощён, но смысл остаётся тот же: мы просто берем HashMap и по классу складываем туда наши factory, а потом по тому же классу достаём и аллоцируем объект. 

Такой код работает сильно медленнее Dagger 2, и тому есть две основные причины:

1. Мы используем Class в качестве ключа.

2. Вызов T::class.java триггерит class loading.

С первым вроде всё понятно: HashMap постоянно вызывает hashCode() и equals(), чтобы класть и доставать объекты. Для Class эти функции не такие быстрые, как хотелось бы.

Со вторым всё интереснее. Когда мы вызываем MyClass::class.java, чтобы получить Class, мы триггерим загрузку этого класса в оперативную память. Это означает, что Java Virtual Machine (или Android Runtime) нужно прочитать класс из JAR- или DEX-файла и запарсить его. Этот процесс очень долгий, и мы бы хотели его избежать как минимум на этапе создания factory.

Как сделать быстрее

Мы можем убить двух птиц одним камнем, заменив ключи нашей мапы на Int. Тогда мы избавимся от HashMap в пользу чего-то более быстрого, а также избавимся от class loading. Осталось самое сложное: как заменить Class на Int?

Первое, что пришло в голову, — написать плагин на компилятор Kotlin, чтобы он заменял Class на Int. Но, во-первых, было совершенно непонятно, как подходить к этой проблеме. Во-вторых, мы опасались, что будем сильно зависеть от версии Kotlin и его внутреннего API, а это непременно доставит нам хлопот в будущем.

Тогда мне пришла в голову идея: а что, если заменить
typealias Key = Class<> на typealias Key = Int и кидать ошибку в функции create? 

Попробуем изменить код: 

typealias Key = Int // Замена с Class<*> на Int  object Keys {     @JvmStatic fun create(keyClass: Class<*>): Key {         throw NotImplementedError() // Кидаем ошибку     } }      // Scope остался неизменным class Scope {     internal val factories = HashMap<Key, () -> Any?>()      inline fun <reified T> factory(noinline factory: () -> T) {         val key: Key = Keys.create(T::class.java)         factories[key] = factory     }              inline fun <reified T> get() {         val key: Key = Keys.create(T::class.java)         return scope.factories[key].invoke()     } }

Вы спросите: «А зачем кидать ошибку? Как тогда это должно работать?»

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

Теперь посмотрим на байт-код нашего примера, преобразованный в Java-код для наглядности:

class MainKt {     public static final Scope scope = new Scope();      static {         int key;          // Это байт-код для factory<String>          key = Keys.create(String.class); // Сейчас этот метод кинет исключение, если всё оставить как есть         scope.factories.put(key, new Lambda1());          // А это байт-код для factory<MyClass>         key = Keys.create(MyClass.class);         scope.factories.put(key, new Lambda2());      }          public static void main(String[] args) {         int key;                  // Это вывод строки         key = Keys.create(String.class);         System.out.println(scope.factories.get(key).invoke());                  // А это вывод MyClass         key = Keys.create(MyClass.class);         System.out.println(scope.factories.get(key).invoke());     } } 

Теперь, когда процессор будет проходиться по байт-коду, он будет встречать вызовы типа Keys.create(SomeClass.class) и заменять их на константы. Условимся, что String — это 1, а MyClass — это 2. В итоге получаем следующий изменённый байт-код:

class MainKt {     public static final Scope scope = new Scope();      static {         int key;          // Это байт-код для factory<String>          key = 1;         scope.factories.put(key, new Lambda1());          // А это байт-код для factory<MyClass>         key = 2;         scope.factories.put(key, new Lambda2());      }          public static void main(String[] args) {         int key;                  // Это вывод строки         key = 1;         System.out.println(scope.factories.get(key).invoke());                  // А это вывод MyClass         key = 2;         System.out.println(scope.factories.get(key).invoke());     } }

Вызовов метода create нигде не осталось, так как мы всё заменили на константы. Исключения не кидаются, и всё работает!

А дальше я расскажу, как написать такой процессор и правильно его применять.

Как устроен байт-код

Байт-код — стандартное промежуточное представление, в которое компьютерная программа может быть переведена автоматическими средствами. Так говорит «Википедия». В случае JVM он хранится в .class-файлах, которые чаще всего можно найти упакованными в JAR-файлы (по сути, это обычный ZIP-архив).

Прежде чем перейдём к устройству байт-кода, расскажу о том, что такое обратная польская запись (ОПЗ). Дело в том, что байт-код работает по такому же принципу. Если вы знаете, что это такое, то можете пропустить эту часть.

(1 + 2) * (3 + 4)

1 2 + 3 4 + *

Идём по математическому выражению с начала.  Если встречаем константу,          Tо кладем её в стек.  Иначе, если встречаем оператор или функцию (в ОПЗ это, в принципе, одни и те же вещи),          То достаём из стека нужное количество операндов, выполняем операцию и кладём в стек получившееся значение.  В конце в стеке остаётся наш результат.

JVM — это навороченная стековая машина, которая читает подряд инструкции, а потом кладёт в стек или достаёт из него значения. Приведу пример:

System.out.println(1 + 2 * 3);

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

// Стек: пустой GETFIELD java.lang.System.out // Загружаем значение из статического поля в стек // Стек: PrintStream LDC 2 // Загружаем константы в стек LDC 3 // Стек: PrintStream, 2, 3 IMULT // Выполняем умножение и кладём результат в стек // Стек: PrintStream, 6 LDC 1 // Загружаем единичку // Стек: PrintStream, 6, 1 IADD // Выполняем сложение // Стек: PrintStream, 7 INVOKEVIRTUAL println // Вызываем нестатический метод println на объекте PrintStream с аргументом 7 // Стек остался пустым, а в консоли напечаталось число 7 

Этот код познакомил нас со следующими базовыми инструкциями:

• LDC — загрузка константы в стек.

• IMULT, IADD — операции сложения и умножения. Буква I в начале означает операцию над целыми числами. Для double есть аналогичные с буквой D.

• INVOKEVIRTUAL — вызов нестатического метода. Для вызова статического используется INVOKESTATIC.

• GETFIELD — получение нестатического поля. Для статического есть GETSTATIC.

Также хочется отметить такие немаловажные инструкции, как LOAD и STORE. Первая загружает переменную в стек, а вторая выгружает из стека в переменную.

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

Например, INVOKESTATIC 42 означает, что в пуле констант под номером 42 лежит ссылка на метод. Так JVM определяет, какой метод нужно вызвать. Для простоты обычно сразу пишут название метода, чтобы не нужно было лезть в пул констант.

Такая же история и с LDC. У неё тоже единственный аргумент — индекс на пул констант, по которому она сможет определить, какую константу загрузить в стек.

Этих базовых знаний будет достаточно, чтобы написать наш процессор.

Как пропатчить байт-код

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

Keys.create(SomeClass.class);

Если посмотреть на байт-код, то тут Java или Kotlin генерируют всего две инструкции:

LDC SomeClass INVOKESTATIC Keys.create

В LDC передаётся ссылка на класс, поэтому она загружает в стек наш Class-объект, а INVOKESTATIC просто вызывает метод create. Вместо этого мы хотим загрузить обычную Int-константу. Для этого подойдёт SIPUSH-инструкция:

SIPUSH 42

Это такая же инструкция, как и LDC, которая загружает константу в стек, но она не использует ссылку на пул констант, а сразу же после себя содержит значение. Минус в том, что её аргумент ограничен 2 байтами, поэтому значения не могут быть больше 2^15 (32 768), но нам пока этого хватит.

Во всей этой магии по замене инструкций нам сильно поможет библиотека BCEL (Byte Code Engineering Library). Она предоставляет очень удобный API, поэтому я выбрал именно её. Сначала напишем метод, который пробежится по всем методам класса и достанет оттуда связанный список инструкций:

fun modifyClass(file: File) {     // Парсим .class-файл     val parser = ClassParser(file.path)     val javaClass = parser.parse()     val classGen = ClassGen(javaClass)      // Получаем пул констант     val constantPoolGen = classGen.constantPool      for (method in classGen.methods) {         val methodGen = MethodGen(method, classGen.className, constantPoolGen)          modifyMethod(methodGen.instructionList, constantPoolGen)          // Сохраняем изменения в методе         classGen.replaceMethod(method, methodGen.method)     }      // Сохраняем изменения в классе     classGen.javaClass.dump(FileOutputStream(file)) }

Теперь реализуем modifyMethod(), который будет заменять инструкции:

fun modifyMethod(     list: InstructionList,     constantPoolGen: ConstantPoolGen ) {     var instruction: InstructionHandle? = null      while (true) {         // Перебираем все инструкции — в этой библиотеке они представлены связанным списком         instruction = (if (instruction == null) list.start else instruction.next) ?: break          // Первая инструкция — это LDC         val ldc = instruction.instruction as? LDC ?: continue         // Убеждаемся, что LDC-инструкция загружает константу класса         val objectType = ldc.getValue(constantPoolGen) as? ObjectType ?: continue         // Узнаём имя этого класса         val className = objectType.className ?: continue          // Следующей инструкцией должна быть INVOKESTATIC         val invokestatic = instruction.next?.instruction as? INVOKESTATIC ?: continue         // Убеждаемся, что мы вызываем метод Keys.create(Class<*>): Int         if (invokestatic.getLoadClassType(constantPoolGen).className != "Keys") continue         if (invokestatic.getName(constantPoolGen) != "create") continue         if (invokestatic.getSignature(constantPoolGen) != "(Ljava/lang/Class)I") continue         // Сигнатура (Ljava/lang/Class)I означает, что метод принимает один аргумент с типом Class, а возвращает Int           // Определяем индекс для нашего класса: просто берём следующий свободный         var index = indexMap[className]         if (index == null) {             index = indexMap.size             indexMap[className] = index // Сохраняем индекс под текущее имя класса         }          // Удаляем инструкцию LDC, заменяем её на инструкцию NOP, которая ничего не делает         instruction.instruction = InstructionConst.NOP         // Заменяем INVOKESTATIC-инструкцию на SIPUSH, которая будет пушить в стек наш индекс         instruction.next.instruction = SIPUSH(index.toShort())     } }

Наш процессор готов! 

Однако теперь возникает вопрос: откуда нам взять .class-файлы для их изменения? С JVM всё просто: JAR-файл содержит нужные .class-файлы. А что касается Android, то внутри APK у нас лежат .dex-файлы. Это тоже байт-код, но не Java, а Dalvik. Он компилируется из байт-кода Java, и в этом случае нам нужно вставить процессор перед этим шагом.

Как интегрировать процессор в систему сборки

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

Чтобы начать работу с Gradle, достаточно понять: основная единица работы в нём — это Task. Таски могут быть связаны: например, задача компиляции байт-кода Dalvik зависит от задачи компиляции Kotlin в байт-код Java. Все задачи в Gradle объединены в проекты, но мы их называем просто модулями.

Акцентируем внимание на таске compileKotlin. Он обрабатывает наши .kt-файлы (исходники Kotlin), создавая .class-файлы (байт-код JVM), а это именно то, что нам нужно.

В корневом файле build.gradle.kts добавим следующий код:

allprojects {     // Проходимся по всем таскам всех проектов.      tasks.configureEach {         // Отфильтровываем все таски с названием compileKotlin         if (name == "compileKotlin") {             // doLast вызывает лямбду сразу после того,             // как таск compileKotlin закончит своё выполнение             doLast {                 // Проходимся по всем выходным папкам таска compileKotlin                 // Там и должны лежать наши .class-файлы                 outputs.files.forEach { output ->                     output.walk()                         .filter { file -> file.isFile && file.extension == "class" } // Фильтруем все .class-файлы                          .forEach { file ->                             // Изменяем наши классы методом, который мы описали ранее                             // Этот метод должен лежать где-то в модуле buildSrc, чтобы он был виден во всех Gradle-скриптах                             modifyClass(file)                         }                 }             }         }     } }

Теперь наш процессор будет запускаться всегда, когда мы компилируем код Kotlin. Неважно, для чего нам это понадобится — для компиляции Android или JVM-приложения, — он всегда будет запускаться и модифицировать байт-код.

Результаты

Вы спросите, как наши изменения повлияли на скорость?

Это скриншот перформанс-тестов Android-приложения Яндекс Маркета. После переноса всего проекта на Int-ключи мы вернулись к прежним значениям времени старта приложения. На скриншоте также видно предыдущую попытку ускорения, о которой можно прочитать в первой статье про Scout.

Работа с байт-кодом оказалась увлекательным занятием. Он достаточно прост для внесения изменений и не так страшен, как чистый ассемблер. JVM проверяет байт-код и в случае ошибки в генерации выдаёт исключение VerifyError с детальным описанием возникшей проблемы.

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