paparazzi: Getting OutOfMemoryErrors after ~500 screenshots

@jmartinesp After conducting some investigations, I identified it is related to the memory leak of context in BridgeContentResolver from animationScale map in WindowRecomposer.android.kt.

In your instance where lots of test cases are executed within same Paparazzi instance, an excessive accumulation of AndroidComposeView occurs, leading to the occurrence of memory leak.

An issue has been filed to the layoutlib team.

cc @jrodbx

This fix works around the memory leak issue, tested on kotlin 1.9.10 This involves a couple of steps:

1. Wrap Paparazzi in your own rule, and use a RunRules chain to run PaparazziCleanupRule at the end.
2. Add PaparazziCleanupRule using a reflection hack to retrieve the private bridgeRenderSession object.
3. Add the extension function that hacks around the RenderSession object.

TestRule wrapper

class PaparazziRule internal constructor(
    private val testDispatcher: TestDispatcher = UnconfinedTestDispatcher(),
    val paparazzi: Paparazzi,
) : TestRule {

    override fun apply(base: Statement, description: Description): Statement =
 RunRules(
            base,
            listOf(MainTestWatcher(testDispatcher), PaparazziCleanupRule(paparazzi), paparazzi),
            description,
        )
 // ** wrapper test rule code including snapshot functions calling into nested paparazzi
}

## Parazzi Reflection Hack
private class PaparazziCleanupRule(
        private val paparazzi: Paparazzi,
    ) : TestWatcher() {

        override fun finished(description: Description?) {
            super.finished(description)
            @Suppress("UNCHECKED_CAST")
            val renderSession: RenderSession = (
                paparazzi::class.memberProperties
                    .first { it.name == "bridgeRenderSession" } as KProperty1<Paparazzi, RenderSession>
                )
                .apply { isAccessible = true }
                .invoke(paparazzi)
            renderSession.disposeHack()
        }
    }

Rendersession Hack

package **

import androidx.compose.ui.platform.AndroidUiDispatcher
import androidx.compose.ui.platform.ComposeView
import app.cash.paparazzi.internal.ComposeViewAdapter
import com.android.ide.common.rendering.api.RenderSession
import com.android.ide.common.rendering.api.ViewInfo
import **.Logger
import java.lang.ref.WeakReference
import java.lang.reflect.InvocationTargetException
import java.util.concurrent.atomic.AtomicReference

private const val WINDOW_RECOMPOSER_ANDROID_KT_FQN =
    "androidx.compose.ui.platform.WindowRecomposer_androidKt"
private const val COMBINED_CONTEXT_FQN = "kotlin.coroutines.CombinedContext"
private const val TAG = "DisposeRenderSession"
private const val SNAPSHOT_KT_FQN = "androidx.compose.runtime.snapshots.SnapshotKt"

/**
 * Initiates a custom [RenderSession] disposal, involving clearing several static collections
 * including some Compose-related objects as well as executing default [RenderSession.dispose].
 */
fun RenderSession.disposeHack() {
    val applyObserversRef = AtomicReference<WeakReference<MutableCollection<*>?>?>(null)
    val toRunTrampolinedRef = AtomicReference<WeakReference<MutableCollection<*>?>?>(null)

    try {
        val windowRecomposer: Class<*> = Class.forName(WINDOW_RECOMPOSER_ANDROID_KT_FQN)
        val animationScaleField = windowRecomposer.getDeclaredField("animationScale")
        animationScaleField.isAccessible = true
        val animationScale = animationScaleField[windowRecomposer]
        if (animationScale is Map<*, *>) {
            (animationScale as MutableMap<*, *>).clear()
        }
    } catch (ex: ReflectiveOperationException) {
        // If the WindowRecomposer does not exist or the animationScale does not exist anymore,
        // ignore.
        Logger.warning(TAG, "Unable to dispose the recompose animationScale $ex")
    }
    applyObserversRef.set(WeakReference(findApplyObservers()))
    toRunTrampolinedRef.set(WeakReference(findToRunTrampolined()))

    execute {
        rootViews
            .filterNotNull()
            .forEach { v -> disposeIfCompose(v) }
    }
    val weakApplyObservers = applyObserversRef.get()
    if (weakApplyObservers != null) {
        val applyObservers = weakApplyObservers.get()
        applyObservers?.clear()
    }
    val weakToRunTrampolined = toRunTrampolinedRef.get()
    if (weakToRunTrampolined != null) {
        val toRunTrampolined = weakToRunTrampolined.get()
        toRunTrampolined?.clear()
    }
    dispose()
}

