Skip to content

Commit

Permalink
perf: Free memory earlier and remove negligible lookup maps
Browse files Browse the repository at this point in the history
Negligible lookup maps used for matching fingerprints have been removed to reduce the likelihood of OOM when the maps are instantiated, commonly observed with 400M RAM. Additionally, lookup maps previously kept for the duration of the patcher instance are now cleared before the classes are compiled. This reduces the likelihood of OOM when compiling classes.
On a related note, a linear increase in memory usage is observed with every compiled class until all classes are compiled implying compiled classes not being freed by GC because they are still referenced. After compiling a class, the class is technically free-able though. The classes are assumed to be referenced in the `multidexlib2` library that takes the list of all classes to compile multiple DEX with and seems to hold the reference to all these classes in memory until all DEX are compiled. A clever fix would involve splitting the list of classes into chunks and getting rid of the list of all classes so that after every DEX compilation, the corresponding split of classes can be freed.
  • Loading branch information
oSumAtrIX committed Oct 2, 2024
1 parent f1615b7 commit 5ec9f39
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 94 deletions.
31 changes: 5 additions & 26 deletions src/main/kotlin/app/revanced/patcher/Fingerprint.kt
Original file line number Diff line number Diff line change
Expand Up @@ -66,36 +66,15 @@ class Fingerprint internal constructor(
}

