Автотрекинг производительности Room Database

Room (абстракция над SQLite) — одна из основных технологий, используемая почти во всех Android-приложениях для кэширования данных, оффлайновости, как cross-process хранилище данных и тому подобное.

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

Встает вопрос: “А не является ли сама работа с БД узким местом скорости старта приложения?”

В Wildberries это особенно актуально, так как приложение построено на парадигме offline-first, когда почти вся информация кэшируется в БД, чтобы приложение работало даже с медленным интернетом или без него.

Для ответа на этот вопрос в статье разберем рантайм реализацию автоматического трекинга скорости выполнения запросов и транзакций в Room Database на основе Java Dynamic Proxy.

Мотивация

Профилировщик Android Studio не даёт ответ на вопрос о производительности БД, так как:

• содержит слишком много ненужной информации (скорость работы всех методов, не только БД)

• работает только локально, а не у конечных пользователей (а как мы знаем: пользователи попадают в такие ситуации, которые тестировщикам только снятся)

Основные проблемы производительности БД связаны с тем, что по сути вся работа с ней является синхронной очередью — чтение блокирует запись и наоборот. Write-Ahead Logging (WAL) позволяет обойти это, создавая буфер для записей и периодически скидывая их в БД, но при этом технология подвержена другим проблемам: каждая N-ная запись приводит к тормозам (из-за сброса буфера), сам буфер при определенных условиях может бесконтрольно расти и т.п..

Среди очевидных проблем производительности БД можно выделить:

1. в database.withTransaction{} выполняется долгий/блокирующий код (например, лок мьютексов, сетевые запросы, тяжелые мэппинги или cross-database взаимодействия),

2. множество чтений/записей, которые выгоднее было бы выполнить в одной транзакции,

3. множество чтений одних и тех же данных (которые можно закэшировать в рантайме),

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

Цель

Научиться узнавать о проблемах производительности БД, а именно:

• Получать информацию как о транзакциях, так и об обычных DAO-запросах, сгенерированных Room.

• В информации должны быть данные о том, что именно замерялось, например, userDao_all для запросов к DAO (так как мы имеем имя DAO и название метода запроса) или UserWorker_runUpdate в случае если был вызван db.withTransaction из метода runUpdate в классе UserWorker (так как транзакции не имеют своего имени).

• Информацию хочется получать и локально, и с реальных пользователей.

• По возможности механизм должен работать без больших переделок существующего кода работы с БД.

Целевой код

Для замеров скорости будем использовать абстрактный Profiler.Trace, который может отправлять данные в firebase, логкат или куда-либо ещё:

interface Profiler {
   fun startTrace(name: String): Trace
   interface Trace {
       fun stop()
   }
}

В качестве Room БД возьмем базу пользователей:

@Entity
data class User(
   @PrimaryKey val uid: Int,
)

@Database(entities = [User::class], version = 1)
abstract class AppDatabase : RoomDatabase() {
   abstract fun userDao(): UserDao
}

@Dao
interface UserDao {
   @Query("SELECT * FROM user")
   fun observe(): Flow<List<User>>

   @Query("SELECT * FROM user")
   suspend fun all(): List<User>

   @Insert(onConflict = REPLACE)
   fun insert(user: User): Long
}

И как она используется в проекте:

// Настройка
val db = Room.databaseBuilder(
   appContext, AppDatabase::class.java, "app-db"
).build()

class Worker {
   suspend fun payload(db: AppDatabase) {
       val userDao = db.userDao()

       // Какие трейсы мы хотим получить на какие вызовы:

       // Должно логировать "UserDao_insert"
       userDao.insert(User(1))

       // Должно логировать "UserDao_all"
       userDao.all()

       // Должно логировать "UserDao_observe", но только для самого первого элемента в отдельном collect{}
       userDao.observe().first()

       // Должно логировать "Worker_payload", так как транзакции не имеют имени - берем имя вызывающего
       db.withTransaction { userDao.all() }
   }
}

Реализация

Трекинг можно реализовать двумя способами — в рантайме и на уровне компиляции.
Остановимся на рантайм реализации, так как:

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

• байткод манипуляции сложны в плане написания и поддержки (для этого нужен старший+ разработчик).

Автотрекинг транзакций: db.withTransaction{}