/**
 * Performs dispose() call against View object associated with [ViewInfo] if that object is an
 * instance of [ComposeViewAdapter]
 *
 * @param viewInfo a [ViewInfo] associated with the View object to be potentially disposed of
 */
private fun disposeIfCompose(viewInfo: ViewInfo) {
    val viewObject: Any? = viewInfo.viewObject
    if (viewObject !is ComposeViewAdapter) {
        return
    }
    try {
        val composeView = viewInfo.children[0].viewObject as ComposeView
        composeView.disposeComposition()
    } catch (ex: IllegalAccessException) {
        Logger.warning(TAG, "Unexpected error while disposing compose view $ex")
    } catch (ex: InvocationTargetException) {
        Logger.warning(TAG, "Unexpected error while disposing compose view $ex")
    }
}

private fun findApplyObservers(): MutableCollection<*>? {
    try {
        val applyObserversField = Class.forName(SNAPSHOT_KT_FQN).getDeclaredField("applyObservers")
        applyObserversField.isAccessible = true
        val applyObservers = applyObserversField[null]
        if (applyObservers is MutableCollection<*>) {
            return applyObservers
        }
        Logger.warning(TAG, "SnapshotsKt.applyObservers found but it is not a List")
    } catch (ex: ReflectiveOperationException) {
        Logger.warning(TAG, "Unable to find SnapshotsKt.applyObservers $ex")
    }
    return null
}

private fun findToRunTrampolined(): MutableCollection<*>? {
    try {
        val uiDispatcher = AndroidUiDispatcher::class.java
        val uiDispatcherCompanion = AndroidUiDispatcher.Companion::class.java
        val uiDispatcherCompanionField = uiDispatcher.getDeclaredField("Companion")
        val uiDispatcherCompanionObj = uiDispatcherCompanionField[null]
        val getMainMethod =
            uiDispatcherCompanion.getDeclaredMethod("getMain").apply { isAccessible = true }
        val mainObj = getMainMethod.invoke(uiDispatcherCompanionObj)
        val combinedContext = Class.forName(COMBINED_CONTEXT_FQN)
        val elementField = combinedContext.getDeclaredField("element").apply { isAccessible = true }
        val uiDispatcherObj = elementField[mainObj]

        val toRunTrampolinedField =
            uiDispatcher.getDeclaredField("toRunTrampolined").apply { isAccessible = true }
        val toRunTrampolinedObj = toRunTrampolinedField[uiDispatcherObj]
        if (toRunTrampolinedObj is MutableCollection<*>) {
            return toRunTrampolinedObj
        }
        Logger.warning(TAG, "AndroidUiDispatcher.toRunTrampolined found but it is not a MutableCollection")
    } catch (ex: ReflectiveOperationException) {
        Logger.warning(TAG, "Unable to find AndroidUiDispatcher.toRunTrampolined $ex")
    }
    return null
}

but there is a reference to a MainTestWatcher that I can’t find definition of, any one else trying this?

@KennyGoers I am using it, I just included the rule in my test class directly as:

    @get:Rule
    val paparazziCleanupRule = PaparazziCleanupRule(paparazzi = paparazzi)

(paparazzi being the Paparazzi rule)

I still can’t see what MainTestWatcher is @KatieBarnett thoughts on that? Desperately trying to get this to work 😕

Thanks for any help in this.

Working on this more, you meant to say you didn’t need the Watcher thing? I am trying that with no change, not sure why but 1.3.2 and 1.3.3 hang gradle with this fix as I’ve attempted it from the description 😕

This issue is still happening with version 1.3.3. The above hack needs to be altered to comment out

    if (viewObject !is ComposeViewAdapter) {
        return
    }

as ComposeViewAdapter is now internal.

The hack still appears to work with this commented out.

Got it working now, thanks @agrosner!