// TODO: If only one string is necessary, why not use a single string for every fingerprint?
fun Fingerprint.lookupByStrings() = strings?.firstNotNullOfOrNull { lookupMaps.methodsByStrings[it] }
if (lookupByStrings()?.let(::match) == true) {
if (strings?.firstNotNullOfOrNull { lookupMaps.methodsByStrings[it] }?.let(::match) == true) {
return true
}

// No strings declared or none matched (partial matches are allowed).
// Use signature matching.
fun Fingerprint.lookupBySignature(): MethodClassPairs {
if (accessFlags == null) return lookupMaps.allMethods

var returnTypeValue = returnType
if (returnTypeValue == null) {
if (AccessFlags.CONSTRUCTOR.isSet(accessFlags)) {
// Constructors always have void return type.
returnTypeValue = "V"
} else {
return lookupMaps.allMethods
}
}

val signature =
buildString {
append(accessFlags)
append(returnTypeValue.first())
appendParameters(parameters ?: return@buildString)
}

return lookupMaps.methodsBySignature[signature] ?: return MethodClassPairs()
context.classes.forEach { classDef ->
if (match(context, classDef)) return true
}
return match(lookupBySignature())

return false
}

/**
Expand Down
12 changes: 5 additions & 7 deletions src/main/kotlin/app/revanced/patcher/Patcher.kt
Original file line number Diff line number Diff line change
Expand Up @@ -98,14 +98,12 @@ class Patcher(private val config: PatcherConfig) : Closeable {

logger.info("Merging extensions")

context.executablePatches.forEachRecursively { patch ->
if (patch is BytecodePatch && patch.extension != null) {
context.bytecodeContext.merge(patch.extension)
}
}
with(context.bytecodeContext) {
context.executablePatches.mergeExtensions()

// Initialize lookup maps.
context.bytecodeContext.lookupMaps
// Initialize lookup maps.
lookupMaps
}

logger.info("Executing patches")

Expand Down
1 change: 1 addition & 0 deletions src/main/kotlin/app/revanced/patcher/PatcherConfig.kt
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ class PatcherConfig(
private val temporaryFilesPath: File = File("revanced-temporary-files"),
aaptBinaryPath: String? = null,
frameworkFileDirectory: String? = null,
@Deprecated("This is going to be removed in the future because it is not needed anymore.")
internal val multithreadingDexFileWriter: Boolean = false,
) {
private val logger = Logger.getLogger(PatcherConfig::class.java.name)
Expand Down
91 changes: 30 additions & 61 deletions src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import lanchon.multidexlib2.MultiDexIO
import lanchon.multidexlib2.RawDexIO
import java.io.Closeable
import java.io.FileFilter
import java.io.InputStream
import java.util.*
import java.util.logging.Logger

Expand Down Expand Up @@ -60,40 +59,41 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
internal val lookupMaps by lazy { LookupMaps(classes) }

/**
* A map for lookup by [merge].
* Merge the extensions for this set of patches.
*/
internal val classesByType = mutableMapOf<String, ClassDef>().apply {
classes.forEach { classDef -> put(classDef.type, classDef) }
}
internal fun MutableSet<Patch<*>>.mergeExtensions() {
// Lookup map for fast checking if a class exists by its type.
val classesByType = mutableMapOf<String, ClassDef>().apply {
classes.forEach { classDef -> put(classDef.type, classDef) }
}

/**
* Merge an extension to [classes].
*
* @param extensionInputStream The input stream of the extension to merge.
*/
internal fun merge(extensionInputStream: InputStream) {
val extension = extensionInputStream.readAllBytes()
forEachRecursively { patch ->
if (patch is BytecodePatch && patch.extension != null) {

RawDexIO.readRawDexFile(extension, 0, null).classes.forEach { classDef ->
val existingClass = classesByType[classDef.type] ?: run {
logger.fine("Adding class \"$classDef\"")
val extension = patch.extension.readAllBytes()

classes += classDef
classesByType[classDef.type] = classDef
RawDexIO.readRawDexFile(extension, 0, null).classes.forEach { classDef ->
val existingClass = classesByType[classDef.type] ?: run {
logger.fine("Adding class \"$classDef\"")

return@forEach
}
classes += classDef
classesByType[classDef.type] = classDef

return@forEach
}

logger.fine("Class \"$classDef\" exists already. Adding missing methods and fields.")
logger.fine("Class \"$classDef\" exists already. Adding missing methods and fields.")

existingClass.merge(classDef, this@BytecodePatchContext).let { mergedClass ->
// If the class was merged, replace the original class with the merged class.
if (mergedClass === existingClass) {
return@let
}
existingClass.merge(classDef, this@BytecodePatchContext).let { mergedClass ->
// If the class was merged, replace the original class with the merged class.
if (mergedClass === existingClass) {
return@let
}

classes -= existingClass
classes += mergedClass
classes -= existingClass
classes += mergedClass
}
}
}
}
}
Expand Down Expand Up @@ -145,6 +145,9 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
override fun get(): Set<PatcherResult.PatchedDexFile> {
logger.info("Compiling patched dex files")

// Free up memory before compiling the dex files.
lookupMaps.close()

val patchedDexFileResults =
config.patchedFiles.resolve("dex").also {
it.deleteRecursively() // Make sure the directory is empty.
Expand Down Expand Up @@ -178,21 +181,6 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
* @param classes The list of classes to create the lookup maps from.
*/
internal class LookupMaps internal constructor(classes: List<ClassDef>) : Closeable {
/**
* Classes associated by their type.
*/
internal val classesByType = classes.associateBy { it.type }.toMutableMap()

/**
* All methods and the class they are a member of.
*/
internal val allMethods = MethodClassPairs()

/**
* Methods associated by its access flags, return type and parameter.
*/
internal val methodsBySignature = MethodClassPairsLookupMap()

/**
* Methods associated by strings referenced in it.
*/
Expand All @@ -203,22 +191,6 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
classDef.methods.forEach { method ->
val methodClassPair: MethodClassPair = method to classDef

// For fingerprints with no access or return type specified.
allMethods += methodClassPair

val accessFlagsReturnKey = method.accessFlags.toString() + method.returnType.first()

// Add <access><returnType> as the key.
methodsBySignature[accessFlagsReturnKey] = methodClassPair

// Add <access><returnType>[parameters] as the key.
methodsBySignature[
buildString {
append(accessFlagsReturnKey)
appendParameters(method.parameterTypes)
},
] = methodClassPair

// Add strings contained in the method as the key.
method.instructionsOrNull?.forEach instructions@{ instruction ->
if (instruction.opcode != Opcode.CONST_STRING && instruction.opcode != Opcode.CONST_STRING_JUMBO) {
Expand Down Expand Up @@ -259,15 +231,12 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
}

override fun close() {
allMethods.clear()
methodsBySignature.clear()
methodsByStrings.clear()
}
}

override fun close() {
lookupMaps.close()
classesByType.clear()
classes.clear()
}
}
Expand Down

0 comments on commit 5ec9f39

Please sign in to comment.