Идея основывается на том, что работа с транзакциями почти всегда идёт через единственный библиотечный метод-экстеншн:

package androidx.room

suspend fun <R> RoomDatabase.withTransaction(block: suspend () -> R): R

При этом в приложениях обычно используются наследники БД (у нас это AppDatabase), а не RoomDatabase напрямую. Значит, мы можем сделать автотрекинг за 3 шага:

1. Создать аналог библиотечного метода в том же пакете, но как экстеншн к собственному интерфейсу (ProfiledDatabase).

suspend fun <R> ProfiledDatabase.withTransaction(block: suspend () -> R): R {
   // Что это и как работает - позднее в статье
   val profiler = (this as ProfilerProvider).profiler()

   val trace = profiler.startTrace(transactionName())
   val result = (this as RoomDatabase).withTransaction(block)
   trace.stop()
   
   return result
}

2. Наследовать этот новый интерфейс во всех БД приложения, изменив иерархию.

interface ProfiledDatabase

// Выделили отдельный интерфейс со всеми Dao
interface AppDatabase : ProfiledDatabase {
   abstract fun userDao(): UserDao
}

@Database(entities = [User::class], version = 1)
abstract class RoomAppDatabase : RoomDatabase(), AppDatabase

Теперь вызовы AppDatabase.withTransaction{} будут приводить к вызову нашего ProfiledDatabase.withTransaction{} вместо RoomDatabase.withTransaction{}.

3. Получать класс и имя метода, в которых был вызван withTransaction.

Идея в том, чтобы у потока получить стектрейс, через который сможем найти информацию о том, какой код вызвал withTransaction.

@Suppress("NOTHING_TO_INLINE") // inline для корректности stackTrace
internal inline fun transactionName(): String? {
   // Получаем стектрейс текущего потока
   val stackTrace = Thread.currentThread().stackTrace

   // Отбрасываем все вызовы библиотеки замера производительности,
   // а так же dalvik что появляется при дебаге
   val stackIndex = (if (stackTrace[0].className.startsWith("dalvik")) 3 else 2)
  
   // Получаем информацию о коде, вызвавшем withTransaction
   val caller = stackTrace.getOrNull(stackIndex)
       ?: return null

   // Нас интересует имя класса без пакета
   val className = caller.className.substringAfterLast('.')

   // Имена сгенерированных классов/методов/лямбды нужно подчистить перед использованием
   return if (className.contains('$')) {
       // Class name:
       // UserRepository$localUser$2$$inlined$$invokeSuspend
       val tokens = className.splitToSequence('$').iterator()
       "${tokens.nextOrNull()}_${tokens.nextOrNull()}"
   } else {
       //
       // Class name: DatabasePerformanceTrackerTest
       // Method name: withTransaction_logged__when_called_from_local_function$lambda$17$doLocalTransaction
       "${className}_${caller.methodName.substringBefore('$')}"
   }
}

Автотрекинг запросов (query)

AppDatabase и UserDao — интерфейсы, а значит можно использовать механизм java dynamic Proxy, который позволяет через рефлексию в рантайме имплементировать интерфейсы (например, его использует Retrofit).
Для этого нужно реализовать InvocationHandler, который будет перехватывать все вызовы методов к Proxy:

public interface InvocationHandler {
   public Object invoke(Object proxy, Method method, Object[] args)
       throws Throwable;
}

То есть, мы сможем создать Proxy AppDatabase, который будет создавать Proxy UserDao, который в свою очередь будет для вызовов всех методов DAO автоматически стартовать и завершать трейсы.

Note: Несмотря на то, что Proxy использует рефлексию, в актуальных андроидах оверхеда почти нет (порядка 0.01мс на один вызов, тогда как вызовы БД обычно занимают 1мс+).

Proxy AppDatabase & UserDao

Сначала создадим Proxy над AppDatabase интерфейсом:

inline fun <reified TDB : ProfiledDatabase> RoomDatabase.withProfiler(profiler: Profiler) : TDB {
   val dbInterface = TDB::class.java
   require(dbInterface.isInstance(this))

   return Proxy.newProxyInstance(
       this::class.java.classLoader,
       arrayOf(dbInterface, ProfilerProvider::class.java),
       DatabaseInvocationHandler(
           db = this,
           profiler = profiler,
       ),
   ) as TDB
}

