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

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

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

В этой статье мы будем подробно рассматривать то, как применение байт-кода помогло сохранить скорость на том же уровне и спасти проект по миграции на 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 с детальным описанием возникшей проблемы.

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