Это позволит одним вызовом db.withProfiler(profiler) включать автотрекинг производительности любой RoomDatabase в проекте (после того, как AppDatabase переделан на интерфейс).

Тут мы объявили, что наша Proxy будет наследовать AppDatabase, чтобы создавать Proxy для отдельных DAO, а также ProfilerProvider, через который мы в транзакциях будем получать доступ до Profiler для создания трейсов.

Далее реализация самой прокси для БД:

class DatabaseInvocationHandler(
   private val db: RoomDatabase,
   private val profiler: Profiler,
) : InvocationHandler {
   // Кэш всех настоящих прокси-Dao (UserDao и тому подобных)
   private val daos = ConcurrentHashMap<Method, Any>()

   // Здесь и далее @Throws нужен для устранения проблем несовместимости
   // работы с исключениями у Kotlin и Java, в частности UndeclaredThrowableException.
   @Throws(Throwable::class)
   override fun invoke(proxy: Any, method: Method, args: Array<*>?): Any? {
       // Если запросили profiler, а не dao - возвращаем его.
       // За счет этого работает withTransaction.
       if (method.name == ProfilerProvider::profiler.name) {
           return profiler
       }

       return daos.computeIfAbsent(method) {
           // Если в кэше нет нужного Dao - получаем его, оборачиваем в прокси и кэшируем
           val realDao = method.invokeSafe(db, args)
           // Мы не можем использовать realDao::class, так как он будет сгенерирован Room - UserDao_Impl
           val daoType = method.returnType
           daoType.createDaoProxy(realDao)
       }
   }

   private fun Class<*>.createDaoProxy(realDao: Any) = Proxy.newProxyInstance(
       DatabaseInvocationHandler::class.java.classLoader,
       arrayOf(this),
       DaoInvocationHandler(
           daoName = simpleName,
           realDao,
           profiler,
       ),
   )
}

Здесь мы создаем (и кэшируем) Proxy DAO для каждого метода в AppDatabase (в нашем случае userDao(), а также возвращаем profiler для поддержки withTransaction.

Для вызовов настоящих методов из Proxy используем специальную обертку над invoke, которая умеет разворачивать оберточные Kotlin-исключения. Это связано с тем, что Kotlin и Java имеют разные системы исключений и, так как Proxy работает на уровне Java рефлексии, необходимо быть особенно осторожным:

@Throws(Throwable::class)
fun Method.invokeSafe(obj: Any, args: Array<*>?) =
   try {
       invoke(obj, *(args ?: emptyArray()))
   } catch (e: InvocationTargetException) {
       throw e.targetException ?: e.cause ?: e
   }

Сама DAO Proxy:

class DaoInvocationHandler(
   private val daoName: String,
   private val dao: Any,
   private val profiler: Profiler,
) : InvocationHandler {

   @Throws(Throwable::class)
   override fun invoke(proxy: Any, method: Method, args: Array<*>?): Any? {
       // Все методы реализуем позднее
       return when {
           isSuspend(args) -> handleSuspendMethod(method, args)
           isObserver(method) -> handleObserverMethod(method, args)
           else -> handleSyncMethod(method, args)
       }
   }

   fun isSuspend(args: Array<*>?): Boolean = TODO()
   @Throws(Throwable::class)
   fun handleSuspendMethod(method: Method, args: Array<*>?): Any = TODO()

   fun isObserver(method: Method): Boolean = TODO()
   @Throws(Throwable::class)
   fun handleObserverMethod(method: Method, args: Array<*>?): Any = TODO()

   @Throws(Throwable::class)
   fun handleSyncMethod(method: Method, args: Array<*>?): Any = TODO()

Всего есть 3 типа методов, которые нужно обработать отдельно:

UserDao.insert()

Сигнатура insert() проста — это обычный синхронный метод:

@Insert(onConflict = REPLACE)
fun insert(user: User): Long

Поэтому invokeSafe() справится с ним сам:

@Throws(Throwable::class)
fun handleSyncMethod(method: Method, args: Array<*>?): Any {
   val trace = profiler.startTrace("${daoName}_${method.name}")

   return method.invokeSafe(dao, args)
       .also { trace.stop() }
}

UserDao.observe()

observe() хоть и обычный синхронный метод, но возвращает Flow.

@Query("SELECT * FROM user")
fun observe(): Flow<List<User>>

Так как нам важно слушать каждый отдельный collect (пользователь может получить вызвать observe один раз, но переиспользовать полученный Flow в разных местах), то единственный подходящий вариант — через наследование:

fun isObserver(method: Method): Boolean =
   method.returnType == Flow::class.java

@Throws(Throwable::class)
fun handleObserverMethod(method: Method, args: Array<*>?): Any {
   val dbFlow = method.invokeSafe(dao, args) as Flow<*>

   return object : Flow<Any?> {
       override suspend fun collect(collector: FlowCollector<Any?>) {
           val trace = profiler.startTrace("${daoName}_${method.name}")

           dbFlow
               .onEach { trace.stop() }
               .collect(collector)
       }
   }
}

Мы замеряем только производительность первого элемента потому что:

• Время между collect и первым элементов включает в себя время создания SQLStatement, что также важная часть производительности БД.

• Если замерять выдачу всех элементов, то количество данных будет слишком большим.

UserDao.all()

all() — это suspend метод, которых как таковых не существует в Java.

@Query("SELECT * FROM user")
suspend fun all(): List<User>

Так как Java Proxy ничего не знает о suspend, необходимо вручную научить его. Для этого нужно разобраться с тем, как работает suspend с точки зрения кодогенерации Kotlin в Java.

Корутины в Java

Код, который будем декомпилировать:

class Worker {
   suspend fun payload1() { println("1") }
   suspend fun payload2() { println("2") }

   suspend fun doWork() {
       payload1()
       payload2()
   }
}

После декомпиляции сигнатура suspend метода doWork в Java выглядит так:

public final Object doWork(Continuation parentContinuation)

В сигнатуру последним параметром был добавлен Continuation — “колбэк”, который будет вызван по окончанию выполнения метода doWork. Это позволяет методу прерывать свою работу на полпути, переключаться между потоками или блокироваться на другие suspend методы.
На самом деле, Continuation вызывается не всегда, об этом далее.

Основная часть java реализации doWork выглядит так (это вручную переписанный декомпилированный код для упрощения понимания):

public final Object doWork(Continuation parentContinuation) {
   // Continuation создается и хранится для каждого отдельного вызова doWork.
   // Для простоты, опускаем реализацию хранения doWork_continuation.
   if (doWork_continuation == null) {
       doWork_continuation = new ContinuationImpl(parentContinuation) {
           int label = 0;

           @Nullable
           public Object invokeSuspend(Object omitted) {
               return doWork(this);
           }
       };
   }

   // О том, что это такое - позднее
   Object SUSPENDED = IntrinsicsKt.getCOROUTINE_SUSPENDED();

   switch (doWork_continuation.label) {
       case 0:
           doWork_continuation.label = 1;
           if (payload1(doWork_continuation) == SUSPENDED) {
               return SUSPENDED;
           }
           // Нет break. Если не было suspend - проваливаемся в следующий case
       case 1:
           doWork_continuation.label = 2;
           if (payload2(doWork_continuation) == SUSPENDED) {
               return SUSPENDED;
           }
           // Нет break. Если не было suspend - проваливаемся в следующий case
       case 2:
           return Unit.INSTANCE;
   }
}

Методы payload вызываются с Continuation-оберткой, которая ведет на текущий doWork метод. Он в свою очередь восстанавливает свою работу за счет label, пропуская с его помощью уже выполненную часть кода. 

То есть, по сути метод doWork может быть неявно вызван 3 раза:

• самый первый вызов,

• вызов после окончания payload1,

• вызов после окончания payload2.

Но есть нюанс. Suspend методы могут быть не suspend. Причины могут быть разные, но это приводит к тому что Continuation может быть не вызван. Например, так выглядит метод payload1:

public final Object payload1(Continuation completion) {
   System.out.println("1");
   return Unit.INSTANCE;
}

Continuation тут не вызывается, а сразу возвращается результат, ведь внутри payload1 нет ничего, что могло бы вызвать прерывание метода.

Такие ситуации обрабатываются за счет специального объекта — Intrinsics.COROUTINE_SUSPENDED. Если suspend функция его вернет, код снаружи знает, что произошло прерывание выполнения, например, смена потока, и нужно выходить из текущего стека вызовов, так как продолжение выполнения со следующего лейбла будет через Continuation.invokeSuspend.

Корутины в Proxy

В итоге мы получаем следующую реализацию трекинга производительности:

fun isSuspend(args: Array<*>?): Boolean =
   args?.lastOrNull() as? Continuation<Any?>

@Throws(Throwable::class)
fun handleSuspendMethod(method: Method, args: Array<*>?): Any {
   val trace = profiler.startTrace("${daoName}_${method.name}")

   val continuation = args?.lastOrNull() as? Continuation<Any?>
   // В случае если внутри метода будет прерывание, оборачиваем
   // старый Continuation в наш новый, добавляя остановку трейса.
   val newContinuation = continuation.withBefore {
       trace.stop()
   }

   val newArgs = args.dropLast(1) + newContinuation
   method.invokeSafe(dao, newArgs.toTypedArray()).also {
       if (it != COROUTINE_SUSPENDED) {
           // Если же прерывания не случилось и метод отработал синхронно -
           // останавливаем трейс сразу.
           trace.stop()
       }
   }
}

Самая большая проблема в этой логике в withBefore — это самописный метод, использующий hidden coroutine api:

private fun Continuation<Any?>.withBefore(
   action: (Result<Any?>) -> Unit,
) = object : Continuation<Any?> {
   override val context = [e-mail]
   override fun resumeWith(result: Result<Any?>) {
       action(result)

       /**
        * Room внутри использует [Executor.asCoroutineDispatcher], который является
        * [ContinuationInterceptor] и отвечает за создание [DispatchedContinuation].
        * [DispatchedContinuation] использует скрытое api [ContinuationImpl.intercepted]
        * для обеспечения редиспатчинга [Continuation] на нужный поток.
        *
        * Здесь мы делаем то, что делает [DispatchedCoroutine.afterResume], используя скрытое api.
        */
       [e-mail]().resumeCancellableWith(result)
   }
}

Связано это с тем, что механизм “интерцептов” в корутинах отвечает за переключение потоков использования такого же механизма,  userDao.all() может продолжить свое выполнение на потоке Room, что сломает приложение.
Изначально было опасение что данная конструкция будет часто ломаться, но на данный момент ей уже год и проблем не наблюдалось. Известно, что ещё такая конструкция использовалась в ktor.

Итог

Автотрекинг работает в рантайме, имеет малый оверхед и код достаточно простой для расширения (за исключением “интерцепт”-магии корутин). При этом позволяет отследить все основные проблемы медленной работы БД.

Как примеры расширений на основе этого механизма, можно упомянуть:

• Возможность отслеживать дедлоки (когда вызван trace.start, но не вызван trace.stop).

• Возможность логировать нонфаталы для слишком долгих транзакций (если не нужны логи обо всех БД операциях).

• Возможность поддержать DAO запросы с одинаковыми именами, но разными параметрами.

Из недостатков:

1. Черная магия корутин (есть вероятность что это просто недосмотр, так как в ContinuationKt утилитах есть возможность создавать свои обертки над Continuation с учетом intercepted, нет лишь нужного нам метода).

2. Для интеграции с обфусцированными приложениями нужно делать дополнительные манипуляции на бэкенде (например, выгружать mapping.txt от R8 и де-обфусцировать названия).

3. Код использования БД в приложении почти не нужно менять (нужно только изменить иерархию БД классов и добавить враппер над RoomDatabase), за исключением db.inTransaction. В этом случае придется делать свой метод типа ProfiledDatabase.isInTransaction() = (this as RoomDatabase).inTransaction() и менять код в проекте.

4. Не поддерживаются ручные db.beginTransaction/endTransaction (однако эти методы использовать не рекомендуется, так как высока цена ошибки).

5. @Transaction запросы в DAO поддерживаются, но если внутри такого метода используются другие запросы из DAO, они залогированы не будут. Это связано с тем, что, когда Room генерирует имплементацию таких методов, он передает внутрь @Transaction метода настоящий DAO, а не наш Proxy. И обойти это в рантайме нельзя. Но это не является большим недостатком, так как время самого @Transaction метода всё ещё будет логироваться, поэтому проблемы производительности не пройдут незамеченными.